from __future__ import annotations

from typing import Any

from unmake.models import OperationRecord, SessionRecord, SessionStatus, new_id, utcnow
from unmake.utils.validation import sanitize_op_type, validate_payload


class InMemoryStorage:
    """In-memory storage backend for session history.

    Warning:
        This storage is NOT thread-safe. Use external locking for concurrent
        access or use SQLiteStorage for multi-threaded workflows.
    """

    def __init__(self) -> None:
        self._sessions: dict[str, SessionRecord] = {}
        self._operations: list[OperationRecord] = []
        self._next_sequence: dict[str, int] = {}

    def start_session(self, *, context: dict[str, Any]) -> SessionRecord:
        record = SessionRecord(
            session_id=new_id(), created_at=utcnow(), status=SessionStatus.RUNNING, context=context
        )
        self._sessions[record.session_id] = record
        self._next_sequence[record.session_id] = 1
        return record

    def append_operation(
        self,
        *,
        session_id: str,
        op_type: str,
        payload: dict[str, Any],
    ) -> OperationRecord:
        op_type = sanitize_op_type(op_type)
        validate_payload(payload)

        if session_id not in self._sessions:
            raise KeyError(
                f"Session not found: {session_id}. Use list_sessions() to see available sessions."
            )
        seq = self._next_sequence[session_id]
        self._next_sequence[session_id] += 1
        op = OperationRecord(
            op_id=new_id(),
            session_id=session_id,
            sequence=seq,
            op_type=op_type,
            payload=payload,
            created_at=utcnow(),
        )
        self._operations.append(op)
        return op

    def close_session(
        self,
        *,
        session_id: str,
        status: SessionStatus,
        error: str | None = None,
    ) -> SessionRecord:
        if session_id not in self._sessions:
            raise KeyError(
                f"Session not found: {session_id}. Use list_sessions() to see available sessions."
            )
        record = self._sessions[session_id]
        record.status = status
        record.error = error
        record.completed_at = utcnow()
        return record

    def list_sessions(self, *, limit: int | None = None) -> list[SessionRecord]:
        sessions = sorted(
            self._sessions.values(), key=lambda session: session.created_at, reverse=True
        )
        return sessions[:limit] if limit else sessions

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

    def get_operations(self, session_id: str) -> list[OperationRecord]:
        return sorted(
            (op for op in self._operations if op.session_id == session_id),
            key=lambda operation: operation.sequence,
        )
