from __future__ import annotations

import asyncio
import datetime as dt
import json
import os
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional

import jwt
import torch
from fastapi import Depends, FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from passlib.context import CryptContext
from pydantic import BaseModel
from sqlmodel import Field, Session, SQLModel, create_engine, select
from torch.serialization import add_safe_globals

from ..game import GameConfig
from ..sdk import Action, ActionType, CatanEnv
from ..ml.agents import PolicyAgent, RandomAgent
from ..ml.encoding import encode_action, encode_observation
from ..ml.selfplay import ActionScoringNetwork, PPOConfig

BASE_DIR = Path(__file__).resolve().parent
PROJECT_DIR = BASE_DIR.parent.parent
STATIC_DIR = BASE_DIR / "static"
DATA_DIR = Path(os.environ.get("CATAN_DATA_DIR", PROJECT_DIR / "data"))
MODELS_DIR = Path(os.environ.get("CATAN_MODELS_DIR", PROJECT_DIR / "models"))
JWT_SECRET = os.environ.get("CATAN_JWT_SECRET", "change-me")
JWT_ALG = "HS256"
TOKEN_EXPIRY_DAYS = 7

DATA_DIR.mkdir(parents=True, exist_ok=True)
MODELS_DIR.mkdir(parents=True, exist_ok=True)

DB_URL = f"sqlite:///{DATA_DIR / 'catan.db'}"
engine = create_engine(DB_URL, echo=False)

pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")


class User(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    username: str = Field(index=True, unique=True)
    password_hash: str
    created_at: dt.datetime = Field(default_factory=lambda: dt.datetime.now(dt.timezone.utc))


class GameRecord(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    game_id: str = Field(index=True)
    created_at: dt.datetime = Field(default_factory=lambda: dt.datetime.now(dt.timezone.utc))
    finished_at: Optional[dt.datetime] = None
    winner: Optional[str] = None
    players: str
    max_players: int
    total_turns: int


class RegisterRequest(BaseModel):
    username: str
    password: str


class LoginRequest(BaseModel):
    username: str
    password: str


class CreateGameRequest(BaseModel):
    name: str
    max_players: int


class AddAIRequest(BaseModel):
    ai_type: str
    model_name: Optional[str] = None


class ActionRequest(BaseModel):
    type: str
    payload: Dict[str, object] = Field(default_factory=dict)


@dataclass
class PlayerSlot:
    slot_id: int
    name: Optional[str] = None
    user_id: Optional[int] = None
    is_ai: bool = False
    ai_kind: Optional[str] = None
    ai_model: Optional[str] = None
    ready: bool = False
    color: Optional[str] = None

    def to_dict(self) -> Dict[str, object]:
        return {
            "slot_id": self.slot_id,
            "name": self.name,
            "user_id": self.user_id,
            "is_ai": self.is_ai,
            "ai_kind": self.ai_kind,
            "ai_model": self.ai_model,
            "ready": self.ready,
            "color": self.color,
        }


@dataclass
class GameSession:
    id: str
    name: str
    max_players: int
    created_by: str
    created_at: dt.datetime
    status: str = "lobby"
    players: List[PlayerSlot] = field(default_factory=list)
    env: Optional[CatanEnv] = None
    ai_agents: Dict[str, object] = field(default_factory=dict)
    action_log: List[Dict[str, object]] = field(default_factory=list)
    winner: Optional[str] = None
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    connections: Dict[WebSocket, str] = field(default_factory=dict)

    def to_lobby(self) -> Dict[str, object]:
        return {
            "id": self.id,
            "name": self.name,
            "status": self.status,
            "max_players": self.max_players,
            "created_by": self.created_by,
            "created_at": self.created_at.isoformat(),
            "players": [slot.to_dict() for slot in self.players],
        }


class ModelRegistry:
    def __init__(self) -> None:
        self._cache: Dict[str, ActionScoringNetwork] = {}
        add_safe_globals([PPOConfig, GameConfig])

    def _load_actor(self, path: Path) -> ActionScoringNetwork:
        state = torch.load(path, map_location="cpu", weights_only=False)
        game_cfg = state.get("game_config")
        cfg = state.get("config")
        if game_cfg is None or cfg is None:
            raise ValueError("Invalid checkpoint format.")
        env = CatanEnv(game_cfg)
        obs_dim = encode_observation(env.observe()).shape[0]
        action_dim = encode_action(Action(ActionType.END_TURN, {})).shape[0]
        actor = ActionScoringNetwork(obs_dim, action_dim, cfg.hidden_sizes)
        actor.load_state_dict(state["actor"])
        actor.eval()
        return actor

    def get_actor(self, model_name: str) -> ActionScoringNetwork:
        path = MODELS_DIR / model_name
        if not path.exists():
            raise FileNotFoundError(f"Model not found: {model_name}")
        if model_name not in self._cache:
            self._cache[model_name] = self._load_actor(path)
        return self._cache[model_name]


model_registry = ModelRegistry()


class GameManager:
    def __init__(self) -> None:
        self.games: Dict[str, GameSession] = {}

    def create_game(self, name: str, max_players: int, creator: str) -> GameSession:
        if max_players < 3 or max_players > 4:
            raise ValueError("max_players must be between 3 and 4.")
        gid = str(uuid.uuid4())
        session = GameSession(
            id=gid,
            name=name,
            max_players=max_players,
            created_by=creator,
            created_at=dt.datetime.now(dt.timezone.utc),
            players=[PlayerSlot(slot_id=i + 1) for i in range(max_players)],
        )
        self.games[gid] = session
        return session

    def list_games(self) -> List[Dict[str, object]]:
        return [session.to_lobby() for session in self.games.values()]

    def get(self, game_id: str) -> GameSession:
        session = self.games.get(game_id)
        if not session:
            raise KeyError(game_id)
        return session


manager = GameManager()


app = FastAPI(title="Catan Arena")
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


@app.on_event("startup")
def _startup() -> None:
    SQLModel.metadata.create_all(engine)


@app.get("/")
def index() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


def _hash_password(password: str) -> str:
    return pwd_context.hash(password)


def _verify_password(password: str, password_hash: str) -> bool:
    return pwd_context.verify(password, password_hash)


def _create_token(user: User) -> str:
    now = dt.datetime.now(dt.timezone.utc)
    payload = {
        "sub": str(user.id),
        "username": user.username,
        "exp": now + dt.timedelta(days=TOKEN_EXPIRY_DAYS),
    }
    return jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALG)


def _get_user_from_token(token: str) -> User:
    try:
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALG])
    except jwt.PyJWTError as exc:
        raise HTTPException(status_code=401, detail="Invalid token") from exc
    user_id = payload.get("sub")
    if not user_id:
        raise HTTPException(status_code=401, detail="Invalid token")
    with Session(engine) as session:
        user = session.get(User, int(user_id))
        if not user:
            raise HTTPException(status_code=401, detail="Unknown user")
        return user


