295 lines
9.5 KiB
Python
295 lines
9.5 KiB
Python
# main.py (updated)
|
|
import discord
|
|
from discord.ext import commands
|
|
from config import ALLOWED_CHANNEL_ID, HIT_CHECK_INTERVAL, REASSIGN_DELAY
|
|
from cogs.assignments import Assignments
|
|
from cogs.commands import HitCommands
|
|
|
|
import asyncio
|
|
import uvicorn
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from fastapi import FastAPI, Request, HTTPException
|
|
from fastapi.responses import HTMLResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.templating import Jinja2Templates
|
|
from pydantic import BaseModel
|
|
|
|
# import server state
|
|
from services.server_state import STATE, Member
|
|
|
|
from services.torn_api import populate_friendly, populate_enemy, start_friendly_status_loop, start_enemy_status_loop
|
|
|
|
# ============================================================
|
|
# FastAPI Setup
|
|
# ============================================================
|
|
app = FastAPI()
|
|
|
|
templates = Jinja2Templates(directory="templates")
|
|
app.mount("/static", StaticFiles(directory="static"), name="static")
|
|
|
|
# ============================================================
|
|
# Dashboard Webpage
|
|
# ============================================================
|
|
@app.get("/", response_class=HTMLResponse)
|
|
async def dashboard(request: Request):
|
|
print(">>> DASHBOARD ROUTE LOADED")
|
|
return templates.TemplateResponse("dashboard.html", {"request": request})
|
|
|
|
# ============================================================
|
|
# Pydantic models for JSON POST input
|
|
# ============================================================
|
|
class FactionRequest(BaseModel):
|
|
faction_id: int
|
|
interval: int = 30
|
|
|
|
class AssignMemberRequest(BaseModel):
|
|
group_id: str # "1", "2", ... "5"
|
|
kind: str # "friendly" or "enemy"
|
|
member_id: int
|
|
|
|
class RemoveAssignmentRequest(BaseModel):
|
|
member_id: int
|
|
|
|
class BotControl(BaseModel):
|
|
action: str # "start" or "stop"
|
|
|
|
# ================================
|
|
# Helper: load JSON file into STATE
|
|
# ================================
|
|
def _load_json_list(path: Path):
|
|
if not path.exists():
|
|
return []
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
async def sync_state_from_file(path: Path, kind: str):
|
|
"""
|
|
Read JSON file (list of members dicts) and upsert into STATE.
|
|
Expected member dict keys: id, name, level, estimate, optionally status/hits.
|
|
"""
|
|
arr = _load_json_list(path)
|
|
received_ids = []
|
|
for m in arr:
|
|
try:
|
|
await STATE.upsert_member(m, kind)
|
|
received_ids.append(int(m["id"]))
|
|
except Exception as e:
|
|
print(f"Skipping bad member entry while syncing: {m} -> {e}")
|
|
await STATE.remove_missing_members(received_ids, kind)
|
|
|
|
# -----------------------------
|
|
# Populate endpoints (populate JSON once)
|
|
# -----------------------------
|
|
@app.post("/api/populate_friendly")
|
|
async def api_populate_friendly(data: FactionRequest):
|
|
await populate_friendly(data.faction_id)
|
|
# sync STATE from file
|
|
await sync_state_from_file(Path("data/friendly_members.json"), "friendly")
|
|
return {"status": "friendly populated", "id": data.faction_id}
|
|
|
|
@app.post("/api/populate_enemy")
|
|
async def api_populate_enemy(data: FactionRequest):
|
|
await populate_enemy(data.faction_id)
|
|
await sync_state_from_file(Path("data/enemy_memberes.json"), "enemy")
|
|
return {"status": "enemy populated", "id": data.faction_id}
|
|
|
|
# -----------------------------
|
|
# Start status refresh loops
|
|
# -----------------------------
|
|
@app.post("/api/start_friendly_status")
|
|
async def api_start_friendly_status(data: FactionRequest):
|
|
from services.torn_api import start_friendly_status_loop
|
|
await start_friendly_status_loop(data.faction_id, data.interval)
|
|
return {"status": "friendly status loop started", "id": data.faction_id, "interval": data.interval}
|
|
|
|
@app.post("/api/start_enemy_status")
|
|
async def api_start_enemy_status(data: FactionRequest):
|
|
from services.torn_api import start_enemy_status_loop
|
|
await start_enemy_status_loop(data.faction_id, data.interval)
|
|
return {"status": "enemy status loop started", "id": data.faction_id, "interval": data.interval}
|
|
|
|
# =============================
|
|
# Member JSON endpoints (unchanged)
|
|
# =============================
|
|
@app.get("/api/friendly_members")
|
|
async def get_friendly_members():
|
|
# Return list, but prefer STATE if populated
|
|
if STATE.friendly:
|
|
return [m.model_dump() for m in STATE.friendly.values()]
|
|
# fallback to file
|
|
path = Path("data/friendly_members.json")
|
|
return _load_json_list(path)
|
|
|
|
@app.get("/api/enemy_members")
|
|
async def get_enemy_members():
|
|
if STATE.enemy:
|
|
return [m.model_dump() for m in STATE.enemy.values()]
|
|
path = Path("data/enemy_members.json")
|
|
return _load_json_list(path)
|
|
|
|
# =============================
|
|
# Status JSON endpoints (unchanged)
|
|
# =============================
|
|
@app.get("/api/friendly_status")
|
|
async def api_friendly_status():
|
|
path = Path("data/friendly_status.json")
|
|
if not path.exists():
|
|
return {}
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
@app.get("/api/enemy_status")
|
|
async def api_enemy_status():
|
|
path = Path("data/enemy_status.json")
|
|
if not path.exists():
|
|
return {}
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
# =============================
|
|
# Assignment endpoints (server-side)
|
|
# =============================
|
|
@app.get("/api/assignments")
|
|
async def api_get_assignments():
|
|
"""
|
|
Return assignments snapshot:
|
|
{ "1": { "friendly": [...], "enemy": [...] }, "2": {...}, ... }
|
|
"""
|
|
snap = await STATE.get_assignments_snapshot()
|
|
return snap
|
|
|
|
@app.post("/api/assign_member")
|
|
async def api_assign_member(req: AssignMemberRequest):
|
|
group_id = req.group_id
|
|
kind = req.kind
|
|
member_id = req.member_id
|
|
|
|
# Validate group
|
|
if group_id not in STATE.groups:
|
|
raise HTTPException(status_code=400, detail="invalid group_id")
|
|
|
|
# Validate member exists in STATE (if not present, attempt to load from file)
|
|
if kind not in ("friendly", "enemy"):
|
|
raise HTTPException(status_code=400, detail="invalid kind")
|
|
|
|
coll = STATE.friendly if kind == "friendly" else STATE.enemy
|
|
if member_id not in coll:
|
|
# Try to load members from file into STATE and re-check
|
|
file_path = Path("data/friendly_members.json") if kind == "friendly" else Path("data/enemy_members.json")
|
|
await sync_state_from_file(file_path, kind)
|
|
if member_id not in coll:
|
|
raise HTTPException(status_code=404, detail="member not found")
|
|
|
|
await STATE.assign_member(member_id, kind, group_id)
|
|
return {"status": "ok", "group": group_id, "kind": kind, "member_id": member_id}
|
|
|
|
@app.post("/api/remove_member_assignment")
|
|
async def api_remove_member_assignment(req: RemoveAssignmentRequest):
|
|
member_id = req.member_id
|
|
await STATE.remove_member_assignment(member_id)
|
|
return {"status": "ok", "member_id": member_id}
|
|
|
|
@app.post("/api/clear_assignments")
|
|
async def api_clear_assignments():
|
|
await STATE.clear_all_assignments()
|
|
return {"status": "ok"}
|
|
|
|
# =============================
|
|
# Bot control endpoint
|
|
# =============================
|
|
@app.post("/api/bot_control")
|
|
async def api_bot_control(req: BotControl):
|
|
if req.action not in ("start", "stop"):
|
|
raise HTTPException(status_code=400, detail="invalid action")
|
|
STATE.bot_running = (req.action == "start")
|
|
return {"status": "ok", "bot_running": STATE.bot_running}
|
|
|
|
|
|
# ============================================================
|
|
# Reset Groups Endpoint
|
|
# ============================================================
|
|
|
|
@app.post("/api/reset_groups")
|
|
async def reset_groups():
|
|
# Load existing data
|
|
data = load_json("assigned_groups.json")
|
|
|
|
# Clear group assignments
|
|
for g in data["groups"].values():
|
|
g.clear()
|
|
|
|
# Reset friendly assigned_group
|
|
for f in data["friendly_members"]:
|
|
f["assigned_group"] = None
|
|
|
|
# Reset enemy assigned_group
|
|
for e in data["enemy_members"]:
|
|
e["assigned_group"] = None
|
|
|
|
# Save back to file
|
|
save_json("assigned_groups.json", data)
|
|
|
|
return { "success": True }
|
|
|
|
|
|
# ============================================================
|
|
# Discord Bot Setup
|
|
# ============================================================
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
|
|
enrolled_attackers = []
|
|
enemy_queue = []
|
|
active_assignments = {}
|
|
round_robin_index = 0
|
|
|
|
class HitDispatchBot(commands.Bot):
|
|
async def setup_hook(self):
|
|
await self.add_cog(
|
|
Assignments(
|
|
self,
|
|
enemy_queue=enemy_queue,
|
|
active_assignments=active_assignments,
|
|
enrolled_attackers=enrolled_attackers,
|
|
hit_check=HIT_CHECK_INTERVAL,
|
|
reassign_delay=REASSIGN_DELAY,
|
|
)
|
|
)
|
|
|
|
await self.add_cog(
|
|
HitCommands(
|
|
self,
|
|
enrolled_attackers=enrolled_attackers,
|
|
enemy_queue=enemy_queue
|
|
)
|
|
)
|
|
|
|
async def cog_check(self, ctx):
|
|
return ctx.channel.id == ALLOWED_CHANNEL_ID
|
|
|
|
bot = HitDispatchBot(command_prefix="!", intents=intents)
|
|
|
|
@bot.event
|
|
async def on_ready():
|
|
print(f"Logged in as {bot.user.name}")
|
|
|
|
TOKEN = "YOUR_DISCORD_TOKEN"
|
|
|
|
async def start_bot():
|
|
await bot.start(TOKEN)
|
|
|
|
# ============================================================
|
|
# Main Entry Point
|
|
# ============================================================
|
|
if __name__ == "__main__":
|
|
loop = asyncio.get_event_loop()
|
|
|
|
# Start Discord bot in background
|
|
loop.create_task(start_bot())
|
|
|
|
# Run FastAPI app — keeps loop alive
|
|
uvicorn.run(app, host="127.0.0.1", port=8000)
|