from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from typing import Any, Literal
from uuid import uuid4

from niannian.chat.parts import parts_from_legacy
from niannian.virtual_persona.chat_mode_gate import assert_chat_mode_allowed_for_vp

Role = Literal["user", "assistant", "system"]
ChatMode = Literal["auto", "default", "companion", "work", "education"]
_VALID_CHAT_MODES = frozenset({"auto", "default", "companion", "work", "education"})


@dataclass
class Turn:
    id: str
    role: Role
    content: str
    created_at: str
    meta: dict[str, Any] = field(default_factory=dict)
    parts: list[dict[str, Any]] | None = None

    def to_dict(self) -> dict[str, Any]:
        out: dict[str, Any] = {
            "id": self.id,
            "role": self.role,
            "content": self.content,
            "created_at": self.created_at,
            "meta": self.meta,
        }
        if self.parts is not None:
            out["parts"] = self.parts
        else:
            legacy = parts_from_legacy(content=self.content, meta=self.meta, parts=None)
            if legacy:
                out["parts"] = legacy
        return out


@dataclass
class Session:
    id: str
    project: str
    turns: list[Turn] = field(default_factory=list)
    meta: dict[str, Any] = field(default_factory=dict)
    tool_scope: dict[str, Any] | None = None
    created_at: str = ""
    updated_at: str = ""

    def to_dict(self) -> dict[str, Any]:
        out: dict[str, Any] = {
            "id": self.id,
            "project": self.project,
            "turns": [t.to_dict() for t in self.turns],
            "meta": self.meta,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
        }
        if self.tool_scope is not None:
            out["tool_scope"] = self.tool_scope
        return out


def text_from_parts(parts: list[dict[str, Any]] | None) -> str:
    """Combined user text from text parts (for @ mention parsing)."""
    if not parts:
        return ""
    from niannian.chat.parts import parts_to_content

    return parts_to_content(parts)


def collect_recent_user_texts(turns: list[Turn], *, limit: int = 3) -> list[str]:
    """Most recent user texts (newest first), skipping recalled and empty content."""
    out: list[str] = []
    for turn in reversed(turns):
        if turn.role != "user":
            continue
        if turn.meta.get("recalled_at"):
            continue
        if not (turn.content or "").strip():
            continue
        out.append(turn.content)
        if len(out) >= limit:
            break
    return out


class TurnInProgressError(Exception):
    """同 session 已有 Turn 正在 GENERATE。"""


class SessionNotFoundError(Exception):
    pass


class TurnNotFoundError(Exception):
    def __init__(self, turn_id: str) -> None:
        self.turn_id = turn_id
        super().__init__(turn_id)


class TurnNotRecallableError(Exception):
    def __init__(self, turn_id: str) -> None:
        self.turn_id = turn_id
        super().__init__(turn_id)


class RecallWindowExpiredError(Exception):
    def __init__(self, turn_id: str) -> None:
        self.turn_id = turn_id
        super().__init__(turn_id)


class PairNotFoundError(Exception):
    def __init__(self, assistant_turn_id: str) -> None:
        self.assistant_turn_id = assistant_turn_id
        super().__init__(assistant_turn_id)


class AlreadyRememberedError(Exception):
    def __init__(self, assistant_turn_id: str) -> None:
        self.assistant_turn_id = assistant_turn_id
        super().__init__(assistant_turn_id)


class TurnRecalledError(Exception):
    def __init__(self, assistant_turn_id: str) -> None:
        self.assistant_turn_id = assistant_turn_id
        super().__init__(assistant_turn_id)


class RecallWindowActiveError(Exception):
    def __init__(self, assistant_turn_id: str) -> None:
        self.assistant_turn_id = assistant_turn_id
        super().__init__(assistant_turn_id)


class NotRememberableError(Exception):
    def __init__(self, assistant_turn_id: str, *, reason: str = "") -> None:
        self.assistant_turn_id = assistant_turn_id
        self.reason = reason
        super().__init__(assistant_turn_id)


class TurnAlreadySupersededError(Exception):
    def __init__(self, assistant_turn_id: str) -> None:
        self.assistant_turn_id = assistant_turn_id
        super().__init__(assistant_turn_id)


class NotAssistantTurnError(Exception):
    def __init__(self, turn_id: str) -> None:
        self.turn_id = turn_id
        super().__init__(turn_id)


