"""Updater for repository update operations."""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, TypeVar, cast
from uuid import UUID

import sqlalchemy as sa
from sqlalchemy.engine import CursorResult

from ai.backend.manager.models.base import Base

from .integrity import match_integrity_error, parse_integrity_error
from .types import IntegrityErrorCheck, QueryCondition

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession as SASession

TRow = TypeVar("TRow", bound=Base)


class UpdaterSpec[TRow: Base](ABC):
    """Abstract base class defining values to update for single-row updates.

    Implementations specify what to update by providing:
    - row_class property for target table and PK detection
    - build_values() for column-value mapping
    """

    @property
    @abstractmethod
    def row_class(self) -> type[TRow]:
        """Return the ORM class for table access and PK detection."""
        raise NotImplementedError

    @abstractmethod
    def build_values(self) -> dict[str, Any]:
        """Build column name to value mapping for update.

        Note: This method should only be called internally by execute_updater
        or apply_to_row. External code should use apply_to_row() instead.

        Returns:
            Dict mapping column names to values
        """
        raise NotImplementedError

    @property
    def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]:
        """Return integrity error checks for declarative error matching.

        Override in subclasses to map constraint violations to domain errors.
        Default returns empty sequence (no checks, fallback behavior).
        """
        return ()

    def apply_to_row(self, row: TRow) -> None:
        """Apply update values to a row object via setattr.

        This method encapsulates the build_values() call and applies
        updates to the row without exposing the internal dict structure.

        Args:
            row: The ORM row object to apply updates to
        """
        for field, value in self.build_values().items():
            setattr(row, field, value)


class BatchUpdaterSpec[TRow: Base](ABC):
    """Abstract base class defining values to update for batch updates.

    Implementations specify what to update by providing:
    - row_class property for target table access
    - build_values() for column-value mapping
    """

    @property
    @abstractmethod
    def row_class(self) -> type[TRow]:
        """Return the ORM class for table access."""
        raise NotImplementedError

    @abstractmethod
    def build_values(self) -> dict[str, Any]:
        """Build column name to value mapping for update.

        Returns:
            Dict mapping column names to values
        """
        raise NotImplementedError

    @property
    def integrity_error_checks(self) -> Sequence[IntegrityErrorCheck]:
        """Return integrity error checks for declarative error matching.

        Override in subclasses to map constraint violations to domain errors.
        Default returns empty sequence (no checks, fallback behavior).
        """
        return ()


@dataclass
class Updater[TRow: Base]:
    """Bundles updater spec with target info for single-row update operations.

    Attributes:
        spec: UpdaterSpec implementation defining row_class and values to update.
        pk_value: Primary key value to identify the target row.
    """

    spec: UpdaterSpec[TRow]
    pk_value: UUID | str | int


@dataclass
class BatchUpdater[TRow: Base]:
    """Bundles batch updater spec with conditions for batch update operations.

    Attributes:
        spec: BatchUpdaterSpec implementation defining row_class and values to update.
        conditions: List of QueryCondition factories for WHERE clause (AND combined).
    """

    spec: BatchUpdaterSpec[TRow]
    conditions: list[QueryCondition]


@dataclass
class UpdaterResult[TRow: Base]:
    """Result of executing a single-row update operation."""

    row: TRow


@dataclass
class BatchUpdaterResult:
    """Result of executing a batch update operation."""

    updated_count: int


@dataclass
class BulkUpdaterError[TRow: Base]:
    """Error information for a failed bulk update operation.

    Contains the spec that failed and the exception for debugging.
    Follows the same pattern as BulkCreatorError.

    Attributes:
        spec: The UpdaterSpec that failed
        exception: The exception that occurred
        index: Original position in specs list for traceability
    """

    spec: UpdaterSpec[TRow]
    exception: Exception
    index: int


async def execute_updater[TRow: Base](
    db_sess: SASession,
    updater: Updater[TRow],
) -> UpdaterResult[TRow] | None:
    """Execute UPDATE for a single row by primary key.

    Args:
        db_sess: Database session (must be writable)
        updater: Updater containing spec and pk_value

    Returns:
        UpdaterResult containing the updated row, or None if no row matched
        or if there are no values to update

    Example:
        class SessionStatusUpdaterSpec(UpdaterSpec[SessionRow]):
            def __init__(self, new_status: str):
                self._new_status = new_status

            @property
            def row_class(self) -> type[SessionRow]:
                return SessionRow

            def build_values(self) -> dict[str, Any]:
                return {"status": self._new_status}

        updater = Updater(
            spec=SessionStatusUpdaterSpec("terminated"),
            pk_value=session_id,
        )
        result = await execute_updater(db_sess, updater)
        if result:
            print(result.row.status)  # "terminated"
    """
    row_class = updater.spec.row_class
    table = row_class.__table__
    pk_columns = list(table.primary_key.columns)
    values = updater.spec.build_values()

    # Return None if there are no values to update
    if not values:
        return None

    if len(pk_columns) != 1:
        raise ValueError("Updater only supports single-column primary keys")

    update_stmt = (
        sa.update(table)
        .values(values)
        .where(pk_columns[0] == updater.pk_value)
        .returning(*table.columns)
    )

    # Use from_statement to properly construct ORM objects
    # This allows SQLAlchemy to handle column mapping correctly
    select_stmt = sa.select(row_class).from_statement(update_stmt)
    try:
        result = await db_sess.execute(select_stmt)
    except sa.exc.IntegrityError as e:
        parsed = parse_integrity_error(e)
        match_integrity_error(parsed, updater.spec.integrity_error_checks)
    updated_row = result.scalar_one_or_none()

    if updated_row is None:
        return None

    return UpdaterResult(row=updated_row)


async def execute_batch_updater[TRow: Base](
    db_sess: SASession,
    updater: BatchUpdater[TRow],
) -> BatchUpdaterResult:
    """Execute UPDATE with multiple conditions (AND combined).

    Args:
        db_sess: Database session (must be writable)
        updater: BatchUpdater containing spec and conditions for the batch update

    Returns:
        BatchUpdaterResult with count of updated rows

    Example:
        class ItemStatusBatchUpdaterSpec(BatchUpdaterSpec[ItemRow]):
            def __init__(self, new_status: str):
                self._new_status = new_status

            @property
            def row_class(self) -> type[ItemRow]:
                return ItemRow

            def build_values(self) -> dict[str, Any]:
                return {"status": self._new_status}

        updater = BatchUpdater(
            spec=ItemStatusBatchUpdaterSpec("processed"),
            conditions=[lambda: ItemRow.__table__.c.status == "pending"],
        )
        result = await execute_batch_updater(db_sess, updater)
        print(result.updated_count)  # N
    """
    spec = updater.spec
    row_class = spec.row_class
    table = row_class.__table__
    values = spec.build_values()

    stmt = sa.update(table).values(values)
    for condition in updater.conditions:
        stmt = stmt.where(condition())

    try:
        result = await db_sess.execute(stmt)
    except sa.exc.IntegrityError as e:
        parsed = parse_integrity_error(e)
        match_integrity_error(parsed, updater.spec.integrity_error_checks)
    return BatchUpdaterResult(updated_count=cast(CursorResult[Any], result).rowcount)