def require_user(request: Request) -> User:
    auth = request.headers.get("authorization")
    if not auth or not auth.lower().startswith("bearer "):
        raise HTTPException(status_code=401, detail="Missing token")
    token = auth.split(" ", 1)[1].strip()
    return _get_user_from_token(token)


def user_from_ws(websocket: WebSocket) -> User:
    token = websocket.query_params.get("token")
    if not token:
        raise HTTPException(status_code=401, detail="Missing token")
    return _get_user_from_token(token)


def _serialize_action(action: Action) -> Dict[str, object]:
    return {"type": action.type.value, "payload": action.payload}


def _serialize_legal_actions(env: CatanEnv) -> List[Dict[str, object]]:
    return [_serialize_action(action) for action in env.legal_actions()]


def _build_state(session: GameSession, username: Optional[str] = None) -> Dict[str, object]:
    if not session.env:
        return {
            "id": session.id,
            "status": session.status,
            "players": [slot.to_dict() for slot in session.players],
            "game": None,
            "board": None,
            "history": session.action_log,
        }
    obs = session.env.observe()
    game = obs["game"]
    if username:
        masked_players = {}
        for name, pdata in game["players"].items():
            pdata_copy = dict(pdata)
            if name != username:
                resources = pdata_copy.get("resources", {})
                total = sum(resources.values()) if isinstance(resources, dict) else 0
                pdata_copy["resources"] = {"hidden": total}
            masked_players[name] = pdata_copy
        game = dict(game)
        game["players"] = masked_players
    legal_actions: List[Dict[str, object]] = []
    if username and session.env:
        current = session.env.game.current_player.name
        if username == current or username in session.env.game.pending_discards:
            legal_actions = _serialize_legal_actions(session.env)
    return {
        "id": session.id,
        "status": session.status,
        "players": [slot.to_dict() for slot in session.players],
        "game": game,
        "board": obs["board"],
        "history": session.action_log[-200:],
        "legal_actions": legal_actions,
    }


def _add_log(session: GameSession, actor: str, action: Action) -> None:
    session.action_log.append(
        {
            "ts": dt.datetime.now(dt.timezone.utc).isoformat(),
            "player": actor,
            "type": action.type.value,
            "payload": action.payload,
        }
    )


def _list_models() -> List[str]:
    if not MODELS_DIR.exists():
        return []
    return sorted([p.name for p in MODELS_DIR.glob("*.pt")])