class ConsolidationNotBatchModeError(Exception):
    pass


class SessionNotIdleError(Exception):
    def __init__(self, session_id: str, *, idle_seconds: float, required_seconds: float) -> None:
        self.session_id = session_id
        self.idle_seconds = idle_seconds
        self.required_seconds = required_seconds
        super().__init__(session_id)


class EnrichmentNotFoundError(Exception):
    def __init__(self, key: str) -> None:
        self.key = key
        super().__init__(key)


class EnrichmentNotActionableError(Exception):
    def __init__(self, key: str, reason: str) -> None:
        self.key = key
        self.reason = reason
        super().__init__(f"{key}:{reason}")


class EnrichmentAlreadyPersistedError(Exception):
    def __init__(self, key: str, *, vault_id: str | None = None, atom_id: str | None = None) -> None:
        self.key = key
        self.vault_id = vault_id
        self.atom_id = atom_id
        super().__init__(key)


class EnrichmentKindNotSupportedError(Exception):
    def __init__(self, key: str, kind: str, action: str) -> None:
        self.key = key
        self.kind = kind
        self.action = action
        super().__init__(f"{kind}:{action}")


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def resolve_turn_pair(session: Session, assistant_turn_id: str) -> tuple[Turn, Turn]:
    """Resolve assistant turn and paired user turn."""
    assistant_turn: Turn | None = None
    assistant_idx = -1
    for i, turn in enumerate(session.turns):
        if turn.id == assistant_turn_id:
            assistant_turn = turn
            assistant_idx = i
            break
    if assistant_turn is None:
        raise TurnNotFoundError(assistant_turn_id)
    if assistant_turn.role != "assistant":
        raise TurnNotFoundError(assistant_turn_id)

    user_turn: Turn | None = None
    user_id = assistant_turn.meta.get("user_turn_id")
    if user_id:
        for turn in session.turns:
            if turn.id == user_id and turn.role == "user":
                user_turn = turn
                break
    if user_turn is None:
        for i in range(assistant_idx - 1, -1, -1):
            turn = session.turns[i]
            if turn.role == "user":
                user_turn = turn
                break
    if user_turn is None:
        raise PairNotFoundError(assistant_turn_id)
    return user_turn, assistant_turn


