from __future__ import annotations

import logging
from dataclasses import dataclass
from functools import lru_cache
from time import perf_counter
from typing import Any

from unmake.core.operations import OperationRegistry
from unmake.core.session import SessionContext
from unmake.exceptions import SessionNotFoundError
from unmake.models import OperationRecord, SessionRecord, SessionStatus
from unmake.observability import MetricsCollector
from unmake.storage import SQLiteStorage, StorageBackend

logger = logging.getLogger("unmake")


@dataclass(slots=True)
class HistoryManager:
    storage: StorageBackend
    metrics: MetricsCollector | None = None

    def start_session(self, **context: Any) -> SessionContext:
        logger.debug("Starting session with context: %s", context)
        return SessionContext(self.storage, context, metrics=self.metrics)

    def list_sessions(self, *, limit: int | None = None) -> list[SessionRecord]:
        return self.storage.list_sessions(limit=limit)

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

    def get_operations(self, session_id: str) -> list[OperationRecord]:
        return self.storage.get_operations(session_id)

    def undo_session(
        self,
        session_id: str,
        registry: OperationRegistry,
        *,
        mark_rolled_back: bool = True,
        allow_multiple_undo: bool = False,
    ) -> list[OperationRecord]:
        session = self.storage.get_session(session_id)
        if session is None:
            candidates = [item.session_id for item in self.storage.list_sessions(limit=5)]
            raise SessionNotFoundError(session_id, suggestions=candidates)

        if not allow_multiple_undo and session.status == SessionStatus.ROLLED_BACK:
            raise ValueError(
                f"Session {session_id} has already been rolled back. "
                "Use allow_multiple_undo=True to undo again."
            )

        operations = self.storage.get_operations(session_id)
        logger.info("Undoing session: %s", session_id)
        logger.debug("Found %d operations to undo", len(operations))
        start = perf_counter()

        for index, op in enumerate(reversed(operations), 1):
            logger.debug("Undoing operation %d/%d: %s", index, len(operations), op.op_type)
            executable = registry.build(op.op_type, op.payload)
            executable.undo()

        if mark_rolled_back and operations:
            self.storage.close_session(session_id=session_id, status=SessionStatus.ROLLED_BACK)

        if isinstance(self.storage, SQLiteStorage):
            self.storage.increment_undo_count(session_id)
        elif session is not None:
            session.undo_count += 1

        if self.metrics is not None:
            self.metrics.record_undo(perf_counter() - start)

        logger.info("Successfully undid session: %s", session_id)
        return operations


class CachedHistoryManager(HistoryManager):
    """History manager with lightweight caching for read-heavy access patterns."""

    @lru_cache(maxsize=128)
    def get_session(self, session_id: str) -> SessionRecord | None:
        return super().get_session(session_id)

    @lru_cache(maxsize=128)
    def get_operations(self, session_id: str) -> list[OperationRecord]:
        return super().get_operations(session_id)

    def _invalidate_cache(self) -> None:
        self.get_session.cache_clear()
        self.get_operations.cache_clear()

    def start_session(self, **context: Any) -> SessionContext:
        self._invalidate_cache()
        return super().start_session(**context)

    def undo_session(
        self,
        session_id: str,
        registry: OperationRegistry,
        *,
        mark_rolled_back: bool = True,
        allow_multiple_undo: bool = False,
    ) -> list[OperationRecord]:
        operations = super().undo_session(
            session_id,
            registry,
            mark_rolled_back=mark_rolled_back,
            allow_multiple_undo=allow_multiple_undo,
        )
        self._invalidate_cache()
        return operations
