# main.py (updated) import discord from discord.ext import commands from config import ALLOWED_CHANNEL_ID, HIT_CHECK_INTERVAL, REASSIGN_DELAY, DISCORD_TOKEN 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 from services.bot_assignment import BotAssignmentManager # ============================================================ # 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" class DiscordMappingRequest(BaseModel): torn_id: int discord_id: int # ================================ # 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) # Return members list for frontend (already in STATE from populate_friendly) members = [m.model_dump() for m in STATE.friendly.values()] return {"status": "friendly populated", "id": data.faction_id, "members": members} @app.post("/api/populate_enemy") async def api_populate_enemy(data: FactionRequest): await populate_enemy(data.faction_id) # Return members list for frontend (already in STATE from populate_enemy) members = [m.model_dump() for m in STATE.enemy.values()] return {"status": "enemy populated", "id": data.faction_id, "members": members} # ----------------------------- # 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.get("/api/bot_status") async def api_bot_status(): """Get current bot status""" active_count = len(assignment_manager.active_targets) if assignment_manager else 0 return { "bot_running": STATE.bot_running, "active_assignments": active_count, "discord_mappings_count": len(assignment_manager.discord_mapping) if assignment_manager else 0 } @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") # Start or stop the assignment manager if assignment_manager: if req.action == "start": await assignment_manager.start() else: await assignment_manager.stop() return {"status": "ok", "bot_running": STATE.bot_running} # ============================= # Discord Mapping endpoints # ============================= @app.get("/api/discord_mappings") async def get_discord_mappings(): """Get all Torn ID to Discord ID mappings""" path = Path("data/discord_mapping.json") if not path.exists(): return {"mappings": {}} with open(path, "r", encoding="utf-8") as f: data = json.load(f) return {"mappings": data.get("mappings", {})} @app.post("/api/discord_mapping") async def add_discord_mapping(req: DiscordMappingRequest): """Add or update a Torn ID to Discord ID mapping""" path = Path("data/discord_mapping.json") # Load existing mappings if path.exists(): with open(path, "r", encoding="utf-8") as f: data = json.load(f) else: data = {"comment": "Map Torn player IDs to Discord user IDs", "mappings": {}} # Update mapping data["mappings"][str(req.torn_id)] = str(req.discord_id) # Save back with open(path, "w", encoding="utf-8") as f: json.dump(data, f, indent=2) # Reload in assignment manager if assignment_manager: assignment_manager.load_discord_mapping() return {"status": "ok", "torn_id": req.torn_id, "discord_id": req.discord_id} @app.delete("/api/discord_mapping/{torn_id}") async def remove_discord_mapping(torn_id: int): """Remove a Discord mapping""" path = Path("data/discord_mapping.json") if not path.exists(): raise HTTPException(status_code=404, detail="No mappings found") with open(path, "r", encoding="utf-8") as f: data = json.load(f) # Remove mapping if str(torn_id) in data.get("mappings", {}): del data["mappings"][str(torn_id)] # Save back with open(path, "w", encoding="utf-8") as f: json.dump(data, f, indent=2) # Reload in assignment manager if assignment_manager: assignment_manager.load_discord_mapping() return {"status": "ok", "torn_id": torn_id} else: raise HTTPException(status_code=404, detail="Mapping not found") # ============================================================ # Reset Groups Endpoint # ============================================================ @app.post("/api/reset_groups") async def reset_groups(): # Clear all assignments in server state await STATE.clear_all_assignments() 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) # Initialize bot assignment manager assignment_manager = None @bot.event async def on_ready(): global assignment_manager print(f"✓ Discord bot logged in as {bot.user.name} (ID: {bot.user.id})") print(f"✓ Bot is in {len(bot.guilds)} server(s)") # Initialize assignment manager assignment_manager = BotAssignmentManager(bot) print("✓ Bot assignment manager initialized") async def start_bot(): try: print("Starting Discord bot...") await bot.start(DISCORD_TOKEN) except discord.LoginFailure: print("ERROR: Invalid Discord token! Please set DISCORD_TOKEN in config.py") except Exception as e: print(f"ERROR starting Discord bot: {e}") # ============================================================ # Main Entry Point # ============================================================ async def main(): # Start Discord bot in background bot_task = asyncio.create_task(start_bot()) # Configure and run FastAPI server config = uvicorn.Config(app, host="127.0.0.1", port=8000, log_level="info") server = uvicorn.Server(config) await server.serve() if __name__ == "__main__": asyncio.run(main())