models/session_models.py

"""
Pydantic models for runtime state: streak store, analysis results, shuffler.
"""
from __future__ import annotations
import datetime
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field, field_validator, model_validator


# ── Module keys (matches original JS constants) ───────────────────────────────
class ModuleKey(str, Enum):
    JAM       = "jam"
    TT        = "tt"
    IMPROMPTU = "impromptu"
    INTERVIEW = "interview"
    OTHER     = "other"


FREE_MODULE_MAP: dict[str, ModuleKey] = {
    "JAM":       ModuleKey.JAM,
    "Tongue Twister": ModuleKey.TT,
    "Impromptu": ModuleKey.IMPROMPTU,
    "Interview": ModuleKey.INTERVIEW,
}


class FreeModuleChoice(str, Enum):
    JAM       = "JAM"
    TT        = "Tongue Twister"
    IMPROMPTU = "Impromptu"
    INTERVIEW = "Interview"


# ── Per-module lifetime stats ─────────────────────────────────────────────────
class ModuleStats(BaseModel):
    count: int = 0
    totalDuration: int = 0   # seconds — matches JS field name for BrowserState compat

    @property
    def avg_duration(self) -> int:
        return round(self.totalDuration / self.count) if self.count else 0

    def avg_str(self) -> str:
        avg = self.avg_duration
        if not avg:
            return "–"
        return f"{avg // 60}:{avg % 60:02d}"


# ── Per-day entry ─────────────────────────────────────────────────────────────
class DayEntry(BaseModel):
    total: int = 0
    modules: dict[str, int] = Field(default_factory=dict)   # module_key → session count that day

    @field_validator("total")
    @classmethod
    def total_non_negative(cls, v: int) -> int:
        if v < 0:
            raise ValueError("total must be >= 0")
        return v


# ── Full streak store (mirrors localStorage schema v2) ────────────────────────
class StreakStore(BaseModel):
    v: int = 2
    days: dict[str, DayEntry] = Field(default_factory=dict)        # ISO date str → DayEntry
    modules: dict[str, ModuleStats] = Field(default_factory=dict)  # ModuleKey value → ModuleStats
    longest: int = 0

    @model_validator(mode="before")
    @classmethod
    def migrate_v1(cls, data: Any) -> Any:
        """Migrate v1 stores (days[k] = int) to v2 (days[k] = DayEntry)."""
        if not isinstance(data, dict):
            return data
        raw_days = data.get("days", {})
        if isinstance(raw_days, dict):
            migrated = {}
            for k, v in raw_days.items():
                if isinstance(v, int):
                    migrated[k] = {"total": v, "modules": {}}
                else:
                    migrated[k] = v
            data["days"] = migrated
        return data

    def day_total(self, key: str) -> int:
        entry = self.days.get(key)
        if entry is None:
            return 0
        return entry.total

    def calc_streak(self) -> int:
        streak = 0
        d = datetime.date.today()
        if not self.day_total(d.isoformat()):
            d -= datetime.timedelta(days=1)
        while self.day_total(d.isoformat()):
            streak += 1
            d -= datetime.timedelta(days=1)
        return streak

    def calc_longest(self) -> int:
        keys = sorted(k for k in self.days if self.day_total(k))
        best = run = 0
        prev: datetime.date | None = None
        for k in keys:
            cur = datetime.date.fromisoformat(k)
            if prev:
                run = run + 1 if (cur - prev).days == 1 else 1
            else:
                run = 1
            best = max(best, run)
            prev = cur
        return best

    def record_session(self, module_key: ModuleKey, duration_sec: int = 0) -> "StreakStore":
        """Immutable-style: returns updated copy."""
        data = self.model_copy(deep=True)
        today = datetime.date.today().isoformat()
        entry = data.days.get(today) or DayEntry()
        entry.total += 1
        mk = module_key.value
        entry.modules[mk] = entry.modules.get(mk, 0) + 1
        data.days[today] = entry

        stats = data.modules.get(mk) or ModuleStats()
        stats.count += 1
        stats.totalDuration += max(0, duration_sec)
        data.modules[mk] = stats

        data.longest = max(data.longest, data.calc_longest())

        # Prune entries older than 400 days
        cutoff = (datetime.date.today() - datetime.timedelta(days=400)).isoformat()
        data.days = {k: v for k, v in data.days.items() if k >= cutoff}
        return data

    @classmethod
    def from_raw(cls, raw: dict | None) -> "StreakStore":
        """Safely deserialise from BrowserState dict, falling back to empty."""
        if not isinstance(raw, dict):
            return cls()
        try:
            return cls.model_validate(raw)
        except Exception:
            return cls()

    def to_raw(self) -> dict:
        """Serialise to plain dict for BrowserState storage."""
        return self.model_dump()


# ── No-repeat shuffler state (persisted per-key in BrowserState) ──────────────
class ShufflerState(BaseModel):
    seen: list[str] = Field(default_factory=list)   # serialised items already shown this cycle


# ── Analysis result ───────────────────────────────────────────────────────────
class AnalysisResult(BaseModel):
    transcript: str | None = None
    pronunciation: int | None = None   # 0-100
    fluency: int | None = None         # 0-100
    error: str | None = None

    @field_validator("pronunciation", "fluency", mode="before")
    @classmethod
    def clamp_score(cls, v: int | None) -> int | None:
        if v is None:
            return None
        return max(0, min(100, v))