from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
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()
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)