"""Async repository abstraction for CRUD data access."""

from __future__ import annotations

from collections.abc import Sequence
from typing import Protocol

from pydantic import BaseModel
from sqlalchemy.exc import IntegrityError
from sqlmodel import SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession

from auen.exceptions import ConflictError
from auen.types import PKValue, SelectQuery


class AsyncCrudRepository(Protocol):
    """Protocol for async CRUD data access."""

    async def get(self, session: AsyncSession, pk: PKValue) -> SQLModel | None: ...

    async def list(
        self,
        session: AsyncSession,
        *,
        query: SelectQuery[SQLModel] | None = None,
        offset: int = 0,
        limit: int = 50,
    ) -> Sequence[SQLModel]: ...

    async def count(
        self, session: AsyncSession, *, query: SelectQuery[SQLModel] | None = None
    ) -> int: ...

    async def create(self, session: AsyncSession, *, obj_in: BaseModel) -> SQLModel: ...

    async def update(
        self,
        session: AsyncSession,
        *,
        db_obj: SQLModel,
        obj_in: BaseModel,
        partial: bool = True,
    ) -> SQLModel: ...

    async def delete(self, session: AsyncSession, *, db_obj: SQLModel) -> None: ...

    async def bulk_create(
        self,
        session: AsyncSession,
        *,
        objects_in: Sequence[BaseModel],
    ) -> Sequence[SQLModel]: ...

    async def bulk_update(
        self,
        session: AsyncSession,
        *,
        items: Sequence[tuple[SQLModel, BaseModel]],
    ) -> Sequence[SQLModel]: ...

    async def bulk_delete(
        self,
        session: AsyncSession,
        *,
        db_objects: Sequence[SQLModel],
    ) -> None: ...


class AsyncSqlModelRepository:
    """Default async repository using SQLModel AsyncSession + select."""

    def __init__(self, model: type[SQLModel]) -> None:
        self.model = model

    async def get(self, session: AsyncSession, pk: PKValue) -> SQLModel | None:
        return await session.get(self.model, pk)

    async def list(
        self,
        session: AsyncSession,
        *,
        query: SelectQuery[SQLModel] | None = None,
        offset: int = 0,
        limit: int = 50,
    ) -> Sequence[SQLModel]:
        if query is None:
            query = select(self.model)
        stmt = query.offset(offset).limit(limit)
        result = await session.exec(stmt)
        return result.all()

    async def count(
        self, session: AsyncSession, *, query: SelectQuery[SQLModel] | None = None
    ) -> int:
        if query is None:
            query = select(self.model)
        from sqlalchemy import func

        stmt = select(func.count()).select_from(query.subquery())
        result = await session.exec(stmt)
        return result.one()

    async def create(self, session: AsyncSession, *, obj_in: BaseModel) -> SQLModel:
        data = obj_in.model_dump()
        db_obj = self.model(**data)
        session.add(db_obj)
        try:
            await session.commit()
        except IntegrityError as e:
            await session.rollback()
            raise ConflictError(detail="Integrity constraint violation") from e
        await session.refresh(db_obj)
        return db_obj

    async def update(
        self,
        session: AsyncSession,
        *,
        db_obj: SQLModel,
        obj_in: BaseModel,
        partial: bool = True,
    ) -> SQLModel:
        if partial:
            update_data = obj_in.model_dump(exclude_unset=True)
        else:
            update_data = obj_in.model_dump()
        for field, value in update_data.items():
            setattr(db_obj, field, value)
        session.add(db_obj)
        try:
            await session.commit()
        except IntegrityError as e:
            await session.rollback()
            raise ConflictError(detail="Integrity constraint violation") from e
        await session.refresh(db_obj)
        return db_obj

    async def delete(self, session: AsyncSession, *, db_obj: SQLModel) -> None:
        await session.delete(db_obj)
        try:
            await session.commit()
        except IntegrityError as e:
            await session.rollback()
            raise ConflictError(detail="Integrity constraint violation") from e

    async def bulk_create(
        self,
        session: AsyncSession,
        *,
        objects_in: Sequence[BaseModel],
    ) -> Sequence[SQLModel]:
        db_objects = [self.model(**obj.model_dump()) for obj in objects_in]
        session.add_all(db_objects)
        try:
            await session.commit()
        except IntegrityError as e:
            await session.rollback()
            raise ConflictError(detail="Integrity constraint violation") from e
        for obj in db_objects:
            await session.refresh(obj)
        return db_objects

    async def bulk_update(
        self,
        session: AsyncSession,
        *,
        items: Sequence[tuple[SQLModel, BaseModel]],
    ) -> Sequence[SQLModel]:
        results = []
        for db_obj, obj_in in items:
            update_data = obj_in.model_dump(exclude_unset=True)
            for field, value in update_data.items():
                setattr(db_obj, field, value)
            session.add(db_obj)
            results.append(db_obj)
        try:
            await session.commit()
        except IntegrityError as e:
            await session.rollback()
            raise ConflictError(detail="Integrity constraint violation") from e
        for obj in results:
            await session.refresh(obj)
        return results

    async def bulk_delete(
        self,
        session: AsyncSession,
        *,
        db_objects: Sequence[SQLModel],
    ) -> None:
        for obj in db_objects:
            await session.delete(obj)
        try:
            await session.commit()
        except IntegrityError as e:
            await session.rollback()
            raise ConflictError(detail="Integrity constraint violation") from e
