from __future__ import annotations

from dataclasses import dataclass, field, asdict
from typing import Any, Dict, Tuple
import json
import os

from .config import STATE_PATH


class GameStatus(str):
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


@dataclass
class Constraints:
    mins: Dict[str, int] = field(default_factory=dict)


@dataclass
class AttributeStats:
    relative_frequencies: Dict[str, float] = field(default_factory=dict)
    correlations: Dict[str, Dict[str, float]] = field(default_factory=dict)


@dataclass
class GameRecord:
    game_id: str
    initial_raw: Dict[str, Any] | None = None
    scenario: int | None = None
    status: str = GameStatus.RUNNING
    constraints: Constraints = field(default_factory=Constraints)
    stats: AttributeStats = field(default_factory=AttributeStats)
    admitted_count: int = 0
    rejected_count: int = 0
    next_person: Dict[str, Any] | None = None
    tallies: Dict[str, int] = field(default_factory=dict)
    policy_path: str | None = None  # <— path to the user-provided policy script

    def bump_tallies(self, attributes: Dict[str, bool]) -> None:
        for k, v in attributes.items():
            if v:
                self.tallies[k] = self.tallies.get(k, 0) + 1

    def raw(self) -> Dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "GameRecord":
        constraints = Constraints(**(d.get("constraints") or {}))
        stats = AttributeStats(**(d.get("stats") or {}))
        return cls(
            game_id=d["game_id"],
            initial_raw=d.get("initial_raw"),
            scenario=d.get("scenario"),
            status=d.get("status", GameStatus.RUNNING),
            constraints=constraints,
            stats=stats,
            admitted_count=d.get("admitted_count", 0),
            rejected_count=d.get("rejected_count", 0),
            next_person=d.get("next_person"),
            tallies=d.get("tallies", {}) or {},
            policy_path=d.get("policy_path"),
        )


@dataclass
class GameState:
    player_id: str | None = None
    current_game_id: str | None = None
    games: Dict[str, GameRecord] = field(default_factory=dict)

    def save(self) -> None:
        payload = {
            "player_id": self.player_id,
            "current_game_id": self.current_game_id,
            "games": {gid: g.raw() for gid, g in self.games.items()},
        }
        os.makedirs(os.path.dirname(STATE_PATH), exist_ok=True)
        with open(STATE_PATH, "w", encoding="utf-8") as f:
            json.dump(payload, f, indent=2)

    @classmethod
    def load(cls) -> "GameState":
        try:
            with open(STATE_PATH, "r", encoding="utf-8") as f:
                raw = json.load(f)
        except FileNotFoundError:
            return cls()
        gs = cls()
        gs.player_id = raw.get("player_id")
        gs.current_game_id = raw.get("current_game_id")
        for gid, graw in (raw.get("games") or {}).items():
            gs.games[gid] = GameRecord.from_dict(graw)
        return gs

    def add_game(self, rec: GameRecord) -> None:
        self.games[rec.game_id] = rec
        self.current_game_id = rec.game_id

    def current_game(self) -> GameRecord | None:
        return self.games.get(self.current_game_id) if self.current_game_id else None

    def set_current_game(self, game_id_prefix: str) -> bool:
        g = self.get_game(game_id_prefix)
        if g:
            self.current_game_id = g.game_id
            return True
        return False

    def get_game(self, game_id_prefix: str) -> GameRecord | None:
        matches = [gid for gid in self.games if gid.startswith(game_id_prefix)]
        if len(matches) == 1:
            return self.games[matches[0]]
        return None

    def list_games(self) -> Dict[str, Tuple[str, int | None, int, int]]:
        return {
            gid: (g.status, g.scenario, g.admitted_count, g.rejected_count)
            for gid, g in self.games.items()
        }