""" 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))