# src/database.py
import functools
import operator
import re
from contextlib import asynccontextmanager, contextmanager
from typing import Annotated, Any, AsyncGenerator

from fastapi import Depends
from pydantic import BaseModel
from sqlalchemy import create_engine, inspect
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.ext.declarative import declarative_base, declared_attr
from sqlalchemy.orm import DeclarativeBase, Session, object_session, sessionmaker
from sqlalchemy.sql.expression import true
from sqlalchemy_utils import get_mapper
from starlette.requests import Request

from src.config import SQLALCHEMY_COLLECTOR_URI, SQLALCHEMY_DATABASE_URI

engine = create_async_engine(SQLALCHEMY_DATABASE_URI, echo=False, future=True)
collector_engine = create_async_engine(SQLALCHEMY_COLLECTOR_URI, echo=False, future=True)



async_session = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autocommit=False,
    autoflush=False,
)

async_aeros_session = async_sessionmaker(
    collector_engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autocommit=False,
    autoflush=False,
)


async def get_aeros_db(request: Request):
    return request.state.aeros_db


def get_db(request: Request):
    return request.state.db


DbSession = Annotated[AsyncSession, Depends(get_db)]
CollectorDbSession = Annotated[AsyncSession, Depends(get_aeros_db)]


class Base(DeclarativeBase):
    @declared_attr.directive
    def __tablename__(cls) -> str:
        return resolve_table_name(cls.__name__)

    def dict(self):
        """Returns a dict representation of a model."""
        if hasattr(self, '__table__'):
            return {c.name: operator.attrgetter(c.name)(self) for c in self.__table__.columns}
        return {}

class CollectorBase(DeclarativeBase):
    @declared_attr.directive
    def __tablename__(cls) -> str:
        return resolve_table_name(cls.__name__)

    def dict(self):
        """Returns a dict representation of a model."""
        if hasattr(self, '__table__'):
            return {c.name: operator.attrgetter(c.name)(self) for c in self.__table__.columns}
        return {}

@asynccontextmanager
async def get_main_session():
    session = async_session()
    try:
        yield session
        await session.commit()
    except:
        await session.rollback()
        raise
    finally:
        await session.close()

@asynccontextmanager
async def get_collector_session():
    session = async_aeros_session()
    try:
        yield session
        await session.commit()
    except:
        await session.rollback()
        raise
    finally:
        await session.close()

def resolve_table_name(name):
    """Resolves table names to their mapped names."""
    names = re.split("(?=[A-Z])", name)  # noqa
    return "_".join([x.lower() for x in names if x])


def get_class_by_tablename(table_fullname: str) -> Any:
    """Return class reference mapped to table."""

    def _find_class(name):
        for c in Base.registry._class_registry.values():
            if hasattr(c, "__table__"):
                if c.__table__.fullname.lower() == name.lower():
                    return c

    mapped_name = resolve_table_name(table_fullname)
    mapped_class = _find_class(mapped_name)

    return mapped_class
