from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.responses import HTMLResponse, JSONResponse
import json, time, random, sqlite3, os
from datetime import datetime

DB_PATH = os.environ.get("DB_PATH", "/data/rr.db")
app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")
templates = Jinja2Templates(directory="templates")

# ── SCORING CONSTANTS ─────────────────────────────────────────────────────────
MAX_SCORE_PER_PLAYER = 4500   # 3 rounds × 1500 pts theoretical max
PRIZE_THRESHOLD      = 3000   # avg score that earns full 25 prize points
MAX_PRIZE_POINTS     = 25

# ── DATABASE ──────────────────────────────────────────────────────────────────
def get_db():
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    with get_db() as conn:
        conn.executescript("""
            CREATE TABLE IF NOT EXISTS sessions (
                id         INTEGER PRIMARY KEY AUTOINCREMENT,
                started_at TEXT NOT NULL,
                ended_at   TEXT,
                rounds     INTEGER DEFAULT 0
            );
            CREATE TABLE IF NOT EXISTS scores (
                id         INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id INTEGER REFERENCES sessions(id),
                player     TEXT NOT NULL,
                score      INTEGER DEFAULT 0,
                rank       INTEGER,
                dept       TEXT NOT NULL DEFAULT '',
                finished_at TEXT NOT NULL
            );
        """)
        # safe migration if dept column missing from older DB
        try:
            conn.execute("ALTER TABLE scores ADD COLUMN dept TEXT NOT NULL DEFAULT ''")
        except Exception:
            pass

init_db()

# ── GRID ──────────────────────────────────────────────────────────────────────
GRID_SIZE = 4

MODIFIER_TYPES = [
    "normal", "normal", "normal", "normal", "normal", "normal",
    "slow", "slow", "boost", "blocked"
]

def generate_grid():
    nodes = []
    for row in range(GRID_SIZE):
        for col in range(GRID_SIZE):
            nodes.append({"id": row * GRID_SIZE + col, "row": row, "col": col})

    corners = [0, 3, 12, 15]
    start = random.choice(corners)
    opposite = {0: 15, 3: 12, 12: 3, 15: 0}
    end = opposite[start]

    for attempt in range(20):
        edges = []
        for row in range(GRID_SIZE):
            for col in range(GRID_SIZE):
                nid = row * GRID_SIZE + col
                if col + 1 < GRID_SIZE:
                    edges.append({"a": nid, "b": nid + 1, "modifier": random.choice(MODIFIER_TYPES)})
                if row + 1 < GRID_SIZE:
                    edges.append({"a": nid, "b": nid + GRID_SIZE, "modifier": random.choice(MODIFIER_TYPES)})

        # Validate: start node must have at least one non-blocked neighbor
        start_edges = [e for e in edges if e["a"] == start or e["b"] == start]
        if any(e["modifier"] != "blocked" for e in start_edges):
            break
    # If all attempts failed (extremely unlikely), force one start edge open
    else:
        for e in edges:
            if e["a"] == start or e["b"] == start:
                e["modifier"] = "normal"
                break

    return {"nodes": nodes, "edges": edges, "start": start, "end": end}

def score_path(path: list, edges: list, elapsed: float, time_limit: float, end_node: int) -> dict:
    if not path or path[-1] != end_node:
        return {"points": 0, "reason": "no_path"}

    edge_map = {}
    for e in edges:
        edge_map[(e["a"], e["b"])] = e["modifier"]
        edge_map[(e["b"], e["a"])] = e["modifier"]

    hop_cost = 0
    valid = True
    for i in range(len(path) - 1):
        a, b = path[i], path[i + 1]
        mod = edge_map.get((a, b))
        if mod is None or mod == "blocked":
            valid = False
            break
        elif mod == "slow":
            hop_cost += 2
        elif mod == "boost":
            hop_cost += 0.5
        else:
            hop_cost += 1

    if not valid:
        return {"points": 0, "reason": "invalid_path", "hops": len(path) - 1}

    base = 1000
    speed_bonus = int(500 * max(0, (time_limit - elapsed) / time_limit))
    efficiency_penalty = int(max(0, hop_cost - (GRID_SIZE - 1)) * 50)
    points = max(0, base + speed_bonus - efficiency_penalty)

    return {
        "points": points,
        "reason": "ok",
        "hops": len(path) - 1,
        "hop_cost": hop_cost,
        "speed_bonus": speed_bonus,
        "efficiency_penalty": efficiency_penalty,
    }

# ── GAME STATE ────────────────────────────────────────────────────────────────
class GameState:
    def __init__(self):
        self.reset()

    def reset(self):
        self.phase = "lobby"
        self.players = {}
        self.current_round = 0
        self.total_rounds = 3
        self.grid = None
        self.round_start = 0
        self.time_limit = 30
        self.prize_threshold = PRIZE_THRESHOLD   # can be overridden by host
        self.submissions = {}
        self.session_id = None

game = GameState()

