"""
LearnerModel — In-memory per-session mastery tracking.
Designed so Unity dev can later persist/restore via get_state/set_state.
"""

import time
from typing import Optional

# ---------------------------------------------------------------------------
# Default mastery structure (all rules start at 0.0)
# ---------------------------------------------------------------------------
DEFAULT_MASTERY = {
    "topic_marker": 0.0,
    "copula": 0.0,
    "negative_copula": 0.0,
    "indirect_quote_dago": 0.0,
    "indirect_quote_commands": 0.0,
    "indirect_quote_questions": 0.0,
    "indirect_quote_suggestions": 0.0,
    "regret_expression": 0.0,
}

# Mastery thresholds for difficulty escalation
MASTERY_THRESHOLD_MID = 0.5    # Unlock difficulty 2
MASTERY_THRESHOLD_HIGH = 0.75  # Unlock difficulty 3

# Score weights
CORRECT_WEIGHT = 0.15     # +15% mastery per correct
INCORRECT_WEIGHT = 0.05   # -5% mastery per incorrect (floor 0)
MAX_HISTORY = 20          # Keep last 20 entries per session


class LearnerModel:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.mastery = dict(DEFAULT_MASTERY)
        self.history: list = []
        self.difficulty: int = 1       # 1 = beginner, 2 = intermediate, 3 = advanced
        self.question_count: int = 0
        self.correct_count: int = 0
        self.streak: int = 0
        self.created_at: float = time.time()
        self.last_active: float = time.time()

    def record_outcome(self, grammar_rule: str, correct: bool, interaction_mode: str = ""):
        """Update mastery score and history for a question outcome."""
        self.last_active = time.time()
        self.question_count += 1

        if correct:
            self.correct_count += 1
            self.streak += 1
            delta = CORRECT_WEIGHT
        else:
            self.streak = 0
            delta = -INCORRECT_WEIGHT

        if grammar_rule in self.mastery:
            new_score = self.mastery[grammar_rule] + delta
            self.mastery[grammar_rule] = max(0.0, min(1.0, new_score))

        self.history.append({
            "grammar_rule": grammar_rule,
            "correct": correct,
            "interaction_mode": interaction_mode,
            "timestamp": time.time(),
        })

        # Trim history
        if len(self.history) > MAX_HISTORY:
            self.history = self.history[-MAX_HISTORY:]

        # Auto-escalate difficulty
        self._update_difficulty()

    def _update_difficulty(self):
        avg_mastery = sum(self.mastery.values()) / len(self.mastery)
        if avg_mastery >= MASTERY_THRESHOLD_HIGH:
            self.difficulty = 3
        elif avg_mastery >= MASTERY_THRESHOLD_MID:
            self.difficulty = 2
        else:
            self.difficulty = 1

    def get_weakest_rule(self) -> str:
        """Return the grammar rule with the lowest mastery score."""
        return min(self.mastery, key=lambda k: self.mastery[k])

    def get_strongest_rule(self) -> str:
        """Return the grammar rule with the highest mastery score."""
        return max(self.mastery, key=lambda k: self.mastery[k])

    def get_recommended_rule(self) -> str:
        """
        Smart rule selection:
        - Early sessions: cycle through all rules
        - Later: weight toward weakest rules with some randomness
        """
        import random

        # First pass: if any rule has never been tested, prioritize it
        untested = [r for r, score in self.mastery.items() if score == 0.0]
        if untested and self.question_count < len(DEFAULT_MASTERY) * 2:
            return random.choice(untested)

        # Otherwise weight selection toward weaker rules
        rules = list(self.mastery.keys())
        weights = [max(0.05, 1.0 - score) for score in self.mastery.values()]
        return random.choices(rules, weights=weights, k=1)[0]

    def get_state(self) -> dict:
        """
        Returns full serializable state.
        Unity dev can store this and send it back via set_state on reconnect.
        """
        return {
            "session_id": self.session_id,
            "mastery": dict(self.mastery),
            "history": list(self.history),
            "difficulty": self.difficulty,
            "question_count": self.question_count,
            "correct_count": self.correct_count,
            "streak": self.streak,
            "created_at": self.created_at,
            "last_active": self.last_active,
        }

    def set_state(self, state: dict):
        """Restore state from a previously saved snapshot (for Unity persistence)."""
        self.mastery = state.get("mastery", dict(DEFAULT_MASTERY))
        self.history = state.get("history", [])
        self.difficulty = state.get("difficulty", 1)
        self.question_count = state.get("question_count", 0)
        self.correct_count = state.get("correct_count", 0)
        self.streak = state.get("streak", 0)

    def reset(self):
        """Full reset for this session."""
        self.mastery = dict(DEFAULT_MASTERY)
        self.history = []
        self.difficulty = 1
        self.question_count = 0
        self.correct_count = 0
        self.streak = 0


# ---------------------------------------------------------------------------
# Session Store — maps session_id → LearnerModel
# ---------------------------------------------------------------------------
_sessions: dict[str, LearnerModel] = {}
SESSION_TIMEOUT = 3600  # 1 hour


def get_or_create_session(session_id: str) -> LearnerModel:
    if session_id not in _sessions:
        _sessions[session_id] = LearnerModel(session_id)
    else:
        _sessions[session_id].last_active = time.time()
    return _sessions[session_id]


def get_session(session_id: str) -> Optional[LearnerModel]:
    return _sessions.get(session_id)


def delete_session(session_id: str):
    _sessions.pop(session_id, None)


def purge_stale_sessions():
    """Remove sessions inactive for more than SESSION_TIMEOUT seconds."""
    now = time.time()
    stale = [sid for sid, model in _sessions.items()
             if now - model.last_active > SESSION_TIMEOUT]