def _build_ai_name(session: GameSession, base: str) -> str:
    existing = {slot.name for slot in session.players if slot.name}
    idx = 1
    name = f"{base}-{idx}"
    while name in existing:
        idx += 1
        name = f"{base}-{idx}"
    return name


async def _broadcast(session: GameSession) -> None:
    to_remove: List[WebSocket] = []
    for ws, username in session.connections.items():
        try:
            payload = json.dumps({"type": "state", "data": _build_state(session, username)})
            await ws.send_text(payload)
        except Exception:
            to_remove.append(ws)
    for ws in to_remove:
        session.connections.pop(ws, None)


def _build_agent(slot: PlayerSlot) -> object:
    if slot.ai_kind == "random":
        return RandomAgent()
    if slot.ai_kind == "model" and slot.ai_model:
        actor = model_registry.get_actor(slot.ai_model)

        class ActorPolicy:
            def __init__(self, model: ActionScoringNetwork) -> None:
                self.model = model

            def __call__(self, obs_tensor, action_vectors):
                with torch.no_grad():
                    return self.model(obs_tensor, action_vectors)

        return PolicyAgent(ActorPolicy(actor), device="cpu", stochastic=True)
    raise ValueError("Unknown AI kind")


def _start_game(session: GameSession) -> None:
    names = [slot.name for slot in session.players if slot.name]
    if len(names) < 3:
        raise ValueError("Not enough players to start.")
    colors = ["red", "blue", "orange", "white"]
    config = GameConfig(player_names=names, colors=colors[: len(names)])
    session.env = CatanEnv(config)
    for slot, color in zip(session.players, colors):
        if slot.name:
            slot.color = color
    session.status = "running"
    session.ai_agents.clear()
    for slot in session.players:
        if slot.is_ai and slot.name:
            session.ai_agents[slot.name] = _build_agent(slot)


def _finish_game(session: GameSession) -> None:
    if not session.env:
        return
    session.status = "finished"
    session.winner = session.env.game.winner
    with Session(engine) as db:
        record = GameRecord(
            game_id=session.id,
            created_at=session.created_at,
            finished_at=dt.datetime.now(dt.timezone.utc),
            winner=session.winner,
            players=",".join([slot.name or "" for slot in session.players]),
            max_players=session.max_players,
            total_turns=len(session.action_log),
        )
        db.add(record)
        db.commit()


async def _run_ai_turns(session: GameSession) -> None:
    if not session.env:
        return
    async with session.lock:
        while session.status == "running":
            current = session.env.game.current_player.name
            if current not in session.ai_agents:
                break
            agent = session.ai_agents[current]
            legal = session.env.legal_actions()
            if not legal:
                break
            action = agent.choose_action(session.env, legal)
            _add_log(session, current, action)
            _, _, done, info = session.env.step(action)
            if info.get("invalid"):
                session.action_log.append(
                    {
                        "ts": dt.datetime.now(dt.timezone.utc).isoformat(),
                        "player": current,
                        "type": "invalid",
                        "payload": {"error": info.get("error")},
                    }
                )
                break
            if done:
                _finish_game(session)
                break
        await _broadcast(session)


@app.post("/api/register")
def register(payload: RegisterRequest):
    with Session(engine) as session:
        existing = session.exec(select(User).where(User.username == payload.username)).first()
        if existing:
            raise HTTPException(status_code=400, detail="Username already exists")
        user = User(username=payload.username, password_hash=_hash_password(payload.password))
        session.add(user)
        session.commit()
        session.refresh(user)
        token = _create_token(user)
        return {"token": token, "username": user.username}


@app.post("/api/login")
def login(payload: LoginRequest):
    with Session(engine) as session:
        user = session.exec(select(User).where(User.username == payload.username)).first()
        if not user or not _verify_password(payload.password, user.password_hash):
            raise HTTPException(status_code=401, detail="Invalid credentials")
        token = _create_token(user)
        return {"token": token, "username": user.username}


@app.get("/api/me")
def me(user: User = Depends(require_user)):
    return {"id": user.id, "username": user.username}


@app.get("/api/models")
def models(user: User = Depends(require_user)):
    return {"models": _list_models()}


@app.get("/api/lobby")
def lobby(user: User = Depends(require_user)):
    return {"games": manager.list_games()}


@app.post("/api/games")
def create_game(payload: CreateGameRequest, user: User = Depends(require_user)):
    try:
        session = manager.create_game(payload.name, payload.max_players, user.username)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return session.to_lobby()


@app.get("/api/games/{game_id}")
def game_state(game_id: str, user: User = Depends(require_user)):
    try:
        session = manager.get(game_id)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="Game not found") from exc
    return _build_state(session, user.username)


