from __future__ import annotations

from typing import AsyncGenerator

from loguru import logger
from omegaconf import OmegaConf
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine

class Database:
    def __init__(
        self,
        name: str,
        postgres_user: str,
        postgres_password: str,
        postgres_host: str,
        postgres_port: int,
        postgres_db: str,
    ) -> None:
        self.name = name
        self.postgres_user = postgres_user
        self.postgres_password = postgres_password
        self.postgres_host = postgres_host
        self.postgres_port = postgres_port
        self.postgres_db = postgres_db
        self.engine = None
        self.session_factory = None

    @property
    def database_url(self) -> str:
        return f"postgresql+asyncpg://{self.postgres_user}:{self.postgres_password}@{self.postgres_host}:{self.postgres_port}/{self.postgres_db}"

    @property
    def database_base_url(self) -> str:
        return f"postgresql+asyncpg://{self.postgres_user}:{self.postgres_password}@{self.postgres_host}:{self.postgres_port}/postgres"

    async def init(self) -> None:
        if self.engine is None:
            self.engine = await self.create_engine()
            self.session_factory = async_sessionmaker(bind=self.engine, class_=AsyncSession, expire_on_commit=False)

    async def close(self) -> None:
        if self.engine:
            await self.engine.dispose()
            self.engine = None
            self.session_factory = None
            logger.info("Database connection closed")

    async def create_engine(self) -> AsyncEngine:
        if not all([self.postgres_user, self.postgres_host, self.postgres_port, self.postgres_db]):
            raise ValueError("Missing required database configuration. Please check your .env file or configuration.")
        logger.info(f"Creating database engine: {self.database_url}")
        return create_async_engine(self.database_url)

    async def get_async_session(self) -> AsyncGenerator[AsyncSession, None]:
        if self.session_factory is None:
            await self.init()
        
        if self.session_factory is None:
            raise RuntimeError("Session factory is not initialized")
        
        async with self.session_factory() as session:
            try:
                yield session
            except Exception:
                await session.rollback()
                raise
            finally:
                await session.close()
            yield session

    async def get_session_factory(self) -> async_sessionmaker[AsyncSession]:
        if self.engine is None:
            await self.init()

        if self.session_factory is None:
            raise RuntimeError("Session factory is not initialized")

        return self.session_factory

class DatabaseManager:
    def __init__(
        self,
        databases: list[Database],
    ) -> None:
        self.databases = databases

    def get_database(self, name: str) -> Database | None:
        for database in self.databases:
            if database.name == name:
                return database
        return None

    def get_default_database(self) -> Database | None:
        if len(self.databases) > 0:
            return self.databases[0]
        return None

    @staticmethod
    def from_config(config_path: str = None, config = None) -> DatabaseManager:
        if config is None:
            config = OmegaConf.to_object(OmegaConf.load(config_path))

        databases_config = config.get('databases', None)
        databases = []
        if databases_config is not None:
            for database_config in databases_config:
                if 'postgres_db' in database_config and database_config['postgres_db'] is not None:
                    databases.append(Database(
                        name=database_config['name'],
                        postgres_user=database_config['postgres_user'],
                        postgres_password=database_config['postgres_password'],
                        postgres_host=database_config['postgres_host'],
                        postgres_port=database_config['postgres_port'],
                        postgres_db=database_config['postgres_db'],
                    ))
        return DatabaseManager(databases=databases)

def create_database_manager(config_path: str = None, config = None) -> DatabaseManager:
    return DatabaseManager.from_config(config_path, config)