# ── CONNECTION MANAGER ────────────────────────────────────────────────────────
class ConnMgr:
    def __init__(self):
        self.players = {}
        self.hosts = []

    async def connect_player(self, pid, ws):
        await ws.accept()
        self.players[pid] = ws

    async def connect_host(self, ws):
        await ws.accept()
        self.hosts.append(ws)

    def disconnect(self, pid):
        self.players.pop(pid, None)
        game.players.pop(pid, None)

    def disconnect_host(self, ws):
        if ws in self.hosts:
            self.hosts.remove(ws)

    async def send(self, pid, data):
        ws = self.players.get(pid)
        if ws:
            try:
                await ws.send_json(data)
            except Exception:
                pass

    async def broadcast_players(self, data):
        dead = []
        for pid, ws in list(self.players.items()):
            try:
                await ws.send_json(data)
            except Exception:
                dead.append(pid)
        for d in dead:
            self.disconnect(d)

    async def broadcast_hosts(self, data):
        dead = []
        for ws in list(self.hosts):
            try:
                await ws.send_json(data)
            except Exception:
                dead.append(ws)
        for d in dead:
            self.disconnect_host(d)

    async def broadcast_all(self, data):
        await self.broadcast_players(data)
        await self.broadcast_hosts(data)

mgr = ConnMgr()

def leaderboard():
    ranked = sorted(game.players.values(), key=lambda p: p["score"], reverse=True)
    return [{"name": p.get("display", p["name"]), "score": p["score"], "dept": p.get("dept", "")} for p in ranked]

def dept_leaderboard():
    """
    Group players by dept and compute:
      - avg_score   : mean individual score across all players in the group
      - pct         : avg_score as % of MAX_SCORE_PER_PLAYER (4500 pts)
      - prize_pts   : prize points earned (max 25, anchored at PRIZE_THRESHOLD=3000)
      - player_count: number of players in the group
    Sorted by avg_score descending.
    """
    groups = {}
    for p in game.players.values():
        dept = p.get("dept", "") or "—"
        if dept not in groups:
            groups[dept] = []
        groups[dept].append(p["score"])

    result = []
    for dept, scores in groups.items():
        n = len(scores)
        avg = round(sum(scores) / n) if n else 0
        pct = round((avg / MAX_SCORE_PER_PLAYER) * 100, 1)
        prize = min(MAX_PRIZE_POINTS, round((avg / game.prize_threshold) * MAX_PRIZE_POINTS))
        result.append({
            "dept": dept,
            "avg_score": avg,
            "pct": pct,
            "prize_pts": prize,
            "player_count": n,
            "total_score": sum(scores),
        })

    result.sort(key=lambda x: x["avg_score"], reverse=True)
    # Add rank
    for i, r in enumerate(result, 1):
        r["rank"] = i
    return result

def save_session():
    now = datetime.utcnow().isoformat()
    with get_db() as conn:
        if not game.session_id:
            cur = conn.execute(
                "INSERT INTO sessions (started_at, ended_at, rounds) VALUES (?,?,?)",
                (now, now, game.current_round)
            )
            game.session_id = cur.lastrowid
        else:
            conn.execute(
                "UPDATE sessions SET ended_at=?, rounds=? WHERE id=?",
                (now, game.current_round, game.session_id)
            )
        ranked = sorted(game.players.values(), key=lambda p: p["score"], reverse=True)
        conn.execute("DELETE FROM scores WHERE session_id=?", (game.session_id,))
        for rank, p in enumerate(ranked, 1):
            conn.execute(
                "INSERT INTO scores (session_id, player, score, rank, dept, finished_at) VALUES (?,?,?,?,?,?)",
                (game.session_id, p.get("display", p["name"]), p["score"], rank, p.get("dept", ""), now)
            )

# ── ROUTES ────────────────────────────────────────────────────────────────────
@app.get("/", response_class=HTMLResponse)
async def player_page(request: Request):
    return templates.TemplateResponse("player.html", {"request": request})

@app.get("/practice", response_class=HTMLResponse)
async def practice_page(request: Request):
    return templates.TemplateResponse("practice.html", {"request": request})

@app.get("/host", response_class=HTMLResponse)
async def host_page(request: Request):
    return templates.TemplateResponse("host.html", {"request": request})

@app.get("/scores", response_class=HTMLResponse)
async def scores_page(request: Request):
    return templates.TemplateResponse("scores.html", {"request": request})

@app.get("/api/scores")
async def api_scores():
    with get_db() as conn:
        sessions = conn.execute(
            "SELECT * FROM sessions ORDER BY started_at DESC LIMIT 20"
        ).fetchall()
        result = []
        for s in sessions:
            players = conn.execute(
                "SELECT player, score, rank, dept FROM scores WHERE session_id=? ORDER BY rank",
                (s["id"],)
            ).fetchall()
            result.append({
                "id": s["id"],
                "started_at": s["started_at"],
                "rounds": s["rounds"],
                "players": [
                    {
                        "player": p["player"],
                        "score": p["score"],
                        "rank": p["rank"],
                        "dept": p["dept"] if "dept" in p.keys() else "",
                    }
                    for p in players
                ],
            })
    return JSONResponse(result)

