"""Resolve authoritative metadata and project routing for each MCP request."""

from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any, AsyncIterator
from uuid import UUID
from sqlalchemy.ext.asyncio import AsyncSession
from psycopg import AsyncConnection
from app.auth.project_dependencies import ProjectContext
from app.infra.db.project_routing import ActiveProjectRouting
from fastapi import HTTPException
from mcp.server.auth.provider import AccessToken

from app.auth.metadata_dependencies import get_current_metadata_user
from app.auth.project_dependencies import (
    resolve_project_context,
    resolve_project_business_routing,
    resolve_project_simulation_routing,
    get_project_pg_connection,
    get_project_timescale_connection,
)
from app.infra.db.metadb.database import get_metadata_session
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository, ProjectDbRouting
from app.infra.db.project_routing import activate_project_routing


class QueryMetadataRepository(MetadataRepository):
    def __init__(self, session: AsyncSession):
        super().__init__(session)
        self._routing: dict[tuple[UUID, str], ProjectDbRouting | None] = {}

    async def get_project_db_routing(self, project_id: UUID, role: str) -> ProjectDbRouting | None:
        key = (project_id, role)
        if key in self._routing:
            return self._routing[key]
        routing = await super().get_project_db_routing(project_id, role)
        # Existing connection generators obtain routing before awaiting a pool.
        # End that implicit SELECT transaction here, before the network wait.
        await self.session.commit()
        self._routing[key] = routing
        return routing


@dataclass
class QueryScope:
    project: ProjectContext
    user: Any
    connection: AsyncConnection
    routing: ActiveProjectRouting
    metadata: QueryMetadataRepository

    @asynccontextmanager
    async def simulation_routing(self) -> AsyncIterator[ActiveProjectRouting]:
        """Resolve both databases before a native calculation can modify its model."""
        routing = await resolve_project_simulation_routing(self.project, self.metadata)
        await self.metadata.session.commit()
        with activate_project_routing(routing):
            yield routing

    @asynccontextmanager
    async def timescale(self):
        # Metadata repository here is used only to resolve a known project DB.
        async with asynccontextmanager(get_project_timescale_connection)(
            self.project, self.metadata
        ) as conn:
            await self.metadata.session.commit()
            yield conn


@asynccontextmanager
async def project_scope(principal: AccessToken, project_id: str):
    if not principal.claims:
        raise HTTPException(401, "MCP identity is missing")
    async with asynccontextmanager(get_metadata_session)() as session:
        metadata = QueryMetadataRepository(session)
        user = await get_current_metadata_user(principal.claims, metadata)
        project = await resolve_project_context(project_id, user, metadata)
        routing = await resolve_project_business_routing(project, metadata)
        # End the metadata read before any external query/IO waits.
        await session.commit()
        async with asynccontextmanager(get_project_pg_connection)(
            project, metadata
        ) as conn:
            await session.commit()
            with activate_project_routing(routing):
                yield QueryScope(project, user, conn, routing, metadata)