class SessionStore:
    """进程内 Session 存储；闲置 24h 可 evict。"""

    def __init__(self, *, idle_hours: float = 24.0) -> None:
        self._sessions: dict[str, Session] = {}
        self._turn_locks: set[str] = set()
        self._idle_hours = idle_hours

    def create(self, project: str = "global") -> Session:
        self.evict_idle()
        now = _now_iso()
        session = Session(id=uuid4().hex, project=project, created_at=now, updated_at=now)
        self._sessions[session.id] = session
        return session

    def get(self, session_id: str) -> Session | None:
        return self._sessions.get(session_id)

    def require(self, session_id: str) -> Session:
        session = self.get(session_id)
        if session is None:
            raise SessionNotFoundError(session_id)
        return session

    def delete(self, session_id: str) -> bool:
        self._turn_locks.discard(session_id)
        return self._sessions.pop(session_id, None) is not None

    def append_turn(self, session_id: str, turn: Turn) -> Session:
        session = self.require(session_id)
        session.turns.append(turn)
        session.updated_at = _now_iso()
        return session

    def list_turns(self, session_id: str) -> list[Turn]:
        return list(self.require(session_id).turns)

    def patch_meta(self, session_id: str, updates: dict[str, Any]) -> Session:
        session = self.require(session_id)
        meta = dict(session.meta)
        if updates.get("clear_pathway_sticky"):
            meta.pop("pathway_sticky", None)
            meta.pop("pathway_sticky_hits", None)
            meta.pop("pathway_sticky_signature", None)
        if "chat_mode" in updates:
            mode = updates["chat_mode"]
            if mode is None or mode == "":
                meta.pop("chat_mode", None)
            else:
                mode_str = str(mode).strip()
                if mode_str not in _VALID_CHAT_MODES:
                    raise ValueError(f"invalid chat_mode: {mode_str}")
                assert_chat_mode_allowed_for_vp(meta, mode_str)
                meta["chat_mode"] = mode_str
        if "regime_lock" in updates:
            lock = updates["regime_lock"]
            if lock is None or lock == "":
                meta.pop("regime_lock", None)
            else:
                meta["regime_lock"] = str(lock)
        if "current_object_id" in updates:
            oid = updates["current_object_id"]
            if oid is None or (isinstance(oid, str) and not oid.strip()):
                meta.pop("current_object_id", None)
            else:
                meta["current_object_id"] = str(oid).strip()
        if "persona_override" in updates:
            override = updates["persona_override"]
            if override is None:
                meta.pop("persona_override", None)
            elif isinstance(override, dict):
                meta["persona_override"] = override
        if "plan_mode" in updates:
            pm = updates["plan_mode"]
            if pm is None:
                meta.pop("plan_mode", None)
            else:
                meta["plan_mode"] = bool(pm)
        if "turn_policy" in updates:
            tp = updates["turn_policy"]
            if tp is None:
                meta.pop("turn_policy", None)
            elif isinstance(tp, dict):
                meta["turn_policy"] = dict(tp)
        if "tool_scope" in updates:
            ts = updates["tool_scope"]
            if ts is None:
                session.tool_scope = None
            elif isinstance(ts, dict):
                session.tool_scope = dict(ts)
        session.meta = meta
        session.updated_at = _now_iso()
        return session

    def acquire_turn_lock(self, session_id: str) -> None:
        if session_id in self._turn_locks:
            raise TurnInProgressError(session_id)
        self._turn_locks.add(session_id)

    def release_turn_lock(self, session_id: str) -> None:
        self._turn_locks.discard(session_id)

    def recall_turn(
        self,
        session_id: str,
        turn_id: str,
        *,
        by_user_id: str,
        recall_window_seconds: float = 120.0,
    ) -> tuple[list[str], str | None]:
        """软删 user turn 及配对 assistant turn；返回 (recalled_ids, ingested_atom_id)。"""
        if session_id in self._turn_locks:
            raise TurnInProgressError(session_id)

        session = self.require(session_id)

        user_turn: Turn | None = None
        user_idx = -1
        for i, turn in enumerate(session.turns):
            if turn.id == turn_id:
                user_turn = turn
                user_idx = i
                break
        if user_turn is None:
            raise TurnNotFoundError(turn_id)
        if user_turn.role != "user" or user_turn.meta.get("recalled_at"):
            raise TurnNotRecallableError(turn_id)

        try:
            created = datetime.fromisoformat(user_turn.created_at)
        except ValueError as exc:
            raise TurnNotRecallableError(turn_id) from exc
        if created.tzinfo is None:
            created = created.replace(tzinfo=timezone.utc)
        elapsed = (datetime.now(timezone.utc) - created).total_seconds()
        if elapsed > recall_window_seconds:
            raise RecallWindowExpiredError(turn_id)

        assistant_turn: Turn | None = None
        assistant_id = user_turn.meta.get("assistant_turn_id")
        if assistant_id:
            for turn in session.turns:
                if turn.id == assistant_id and turn.role == "assistant":
                    assistant_turn = turn
                    break
        if assistant_turn is None and user_idx + 1 < len(session.turns):
            next_turn = session.turns[user_idx + 1]
            if next_turn.role == "assistant" and not next_turn.meta.get("recalled_at"):
                assistant_turn = next_turn

        ingested_atom_id = user_turn.meta.get("ingested_atom_id")
        if ingested_atom_id is not None:
            ingested_atom_id = str(ingested_atom_id)

        recalled_at = _now_iso()
        recalled: list[str] = []

        def _soft_recall(turn: Turn) -> None:
            meta = dict(turn.meta)
            meta["recalled_at"] = recalled_at
            meta["recalled_by"] = by_user_id
            turn.meta = meta
            turn.content = ""

        _soft_recall(user_turn)
        recalled.append(user_turn.id)
        if assistant_turn is not None and not assistant_turn.meta.get("recalled_at"):
            _soft_recall(assistant_turn)
            recalled.append(assistant_turn.id)

        session.updated_at = _now_iso()
        return recalled, ingested_atom_id

    def evict_idle(self) -> int:
        cutoff = datetime.now(timezone.utc) - timedelta(hours=self._idle_hours)
        evicted = 0
        for sid, session in list(self._sessions.items()):
            try:
                updated = datetime.fromisoformat(session.updated_at)
            except ValueError:
                continue
            if updated.tzinfo is None:
                updated = updated.replace(tzinfo=timezone.utc)
            if updated < cutoff:
                self.delete(sid)
                evicted += 1
        return evicted
