from __future__ import annotations

from pathlib import Path
from typing import Iterable

from ..schemas import ChatMessage, SessionState
from ..util import ensure_dir, read_json, write_json


class SessionStore:
    def __init__(
        self,
        data_dir: str | Path,
        *,
        default_provider: str,
        default_model: str,
        history_limit: int = 40,
    ):
        self.path = ensure_dir(data_dir) / "sessions.json"
        self.default_provider = default_provider
        self.default_model = default_model
        self.history_limit = None if history_limit in (None, 0) else history_limit
        self._cache: dict[str, dict] = read_json(self.path, default={})

    def list_sessions(self) -> list[SessionState]:
        return [SessionState.from_dict(payload) for payload in self._cache.values()]

    def get_or_create(self, session_id: str, *, chat_id: str, user_id: int | None) -> SessionState:
        payload = self._cache.get(session_id)
        if payload is None:
            state = SessionState(
                session_id=session_id,
                chat_id=chat_id,
                user_id=user_id,
                provider=self.default_provider,
                model=self.default_model,
            )
            self.save_state(state)
            return state
        state = SessionState.from_dict(payload)
        state.chat_id = chat_id
        if user_id is not None:
            state.user_id = user_id
        return state

    def save_state(self, state: SessionState) -> None:
        self._cache[state.session_id] = state.to_dict()
        write_json(self.path, self._cache)

    def append_messages(self, state: SessionState, messages: Iterable[ChatMessage]) -> SessionState:
        state.history.extend(messages)
        if self.history_limit and len(state.history) > self.history_limit:
            state.history = state.history[-self.history_limit :]
        self.save_state(state)
        return state

    def close(self, session_id: str) -> None:
        if session_id in self._cache:
            payload = self._cache[session_id]
            payload["closed"] = True
            self._cache[session_id] = payload
            write_json(self.path, self._cache)

    def reopen(self, session_id: str) -> None:
        if session_id in self._cache:
            payload = self._cache[session_id]
            payload["closed"] = False
            self._cache[session_id] = payload
            write_json(self.path, self._cache)

    def reset(self, session_id: str) -> None:
        if session_id in self._cache:
            del self._cache[session_id]
            write_json(self.path, self._cache)
