# db_engine.py
import asyncio
from typing import Any

from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.pool import NullPool


class TenantDatabaseManager:
    def __init__(self, config: Any):
        self.config = config
        self._engines: dict[str, AsyncEngine] = {}
        self._db_factories: dict[str, async_sessionmaker[AsyncSession]] = {}
        self._locks: dict[str, asyncio.Lock] = {}
        self._locks_guard = asyncio.Lock()

    async def _get_lock(self, tenant: str) -> asyncio.Lock:
        lock = self._locks.get(tenant)
        if lock is not None:
            return lock

        async with self._locks_guard:
            lock = self._locks.get(tenant)
            if lock is None:
                lock = asyncio.Lock()
                self._locks[tenant] = lock
            return lock

    def _build_connect_args(self, tenant: str) -> dict[str, Any]:
        return {
            "options": f"-c application_name={self.config.app_name_prefix}:{tenant}"
        }

    def _create_engine(self, tenant: str) -> AsyncEngine:
        try:
            db_url = self.config.service.get(tenant).db.url.get_secret_value()
            return create_async_engine(
                db_url,
                connect_args=self._build_connect_args(tenant),
                **self._engine_args(),
            )
        except Exception as exc:
            raise RuntimeError(
                f"Failed to create database engine for tenant '{tenant}'"
            ) from exc

    def _engine_args(self) -> dict[str, Any]:
        if self.config.mode == "testing":
            return {
                "echo": self.config.echo_sql,
                "poolclass": NullPool,
            }

        return {
            "echo": self.config.echo_sql,
            "pool_pre_ping": self.config.pool_pre_ping,
            "pool_size": self.config.pool_size,
            "max_overflow": self.config.max_overflow,
            "pool_timeout": self.config.pool_timeout,
        }

    async def get_engine(self, tenant: str) -> AsyncEngine:
        engine = self._engines.get(tenant)
        if engine is not None:
            return engine

        lock = await self._get_lock(tenant)
        async with lock:
            engine = self._engines.get(tenant)
            if engine is None:
                engine = self._create_engine(tenant)
                self._engines[tenant] = engine
            return engine

    async def get_sessionmaker(self, tenant: str) -> async_sessionmaker[AsyncSession]:
        factory = self._db_factories.get(tenant)
        if factory is not None:
            return factory

        lock = await self._get_lock(tenant)
        async with lock:
            factory = self._db_factories.get(tenant)
            if factory is not None:
                return factory

            engine = self._engines.get(tenant)
            if engine is None:
                engine = self._create_engine(tenant)
                self._engines[tenant] = engine

            factory = async_sessionmaker(
                engine,
                class_=AsyncSession,
                expire_on_commit=False,
            )
            self._db_factories[tenant] = factory
            return factory

    async def dispose_tenant(self, tenant: str) -> None:
        engine = self._engines.pop(tenant, None)
        self._db_factories.pop(tenant, None)
        self._locks.pop(tenant, None)

        if engine is not None:
            await engine.dispose()

    async def dispose_all(self) -> None:
        engines = list(self._engines.values())

        self._engines.clear()
        self._db_factories.clear()
        self._locks.clear()

        for engine in engines:
            await engine.dispose()