Code organization and refactoring

This commit is contained in:
2026-01-26 16:13:29 -05:00
parent 9c3b4c8335
commit a64f9a3d74
23 changed files with 495 additions and 396 deletions

414
main.py
View File

@@ -1,4 +1,4 @@
# main.py (updated)
# main.py
import discord
from discord.ext import commands
from config import ALLOWED_CHANNEL_ID, HIT_CHECK_INTERVAL, REASSIGN_DELAY, DISCORD_TOKEN
@@ -7,413 +7,30 @@ 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 import FastAPI
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
# Import routers
from routers import pages, factions, assignments, discord_mappings, config
from routers import bot as bot_router
# ============================================================
# 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})
@app.get("/config", response_class=HTMLResponse)
async def config_page(request: Request):
return templates.TemplateResponse("config.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
class ConfigUpdateRequest(BaseModel):
key: str
value: Optional[str | 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")
# =============================
# Config endpoints
# =============================
def reload_config_from_file():
"""Reload config values from JSON into module globals"""
path = Path("data/config.json")
if not path.exists():
return
try:
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
# Update config module globals
import config
for key, value in data.get("config", {}).items():
if hasattr(config, key):
setattr(config, key, value)
except Exception as e:
print(f"Error reloading config from file: {e}")
@app.get("/api/config")
async def get_config():
"""Get all config values (with sensitive values masked)"""
import config
path = Path("data/config.json")
if not path.exists():
# Return defaults from config.py (masked)
config_values = {
"TORN_API_KEY": config.TORN_API_KEY,
"FFSCOUTER_KEY": config.FFSCOUTER_KEY,
"DISCORD_TOKEN": config.DISCORD_TOKEN,
"ALLOWED_CHANNEL_ID": config.ALLOWED_CHANNEL_ID,
"POLL_INTERVAL": config.POLL_INTERVAL,
"HIT_CHECK_INTERVAL": config.HIT_CHECK_INTERVAL,
"REASSIGN_DELAY": config.REASSIGN_DELAY,
"ASSIGNMENT_TIMEOUT": config.ASSIGNMENT_TIMEOUT,
"ASSIGNMENT_REMINDER": config.ASSIGNMENT_REMINDER
}
else:
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
config_values = data.get("config", {})
# Mask sensitive values
masked_config = config_values.copy()
sensitive = ["TORN_API_KEY", "FFSCOUTER_KEY", "DISCORD_TOKEN"]
for key in sensitive:
if key in masked_config and masked_config[key]:
val = str(masked_config[key])
masked_config[key] = "****" + val[-4:] if len(val) > 4 else "****"
return {"config": masked_config, "sensitive_fields": sensitive}
@app.post("/api/config")
async def update_config(req: ConfigUpdateRequest):
"""Update a single config value"""
path = Path("data/config.json")
# Load existing or create from current config
if path.exists():
with open(path, "r", encoding="utf-8") as f:
data = json.load(f)
else:
import config
data = {
"comment": "Application configuration settings",
"config": {
"TORN_API_KEY": config.TORN_API_KEY,
"FFSCOUTER_KEY": config.FFSCOUTER_KEY,
"DISCORD_TOKEN": config.DISCORD_TOKEN,
"ALLOWED_CHANNEL_ID": config.ALLOWED_CHANNEL_ID,
"POLL_INTERVAL": config.POLL_INTERVAL,
"HIT_CHECK_INTERVAL": config.HIT_CHECK_INTERVAL,
"REASSIGN_DELAY": config.REASSIGN_DELAY,
"ASSIGNMENT_TIMEOUT": config.ASSIGNMENT_TIMEOUT,
"ASSIGNMENT_REMINDER": config.ASSIGNMENT_REMINDER
}
}
# Validate key exists
if req.key not in data["config"]:
raise HTTPException(status_code=400, detail="Invalid config key")
# Update value
data["config"][req.key] = req.value
# Save to file
with open(path, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2)
# Reload config in memory
reload_config_from_file()
return {"status": "ok", "key": req.key}
# ============================================================
# 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}
# Include all routers
app.include_router(pages.router)
app.include_router(factions.router)
app.include_router(assignments.router)
app.include_router(bot_router.router)
app.include_router(discord_mappings.router)
app.include_router(config.router)
# ============================================================
# Discord Bot Setup
@@ -463,6 +80,11 @@ async def on_ready():
# Initialize assignment manager
assignment_manager = BotAssignmentManager(bot)
# Set assignment manager in routers
bot_router.set_assignment_manager(assignment_manager)
discord_mappings.set_assignment_manager(assignment_manager)
print("Bot assignment manager initialized")
async def start_bot():