@app.delete("/api/scores/clear")
async def clear_scores():
    with get_db() as conn:
        conn.execute("DELETE FROM scores")
        conn.execute("DELETE FROM sessions")
    return {"ok": True}

# ── PLAYER WEBSOCKET ──────────────────────────────────────────────────────────
@app.websocket("/ws/player/{pid}")
async def player_ws(ws: WebSocket, pid: str):
    await mgr.connect_player(pid, ws)
    try:
        await mgr.send(pid, {"type": "phase", "phase": game.phase})
        async for raw in ws.iter_text():
            msg = json.loads(raw)

            if msg["type"] == "join":
                name = msg["name"].strip()[:30]
                dept = msg.get("dept", "").strip()[:40]
                display = msg.get("display", name).strip()[:30]
                if not name or len(game.players) >= 4:
                    await mgr.send(pid, {"type": "error", "msg": "Game full or invalid name"})
                    continue
                game.players[pid] = {
                    "name": name,
                    "display": display,
                    "dept": dept,
                    "score": 0,
                    "submitted": False,
                }
                await mgr.send(pid, {"type": "joined", "name": name, "display": display, "dept": dept})
                await mgr.broadcast_hosts({
                    "type": "lobby_update",
                    "players": [p.get("display", p["name"]) for p in game.players.values()],
                    "count": len(game.players),
                })

            elif msg["type"] == "submit_path":
                if game.phase != "playing":
                    continue
                player = game.players.get(pid)
                if not player or player["submitted"]:
                    continue
                elapsed = time.time() - game.round_start
                path = msg.get("path", [])
                result = score_path(path, game.grid["edges"], elapsed, game.time_limit, game.grid["end"])
                player["score"] += result["points"]
                player["submitted"] = True
                game.submissions[pid] = {
                    "name": player.get("display", player["name"]),
                    "path": path,
                    "points": result["points"],
                    "elapsed": round(elapsed, 2),
                    "reason": result["reason"],
                }
                await mgr.send(pid, {
                    "type": "path_result",
                    "points": result["points"],
                    "total": player["score"],
                    "reason": result["reason"],
                    "elapsed": round(elapsed, 2),
                })
                await mgr.broadcast_hosts({
                    "type": "submission_update",
                    "submitted": sum(1 for p in game.players.values() if p["submitted"]),
                    "total": len(game.players),
                    "submissions": list(game.submissions.values()),
                })

    except WebSocketDisconnect:
        mgr.disconnect(pid)
        await mgr.broadcast_hosts({
            "type": "lobby_update",
            "players": [p.get("display", p["name"]) for p in game.players.values()],
            "count": len(game.players),
        })

# ── HOST WEBSOCKET ────────────────────────────────────────────────────────────
@app.websocket("/ws/host")
async def host_ws(ws: WebSocket):
    await mgr.connect_host(ws)
    try:
        await ws.send_json({
            "type": "lobby_update",
            "players": [p.get("display", p["name"]) for p in game.players.values()],
            "count": len(game.players),
        })

        async for raw in ws.iter_text():
            msg = json.loads(raw)

            if msg["type"] == "start_round":
                # accept config params from host on first round
                if "total_rounds" in msg:
                    game.total_rounds = int(msg["total_rounds"])
                if "time_limit" in msg:
                    game.time_limit = int(msg["time_limit"])
                if "prize_threshold" in msg:
                    pt = int(msg["prize_threshold"])
                    game.prize_threshold = max(500, min(4500, pt))  # clamp to valid range
                game.grid = generate_grid()
                game.phase = "playing"
                game.round_start = time.time()
                game.current_round += 1
                game.submissions = {}
                for p in game.players.values():
                    p["submitted"] = False
                await mgr.broadcast_all({
                    "type": "round_start",
                    "round": game.current_round,
                    "total_rounds": game.total_rounds,
                    "grid": game.grid,
                    "time_limit": game.time_limit,
                })

            elif msg["type"] == "end_round":
                game.phase = "results"
                lb = leaderboard()
                dept_lb = dept_leaderboard()
                save_session()
                await mgr.broadcast_all({
                    "type": "round_results",
                    "round": game.current_round,
                    "grid": game.grid,
                    "submissions": list(game.submissions.values()),
                    "leaderboard": lb,
                    "dept_leaderboard": dept_lb,
                    "is_final": game.current_round >= game.total_rounds,
                    "max_score": MAX_SCORE_PER_PLAYER,
                    "prize_threshold": PRIZE_THRESHOLD,
                    "max_prize_pts": MAX_PRIZE_POINTS,
                })

            elif msg["type"] == "reset":
                await mgr.broadcast_players({"type": "kicked"})
                game.reset()
                await mgr.broadcast_hosts({"type": "lobby_update", "players": [], "count": 0})

    except WebSocketDisconnect:
        mgr.disconnect_host(ws)