"""Unit of Work pattern implementation for CleverAgents.

Based on ADR-007 (Repository Pattern) - provides transaction management.
"""

from __future__ import annotations

from collections.abc import Callable, Generator
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any

from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker

from cleveragents.infrastructure.database.engine_cache import MEMORY_ENGINES
from cleveragents.infrastructure.database.repositories import (
    ActionRepository,
    ActorRepository,
    ChangeRepository,
    CheckpointRepository,
    ContextRepository,
    CorrectionAttemptRepository,
    DebugAttemptRepository,
    DecisionRepository,
    LifecyclePlanRepository,
    PlanRepository,
    ProjectRepository,
)

if TYPE_CHECKING:
    from sqlalchemy.engine import Engine


class UnitOfWork:
    """Unit of Work pattern for managing database transactions.

    Ensures all operations within a unit of work are committed together
    or rolled back on failure.
    """

    def __init__(
        self,
        database_url: str,
        prompt_for_migration: Callable[[str], bool] | None = None,
        require_confirmation: bool = True,
    ) -> None:
        """Initialize Unit of Work with database connection.

        Args:
            database_url: SQLAlchemy database URL
            prompt_for_migration: Optional callback to confirm migrations
                before applying them automatically.
            require_confirmation: When False, migrations are applied without
                prompting the user.  Set to False when the caller has
                already obtained blanket approval (e.g. ``--yes`` flag).
        """
        self.database_url = database_url
        self._engine: Engine | None = None
        self._session_factory: sessionmaker[Session] | None = None
        self._database_initialized = False
        self._prompt_for_migration = prompt_for_migration
        self._require_confirmation = require_confirmation

    @property
    def engine(self) -> Engine:
        """Get or create database engine."""
        if self._engine is None:
            self._ensure_database_initialized()
            # For SQLite, we need to ensure proper transaction isolation
            if self.database_url.startswith("sqlite"):
                # For in-memory SQLite databases, we must reuse the same engine
                # across all instances to ensure the in-memory database persists.
                # Each new engine to :memory: creates a separate database.
                if self.database_url == "sqlite:///:memory:":
                    # Check if we have a cached engine for this in-memory database
                    if self.database_url not in MEMORY_ENGINES:
                        MEMORY_ENGINES[self.database_url] = create_engine(
                            self.database_url,
                            echo=False,  # Set to True for SQL debugging
                            future=True,  # Use SQLAlchemy 2.0 style
                            isolation_level="SERIALIZABLE",
                            connect_args={"check_same_thread": False},
                        )
                    self._engine = MEMORY_ENGINES[self.database_url]
                else:
                    # File-based SQLite
                    # Use SERIALIZABLE isolation for SQLite to ensure proper rollback
                    self._engine = create_engine(
                        self.database_url,
                        echo=False,  # Set to True for SQL debugging
                        future=True,  # Use SQLAlchemy 2.0 style
                        isolation_level="SERIALIZABLE",
                        connect_args={"check_same_thread": False},
                    )
            else:
                self._engine = create_engine(
                    self.database_url,
                    echo=False,  # Set to True for SQL debugging
                    future=True,  # Use SQLAlchemy 2.0 style
                )
        return self._engine

    @property
    def session_factory(self) -> sessionmaker[Session]:
        """Get or create session factory."""
        if self._session_factory is None:
            self._session_factory = sessionmaker(
                bind=self.engine,
                expire_on_commit=False,
                autoflush=False,  # Don't auto-flush to database
                autocommit=False,  # Explicit transactions only
                class_=Session,
            )
        return self._session_factory

    @contextmanager
    def transaction(self) -> Generator[UnitOfWorkContext]:
        """Create a new transaction context.

        Usage:
            with uow.transaction() as ctx:
                project = ctx.projects.create(...)
                plan = ctx.plans.create(...)
                # Both operations committed together

        If an exception occurs, all changes are rolled back.
        """
        session = self.session_factory()
        try:
            context = UnitOfWorkContext(session)
            yield context
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def _ensure_database_initialized(self, force: bool = False) -> None:
        """Run migrations if needed before database access."""
        if self._database_initialized and not force:
            return

        from cleveragents.infrastructure.database.migration_runner import (
            MigrationRunner,
        )

        runner = MigrationRunner(self.database_url)
        runner.init_or_upgrade(
            require_confirmation=self._require_confirmation,
            prompt_for_migration=self._prompt_for_migration,
        )
        self._database_initialized = True

    def init_database(self) -> None:
        """Initialize database schema.

        Runs Alembic migrations to ensure database is up to date.
        Should be called once on application startup.
        """

        # Use migration runner to handle database setup
        self._ensure_database_initialized(force=True)

        # For in-memory databases, we must keep the engine alive so the
        # database persists. For file-based databases, dispose to ensure
        # the database file is properly synchronized after migrations.
        # Don't dispose in-memory databases as each new engine is a new database
        if self._engine is not None and self.database_url != "sqlite:///:memory:":
            self._engine.dispose()
            self._engine = None
            self._session_factory = None


