Code organization and refactoring
This commit is contained in:
414
main.py
414
main.py
@@ -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():
|
||||
|
||||
Reference in New Issue
Block a user