"""SQLAlchemy-backed model-capability store."""

from __future__ import annotations

from datetime import UTC, datetime
from typing import Any

from sqlalchemy import delete, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from deerflow.persistence.model_capabilities.base import ModelCapabilityStore
from deerflow.persistence.model_capabilities.model import ModelCapabilityRow


def _to_dict(row: ModelCapabilityRow) -> dict[str, Any]:
    return {
        "model_name": row.model_name,
        "supports_reasoning_effort": bool(row.supports_reasoning_effort),
        "default_reasoning_effort": row.default_reasoning_effort or "xhigh",
        "updated_by": row.updated_by,
        "updated_at": row.updated_at.isoformat() if isinstance(row.updated_at, datetime) else row.updated_at,
    }


class ModelCapabilityRepository(ModelCapabilityStore):
    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
        self._sf = session_factory

    async def list_capabilities(self) -> list[dict[str, Any]]:
        async with self._sf() as session:
            result = await session.execute(select(ModelCapabilityRow).order_by(ModelCapabilityRow.model_name.asc()))
            return [_to_dict(row) for row in result.scalars()]

    async def get_capability(self, model_name: str) -> dict[str, Any] | None:
        async with self._sf() as session:
            row = await session.get(ModelCapabilityRow, model_name)
            return _to_dict(row) if row is not None else None

    async def set_reasoning_effort_support(
        self,
        model_name: str,
        *,
        supports_reasoning_effort: bool,
        default_reasoning_effort: str | None = None,
        updated_by: str | None = None,
    ) -> dict[str, Any]:
        async with self._sf() as session:
            row = await session.get(ModelCapabilityRow, model_name)
            if row is None:
                row = ModelCapabilityRow(model_name=model_name)
                session.add(row)
            row.supports_reasoning_effort = bool(supports_reasoning_effort)
            if default_reasoning_effort is not None:
                row.default_reasoning_effort = default_reasoning_effort
            row.updated_by = updated_by
            row.updated_at = datetime.now(UTC)
            await session.commit()
            await session.refresh(row)
            return _to_dict(row)

    async def seed_reasoning_effort_support(
        self,
        model_name: str,
        *,
        updated_by: str | None = None,
    ) -> bool:
        """Insert a legacy ``true`` without ever overwriting an existing row.

        The primary-key conflict handles concurrent Gateway startup safely: if
        another worker or an administrator creates the row after our read, the
        losing seed transaction rolls back and preserves that authoritative
        value, including ``false``.
        """

        async with self._sf() as session:
            if await session.get(ModelCapabilityRow, model_name) is not None:
                return False
            session.add(
                ModelCapabilityRow(
                    model_name=model_name,
                    supports_reasoning_effort=True,
                    updated_by=updated_by,
                )
            )
            try:
                await session.commit()
            except IntegrityError:
                await session.rollback()
                return False
            return True

    async def delete_capability(self, model_name: str) -> None:
        async with self._sf() as session:
            await session.execute(delete(ModelCapabilityRow).where(ModelCapabilityRow.model_name == model_name))
            await session.commit()