class UnitOfWorkContext:
    """Context for a unit of work transaction.

    Provides access to repositories within a transaction.
    """

    def __init__(self, session: Session) -> None:
        """Initialize context with database session.

        Args:
            session: SQLAlchemy session for this transaction
        """
        self._session = session
        self._projects: ProjectRepository | None = None
        self._plans: PlanRepository | None = None
        self._contexts: ContextRepository | None = None
        self._changes: ChangeRepository | None = None
        self._debug_attempts: DebugAttemptRepository | None = None
        self._actors: ActorRepository | None = None
        self._actions: ActionRepository | None = None
        self._lifecycle_plans: LifecyclePlanRepository | None = None
        self._decisions: DecisionRepository | None = None
        self._checkpoints: CheckpointRepository | None = None
        self._correction_attempts: CorrectionAttemptRepository | None = None

    def _session_factory(self) -> Session:
        """Return the transaction's session for factory-pattern repositories."""
        return self._session

    @property
    def projects(self) -> ProjectRepository:
        """Get project repository for this transaction."""
        if self._projects is None:
            self._projects = ProjectRepository(self._session)
        return self._projects

    @property
    def plans(self) -> PlanRepository:
        """Get legacy plan repository for this transaction.

        .. deprecated::
            Use :attr:`lifecycle_plans` for v3 lifecycle plans.
            This accessor targets the legacy ``plans`` table and will
            be removed in a future release.
        """
        if self._plans is None:
            self._plans = PlanRepository(self._session)
        return self._plans

    @property
    def actions(self) -> ActionRepository:
        """Get v3 action repository for this transaction.

        Uses the session-factory pattern required by ActionRepository.
        """
        if self._actions is None:
            self._actions = ActionRepository(session_factory=self._session_factory)
        return self._actions

    @property
    def lifecycle_plans(self) -> LifecyclePlanRepository:
        """Get v3 lifecycle plan repository for this transaction.

        Uses the session-factory pattern required by LifecyclePlanRepository.
        """
        if self._lifecycle_plans is None:
            self._lifecycle_plans = LifecyclePlanRepository(
                session_factory=self._session_factory,
            )
        return self._lifecycle_plans

    @property
    def checkpoints(self) -> CheckpointRepository:
        """Get checkpoint repository for this transaction.

        Uses the session-factory pattern required by CheckpointRepository.
        """
        if self._checkpoints is None:
            self._checkpoints = CheckpointRepository(
                session_factory=self._session_factory,
            )
        return self._checkpoints

    @property
    def contexts(self) -> ContextRepository:
        """Get context repository for this transaction."""
        if self._contexts is None:
            self._contexts = ContextRepository(self._session)
        return self._contexts

    @property
    def changes(self) -> ChangeRepository:
        """Get change repository for this transaction."""
        if self._changes is None:
            self._changes = ChangeRepository(self._session)
        return self._changes

    @property
    def debug_attempts(self) -> DebugAttemptRepository:
        """Get debug attempt repository for this transaction."""
        if self._debug_attempts is None:
            self._debug_attempts = DebugAttemptRepository(self._session)
        return self._debug_attempts

    @property
    def actors(self) -> ActorRepository:
        """Get actor repository for this transaction."""
        if self._actors is None:
            self._actors = ActorRepository(self._session)
        return self._actors

    @property
    def decisions(self) -> DecisionRepository:
        """Get decision repository for this transaction.

        Uses the session-factory pattern required by DecisionRepository.
        """
        if self._decisions is None:
            self._decisions = DecisionRepository(
                session_factory=self._session_factory,
            )
        return self._decisions

    @property
    def correction_attempts(self) -> CorrectionAttemptRepository:
        """Get correction attempt repository for this transaction.

        Uses the session-factory pattern required by CorrectionAttemptRepository.
        """
        if self._correction_attempts is None:
            self._correction_attempts = CorrectionAttemptRepository(
                session_factory=self._session_factory,
            )
        return self._correction_attempts

    def add(self, entity: Any) -> None:
        """Add an entity to the session.

        Args:
            entity: Entity to add
        """
        self._session.add(entity)

    def flush(self) -> None:
        """Flush pending changes to database.

        Sends SQL to database but doesn't commit transaction.
        Useful for getting auto-generated IDs.
        """
        self._session.flush()

    def refresh(self, entity: Any) -> None:
        """Refresh an entity from database.

        Args:
            entity: Entity to refresh
        """
        self._session.refresh(entity)
