server.py

"""AIguessr — web dashboard server.

Serves the React dashboard and proxies commands to the geo.py TCP server.
Handles autonomous play via Anthropic/OpenRouter APIs with parallel agents.

Usage:
    python server.py                    — start dashboard on http://127.0.0.1:8080
    python server.py --port 9090        — custom port
"""

import asyncio
import json
import os
import socket
import subprocess
import sys
from pathlib import Path

from dotenv import load_dotenv
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

from brain import GeoBrain
from browser_detect import detect_browsers, launch_browser, is_cdp_available
from config import load_config, save_config
from match_tracker import MatchTracker

load_dotenv()

APP_DIR = Path(__file__).parent
SCREENSHOTS_DIR = APP_DIR / "screenshots"
DASHBOARD_DIST = APP_DIR / "dashboard" / "dist"
GEO_HOST = "127.0.0.1"
GEO_PORT = 5555

app = FastAPI(title="AIguessr Dashboard")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

ws_clients: list[WebSocket] = []
active_brain: GeoBrain | None = None
play_task: asyncio.Task | None = None
browser_process: subprocess.Popen | None = None
active_match: MatchTracker | None = None


def send_geo_command(cmd: str, timeout: float = 30) -> str:
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(timeout)
        s.connect((GEO_HOST, GEO_PORT))
        s.sendall((cmd + "\n").encode())
        data = b""
        while True:
            chunk = s.recv(4096)
            if not chunk:
                break
            data += chunk
        s.close()
        return data.decode().strip()
    except ConnectionRefusedError:
        return "ERROR geo.py server not running (start with: python geo.py serve)"
    except socket.timeout:
        return "ERROR command timed out"
    except Exception as e:
        return f"ERROR {e}"


async def broadcast(message: dict):
    dead = []
    for ws in ws_clients:
        try:
            await ws.send_json(message)
        except Exception:
            dead.append(ws)
    for ws in dead:
        ws_clients.remove(ws)


async def brain_event_handler(event_type: str, data: dict):
    await broadcast({"type": "brain_event", "event": event_type, "data": data})
    if event_type == "overlay_recommendation":
        guess = data.get("guess", {})
        overlay_data = json.dumps({
            "country": guess.get("country", ""),
            "country_navigation": guess.get("country_navigation", ""),
            "region": guess.get("region", ""),
            "region_position": guess.get("region_position", ""),
            "pin_placement": guess.get("pin_placement", ""),
            "confidence": guess.get("confidence", "low"),
        })
        await asyncio.to_thread(send_geo_command, "overlay_inject")
        await asyncio.to_thread(send_geo_command, f"overlay_update {overlay_data}")


class CommandRequest(BaseModel):
    command: str


class SettingsRequest(BaseModel):
    provider: str
    apiKey: str = ""
    model: str = ""


class PlayRequest(BaseModel):
    mode: str = "classic"
    play_style: str = "aggressive"


@app.post("/api/command")
async def run_command(req: CommandRequest):
    result = await asyncio.to_thread(send_geo_command, req.command)
    await broadcast({"type": "log", "message": result, "level": "error" if result.startswith("ERROR") else "success"})
    return {"result": result}


@app.get("/api/state")
async def get_state():
    result = await asyncio.to_thread(send_geo_command, "state")
    try:
        return json.loads(result)
    except json.JSONDecodeError:
        return {"error": result}


@app.get("/api/screenshots/{filename}")
async def get_screenshot(filename: str):
    path = SCREENSHOTS_DIR / filename
    if not path.exists() or not path.is_file():
        return JSONResponse({"error": "not found"}, status_code=404)
    if not str(path.resolve()).startswith(str(SCREENSHOTS_DIR.resolve())):
        return JSONResponse({"error": "forbidden"}, status_code=403)
    return FileResponse(path, media_type="image/png")


@app.get("/api/health")
async def health():
    geo_ok = await asyncio.to_thread(send_geo_command, "state")
    return {
        "server": "ok",
        "geo_server": "ok" if not geo_ok.startswith("ERROR") else "disconnected",
        "brain_active": active_brain is not None and active_brain.playing,
    }


@app.post("/api/settings")
async def save_settings(req: SettingsRequest):
    env_path = APP_DIR / ".env"
    lines = []
    if env_path.exists():
        lines = env_path.read_text().splitlines()

    env_vars = {
        "GEOGUESSR_PROVIDER": req.provider,
        "GEOGUESSR_MODEL": req.model,
    }
    if req.apiKey:
        if req.provider == "anthropic":
            env_vars["ANTHROPIC_API_KEY"] = req.apiKey
        elif req.provider == "openrouter":
            env_vars["OPENROUTER_API_KEY"] = req.apiKey

    updated_keys = set()
    new_lines = []
    for line in lines:
        key = line.split("=", 1)[0].strip() if "=" in line else ""
        if key in env_vars:
            new_lines.append(f"{key}={env_vars[key]}")
            updated_keys.add(key)
        else:
            new_lines.append(line)

    for key, val in env_vars.items():
        if key not in updated_keys:
            new_lines.append(f"{key}={val}")

    env_path.write_text("\n".join(new_lines) + "\n")

    for key, val in env_vars.items():
        os.environ[key] = val

    return {"status": "saved"}


