import inspect
from collections.abc import Iterator
from contextlib import contextmanager
from functools import wraps
from typing import Any, Callable, ParamSpec, TypeVar, cast

from sqlalchemy import Integer, create_engine
from sqlalchemy.orm import (
    DeclarativeBase,
    Mapped,
    Session,
    mapped_column,
    sessionmaker,
)

from vdash.config import Environment, Settings


class BaseModel(DeclarativeBase):
    __abstract__ = True

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)


class Database:
    def __init__(self, settings: Settings):
        self.engine = create_engine(
            settings.db_url, echo=settings.environment == Environment.DEV
        )

        self.session_factory = sessionmaker(
            bind=self.engine,
            class_=Session,
            expire_on_commit=False,
            autoflush=False,
            autocommit=False,
        )

    def disconnect(self) -> None:
        self.engine.dispose()

    @contextmanager
    def get_session(self) -> Iterator[Session]:
        session = self.session_factory()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()


P = ParamSpec("P")
R = TypeVar("R")


def inject_session(func: Callable[P, R]) -> Callable[P, R]:
    if inspect.iscoroutinefunction(func):

        @wraps(func)
        async def async_wrapper(
            self: Any, *args: Any, session: Session | None = None, **kwargs: Any
        ) -> Any:
            if session is not None:
                return await func(self, *args, session=session, **kwargs)  # type: ignore
            if self._db is None:
                raise RuntimeError("Database is not set")
            with self._db.get_session() as db_session:
                return await func(self, *args, session=db_session, **kwargs)  # type: ignore

        return cast(Callable[P, R], async_wrapper)

    @wraps(func)
    def sync_wrapper(
        self: Any, *args: Any, session: Session | None = None, **kwargs: Any
    ) -> Any:
        if session is not None:
            return func(self, *args, session=session, **kwargs)  # type: ignore
        if self._db is None:
            raise RuntimeError("Database is not set")
        with self._db.get_session() as db_session:
            return func(self, *args, session=db_session, **kwargs)  # type: ignore

    return cast(Callable[P, R], sync_wrapper)
