import re
import operator
from contextlib import asynccontextmanager
from typing import Annotated, Any
from fastapi import Depends
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy.orm import DeclarativeBase
from starlette.requests import Request
from src.config import SQLALCHEMY_DATABASE_URI, SQLALCHEMY_COLLECTOR_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_collector_session = async_sessionmaker(collector_engine, class_=AsyncSession, expire_on_commit=False, autocommit=False, autoflush=False)

async def get_collector_db(request: Request):
    return request.state.collector_db

def get_db(request: Request):
    return request.state.db
DbSession = Annotated[AsyncSession, Depends(get_db)]
CollectorDbSession = Annotated[AsyncSession, Depends(get_collector_db)]

class Base(DeclarativeBase):

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

    def dict(self):
        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):
        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_collector_session()
    try:
        yield session
        await session.commit()
    except:
        await session.rollback()
        raise
    finally:
        await session.close()

def resolve_table_name(name):
    names = re.split('(?=[A-Z])', name)
    return '_'.join([x.lower() for x in names if x])

def get_class_by_tablename(table_fullname: str) -> Any:

    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