from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.model.configuration import ModelEndpointConfig


class ModelRouterService:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session
        self._unhealthy: set[str] = set()

    def mark_health(self, endpoint_code: str, healthy: bool) -> None:
        if healthy:
            self._unhealthy.discard(endpoint_code)
        else:
            self._unhealthy.add(endpoint_code)

    async def select_endpoint(
        self, *, agent_type: str, task_type: str, required_capability: str | None = None
    ) -> ModelEndpointConfig | None:
        statement = select(ModelEndpointConfig).where(ModelEndpointConfig.status == "active")
        endpoints = list(await self.session.scalars(statement))
        if required_capability:
            endpoints = [e for e in endpoints if required_capability in e.capabilities]
        endpoints = [e for e in endpoints if e.endpoint_code not in self._unhealthy]
        return endpoints[0] if endpoints else None

    async def select_with_fallback(
        self,
        *,
        agent_type: str,
        task_type: str,
        fallback_codes: list[str] | None = None,
        required_capability: str | None = None,
    ) -> list[ModelEndpointConfig]:
        endpoints = list(
            await self.session.scalars(
                select(ModelEndpointConfig).where(ModelEndpointConfig.status == "active")
            )
        )
        if required_capability:
            endpoints = [e for e in endpoints if required_capability in e.capabilities]
        endpoints = [e for e in endpoints if e.endpoint_code not in self._unhealthy]
        by_code = {endpoint.endpoint_code: endpoint for endpoint in endpoints}
        ordered: list[ModelEndpointConfig] = []
        for code in fallback_codes or []:
            endpoint = by_code.get(code)
            if endpoint is not None and endpoint not in ordered:
                ordered.append(endpoint)
        for endpoint in endpoints:
            if endpoint not in ordered:
                ordered.append(endpoint)
        return ordered