@app.post("/api/games/{game_id}/join")
def join_game(game_id: str, user: User = Depends(require_user)):
    try:
        session = manager.get(game_id)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="Game not found") from exc
    for slot in session.players:
        if slot.user_id == user.id:
            return session.to_lobby()
    open_slot = next((slot for slot in session.players if slot.name is None), None)
    if not open_slot:
        raise HTTPException(status_code=400, detail="No available slots")
    open_slot.name = user.username
    open_slot.user_id = user.id
    open_slot.ready = True
    return session.to_lobby()


@app.post("/api/games/{game_id}/leave")
def leave_game(game_id: str, user: User = Depends(require_user)):
    try:
        session = manager.get(game_id)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="Game not found") from exc
    for slot in session.players:
        if slot.user_id == user.id:
            slot.name = None
            slot.user_id = None
            slot.is_ai = False
            slot.ai_kind = None
            slot.ai_model = None
            slot.ready = False
            slot.color = None
    return session.to_lobby()


@app.post("/api/games/{game_id}/add_ai")
def add_ai(game_id: str, payload: AddAIRequest, user: User = Depends(require_user)):
    try:
        session = manager.get(game_id)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="Game not found") from exc
    open_slot = next((slot for slot in session.players if slot.name is None), None)
    if not open_slot:
        raise HTTPException(status_code=400, detail="No available slots")
    ai_type = payload.ai_type.lower()
    if ai_type == "random":
        open_slot.name = _build_ai_name(session, "Random")
        open_slot.is_ai = True
        open_slot.ai_kind = "random"
        open_slot.ready = True
    elif ai_type == "model":
        if not payload.model_name:
            raise HTTPException(status_code=400, detail="Model name required")
        if payload.model_name not in _list_models():
            raise HTTPException(status_code=404, detail="Model not found")
        open_slot.name = _build_ai_name(session, "Model")
        open_slot.is_ai = True
        open_slot.ai_kind = "model"
        open_slot.ai_model = payload.model_name
        open_slot.ready = True
    else:
        raise HTTPException(status_code=400, detail="Unknown AI type")
    return session.to_lobby()


@app.post("/api/games/{game_id}/start")
async def start_game(game_id: str, user: User = Depends(require_user)):
    try:
        session = manager.get(game_id)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="Game not found") from exc
    if session.status != "lobby":
        raise HTTPException(status_code=400, detail="Game already started")
    if user.username != session.created_by:
        raise HTTPException(status_code=403, detail="Only host can start")
    try:
        _start_game(session)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    await _broadcast(session)
    await _run_ai_turns(session)
    return _build_state(session, user.username)


@app.post("/api/games/{game_id}/action")
async def take_action(game_id: str, payload: ActionRequest, user: User = Depends(require_user)):
    try:
        session = manager.get(game_id)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="Game not found") from exc
    if session.status != "running" or not session.env:
        raise HTTPException(status_code=400, detail="Game not running")
    try:
        action_type = ActionType(payload.type)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail="Unknown action type") from exc
    action = Action(type=action_type, payload=payload.payload)
    actor = user.username
    current = session.env.game.current_player.name
    if action_type == ActionType.DISCARD:
        target = payload.payload.get("player")
        if target != actor:
            raise HTTPException(status_code=403, detail="You can only discard for yourself")
    elif actor != current:
        raise HTTPException(status_code=403, detail="Not your turn")

    async with session.lock:
        _add_log(session, actor, action)
        _, _, done, info = session.env.step(action)
        if info.get("invalid"):
            raise HTTPException(status_code=400, detail=info.get("error", "Invalid action"))
        if done:
            _finish_game(session)
    await _broadcast(session)
    await _run_ai_turns(session)
    return _build_state(session, user.username)


@app.get("/api/stats")
def stats(user: User = Depends(require_user)):
    with Session(engine) as session:
        records = session.exec(select(GameRecord)).all()
    total_games = len(records)
    user_games = [r for r in records if user.username in r.players]
    user_wins = sum(1 for r in records if r.winner == user.username)
    avg_turns = sum(r.total_turns for r in records) / total_games if total_games else 0
    return {
        "total_games": total_games,
        "user_games": len(user_games),
        "user_wins": user_wins,
        "avg_turns": avg_turns,
    }


@app.websocket("/ws/games/{game_id}")
async def game_ws(websocket: WebSocket, game_id: str):
    await websocket.accept()
    try:
        user = user_from_ws(websocket)
    except HTTPException:
        await websocket.close(code=1008)
        return
    try:
        session = manager.get(game_id)
    except KeyError:
        await websocket.close(code=1008)
        return
    session.connections[websocket] = user.username
    await websocket.send_text(json.dumps({"type": "state", "data": _build_state(session, user.username)}))
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        session.connections.pop(websocket, None)
