"""Session memory backends and compaction settings."""

from __future__ import annotations

import json
from hashlib import sha256
from pathlib import Path
from typing import Protocol

from pydantic import BaseModel

from yier_agents.src.message import Message


class CompactionConfig(BaseModel):
    """Controls when and how conversation history is compacted."""

    enabled: bool = False
    trigger_message_count: int = 24
    preserve_recent_messages: int = 8
    summary_max_tokens: int = 512


class SessionStore(Protocol):
    """Storage contract for session memory."""

    def load(self, session_id: str) -> list[Message] | None:
        """Load all messages for a session."""

    def save(self, session_id: str, messages: list[Message]) -> None:
        """Persist all messages for a session."""

    def clear_session(self, session_id: str) -> bool:
        """Delete one session."""

    def clear_all(self) -> int:
        """Delete every stored session and return the deleted count."""

    def get_session_messages(self, session_id: str) -> list[Message] | None:
        """Return the current messages for one session."""

    def get_session_count(self) -> int:
        """Return the number of stored sessions."""


class InMemorySessionStore:
    """Simple in-process session store."""

    def __init__(self) -> None:
        self._sessions: dict[str, list[Message]] = {}

    def load(self, session_id: str) -> list[Message] | None:
        messages = self._sessions.get(session_id)
        return messages.copy() if messages else None

    def save(self, session_id: str, messages: list[Message]) -> None:
        self._sessions[session_id] = messages.copy()

    def clear_session(self, session_id: str) -> bool:
        if session_id not in self._sessions:
            return False

        del self._sessions[session_id]
        return True

    def clear_all(self) -> int:
        count = len(self._sessions)
        self._sessions.clear()
        return count

    def get_session_messages(self, session_id: str) -> list[Message] | None:
        return self.load(session_id)

    def get_session_count(self) -> int:
        return len(self._sessions)


class JSONSessionStore:
    """Filesystem-backed session store using one JSON file per session."""

    def __init__(self, base_path: Path) -> None:
        self.base_path = base_path
        self.base_path.mkdir(parents=True, exist_ok=True)

    def load(self, session_id: str) -> list[Message] | None:
        session_file = self._session_file(session_id)
        if not session_file.exists():
            return None

        data = json.loads(session_file.read_text(encoding="utf-8"))
        return [Message.model_validate(item) for item in data["messages"]]

    def save(self, session_id: str, messages: list[Message]) -> None:
        session_file = self._session_file(session_id)
        payload = {
            "session_id": session_id,
            "messages": [message.model_dump(exclude_none=True) for message in messages],
        }
        session_file.write_text(
            json.dumps(payload, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )

    def clear_session(self, session_id: str) -> bool:
        session_file = self._session_file(session_id)
        if not session_file.exists():
            return False

        session_file.unlink()
        return True

    def clear_all(self) -> int:
        count = 0
        for session_file in self.base_path.glob("*.json"):
            session_file.unlink()
            count += 1
        return count

    def get_session_messages(self, session_id: str) -> list[Message] | None:
        return self.load(session_id)

    def get_session_count(self) -> int:
        return sum(1 for _ in self.base_path.glob("*.json"))

    def _session_file(self, session_id: str) -> Path:
        digest = sha256(session_id.encode("utf-8")).hexdigest()
        return self.base_path / f"{digest}.json"
