# tenant_db.py

from contextlib import asynccontextmanager
from contextvars import ContextVar, Token
from typing import AsyncIterator

from fastapi import HTTPException, Request
from sqlalchemy.ext.asyncio import AsyncSession

from .db_engine import TenantDatabaseManager
from .tenant_context import TenantContextConfig
from .exceptions import TenantDatabaseConfigError, TenantDatabaseError


_current_db: ContextVar[AsyncSession | None] = ContextVar(
    "current_db",
    default=None,
)


class TenantDbConfig(TenantContextConfig):
    def __init__(
        self,
        tenants_source: str,
        *,
        tenant_header: str = "X-Tenant",
        user_id_header: str = "X-UserId",
        request_id_header: str = "X-Request-Id",
        default_tenant: str | None = None,
        default_user_id: str | None = None,
        allowed_tenants: set[str] | None = None,
        strict_whitelist: bool = False,
        mode: str = "testing",
        echo_sql: bool = False,
        pool_size: int = 1,
        max_overflow: int = 2,
        pool_timeout: int = 5,
        pool_pre_ping: bool = True,
        app_name_prefix: str = "fastapi",
    ):
        super().__init__(
            tenants_source=tenants_source,
            tenant_header=tenant_header,
            user_id_header=user_id_header,
            request_id_header=request_id_header,
            default_tenant=default_tenant,
            default_user_id=default_user_id,
            allowed_tenants=allowed_tenants,
            strict_whitelist=strict_whitelist,
        )

        self.mode = mode
        self.echo_sql = echo_sql
        self.pool_size = pool_size
        self.max_overflow = max_overflow
        self.pool_timeout = pool_timeout
        self.pool_pre_ping = pool_pre_ping
        self.app_name_prefix = app_name_prefix

        self.db = TenantDatabaseManager(self)

    async def open_db(self, request: Request) -> AsyncIterator[None]:
        tenant = self.get_tenant(request)
        user_id = self.get_user_id(request)
        request_id = self.get_request_id(request)
        try:
            sessionmaker = await self.db.get_sessionmaker(tenant)
        except TenantDatabaseError as exc:
            raise HTTPException(
                status_code=503,
                detail=f"Database non disponibile per il tenant '{tenant}'",
            ) from exc
            
        with (
            self.tenant_context(tenant),
            self.user_id_context(user_id),
            self.request_id_context(request_id),
        ):
            async with sessionmaker() as db:
                token: Token = _current_db.set(db)
                try:
                    yield
                except Exception:
                    if db.in_transaction():
                        await db.rollback()
                    raise
                finally:
                    _current_db.reset(token)

    def current_db(self) -> AsyncSession:
        db = _current_db.get()
        if db is None:
            raise RuntimeError("Nessun DB attivo nel contesto corrente")
        return db

    async def commit(self) -> None:
        await self.current_db().commit()

    async def rollback(self) -> None:
        db = self.current_db()
        if db.in_transaction():
            await db.rollback()

    async def flush(self) -> None:
        await self.current_db().flush()

    @asynccontextmanager
    async def transaction(self):
        try:
            yield self.current_db()
            await self.commit()
        except Exception:
            await self.rollback()
            raise

    async def dispose(self) -> None:
        await self.db.dispose_all()