from __future__ import annotations

from contextlib import AbstractContextManager
from dataclasses import dataclass
from types import TracebackType
from typing import Any, Literal

from unmake.models import OperationRecord, SessionRecord, SessionStatus
from unmake.observability import MetricsCollector
from unmake.storage import StorageBackend


@dataclass(slots=True)
class SessionHandle:
    storage: StorageBackend
    record: SessionRecord
    metrics: MetricsCollector | None = None

    @property
    def session_id(self) -> str:
        return self.record.session_id

    def record_operation(self, op_type: str, payload: dict[str, Any]) -> OperationRecord:
        op = self.storage.append_operation(
            session_id=self.session_id, op_type=op_type, payload=payload
        )
        if self.metrics is not None:
            self.metrics.record_operation()
        return op


class SessionContext(AbstractContextManager[SessionHandle]):
    def __init__(
        self,
        storage: StorageBackend,
        context: dict[str, Any],
        metrics: MetricsCollector | None = None,
    ) -> None:
        self.storage = storage
        self.context = context
        self.metrics = metrics
        self.handle: SessionHandle | None = None

    def __enter__(self) -> SessionHandle:
        record = self.storage.start_session(context=self.context)
        if self.metrics is not None:
            self.metrics.record_session_created()
        self.handle = SessionHandle(storage=self.storage, record=record, metrics=self.metrics)
        return self.handle

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> Literal[False]:
        if self.handle is None:
            return False
        if exc is None:
            self.storage.close_session(
                session_id=self.handle.session_id, status=SessionStatus.COMPLETED
            )
            return False
        self.storage.close_session(
            session_id=self.handle.session_id,
            status=SessionStatus.FAILED,
            error=str(exc),
        )
        return False
