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

from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any, AsyncIterator
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
from app.infra.db.project_routing import activate_project_routing


class QueryMetadataRepository(MetadataRepository):
    async def get_project_db_routing(self, project_id, role):
        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()
        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)