@app.post("/api/play/start")
async def start_play(req: PlayRequest):
    global active_brain, play_task

    if active_brain and active_brain.playing:
        return {"status": "already_playing"}

    provider = os.environ.get("GEOGUESSR_PROVIDER", "anthropic")
    model = os.environ.get("GEOGUESSR_MODEL", "claude-sonnet-4-6")

    if provider == "anthropic":
        api_key = os.environ.get("ANTHROPIC_API_KEY", "")
    elif provider == "openrouter":
        api_key = os.environ.get("OPENROUTER_API_KEY", "")
    else:
        return {"status": "error", "message": "Claude Code mode doesn't support autonomous play. Use Anthropic API or OpenRouter."}

    if not api_key:
        return {"status": "error", "message": f"No API key configured for {provider}. Go to Settings."}

    active_brain = GeoBrain(provider, api_key, model)
    active_brain.on_event(brain_event_handler)

    duels = req.mode == "duels"
    play_task = asyncio.create_task(active_brain.auto_play(duels=duels, play_style=req.play_style))
    return {"status": "started", "mode": req.mode, "play_style": req.play_style}


@app.post("/api/play/stop")
async def stop_play():
    global active_brain, play_task

    if active_brain:
        active_brain.stop()
    if play_task and not play_task.done():
        play_task.cancel()
        try:
            await play_task
        except asyncio.CancelledError:
            pass

    active_brain = None
    play_task = None
    return {"status": "stopped"}


@app.get("/api/play/status")
async def play_status():
    if active_brain and active_brain.playing:
        return {"status": "playing", "mode": "duels" if active_brain.duels_mode else "classic"}
    return {"status": "idle"}


@app.post("/api/duels/start")
async def duels_start():
    return await start_play(PlayRequest(mode="duels"))


@app.post("/api/duels/stop")
async def duels_stop():
    return await stop_play()


@app.get("/api/browsers")
async def list_browsers():
    browsers = detect_browsers()
    cdp_ok = is_cdp_available()
    return {"browsers": browsers, "cdp_connected": cdp_ok}


@app.post("/api/browser/launch")
async def launch_browser_endpoint(req: dict):
    global browser_process
    path = req.get("path", "")
    if not path:
        return JSONResponse({"error": "No browser path provided"}, status_code=400)
    browser_process = launch_browser(path)
    return {"status": "launched", "pid": browser_process.pid}


@app.get("/api/browser/status")
async def browser_status():
    cdp_ok = is_cdp_available()
    running = browser_process is not None and browser_process.poll() is None
    return {"cdp_connected": cdp_ok, "browser_running": running}


@app.get("/api/config")
async def get_config():
    return load_config()


@app.put("/api/config")
async def update_config(req: dict):
    config = load_config()
    config.update(req)
    config["first_run"] = False
    save_config(config)
    return {"status": "saved"}


@app.get("/api/match/stats")
async def match_stats():
    return MatchTracker.get_stats()


@app.get("/api/match/history")
async def match_history():
    return {"matches": MatchTracker.get_history()}


@app.get("/api/match/state")
async def match_state():
    if active_match and active_match.active:
        return active_match.get_state()
    return {"active": False}


@app.get("/api/match/{match_id}")
async def match_detail(match_id: int):
    match = MatchTracker.get_match(match_id)
    if not match:
        return JSONResponse({"error": "not found"}, status_code=404)
    return match


@app.websocket("/ws")
async def websocket_endpoint(ws: WebSocket):
    await ws.accept()
    ws_clients.append(ws)

    geo_result = await asyncio.to_thread(send_geo_command, "state")
    geo_ok = not geo_result.startswith("ERROR")
    await ws.send_json({
        "type": "connection",
        "status": "connected" if geo_ok else "disconnected",
    })

    try:
        while True:
            data = await ws.receive_text()
            try:
                msg = json.loads(data)
                if msg.get("type") == "command":
                    result = await asyncio.to_thread(send_geo_command, msg["command"])
                    await ws.send_json({"type": "log", "message": result})
            except json.JSONDecodeError:
                pass
    except WebSocketDisconnect:
        pass
    finally:
        if ws in ws_clients:
            ws_clients.remove(ws)


if DASHBOARD_DIST.exists():
    app.mount("/assets", StaticFiles(directory=str(DASHBOARD_DIST / "assets")), name="assets")

    @app.get("/{path:path}")
    async def serve_spa(path: str):
        file_path = DASHBOARD_DIST / path
        if file_path.exists() and file_path.is_file():
            return FileResponse(file_path)
        return FileResponse(DASHBOARD_DIST / "index.html")


if __name__ == "__main__":
    import uvicorn

    port = 8080
    if "--port" in sys.argv:
        idx = sys.argv.index("--port")
        port = int(sys.argv[idx + 1])

    print(f"Dashboard: http://127.0.0.1:{port}")
    print(f"Proxying commands to geo.py on {GEO_HOST}:{GEO_PORT}")
    uvicorn.run(app, host="127.0.0.1", port=port, log_level="info")