import functools
import typing as t

import sqlalchemy as sa
from ellar.common import IExecutionContext, IModuleSetup, Module, middleware
from ellar.core import Config, DynamicModule, ModuleBase, ModuleSetup
from ellar.di import ProviderConfig, request_or_transient_scope
from ellar.events import app_context_teardown
from ellar.utils.importer import get_main_directory_by_stack
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
)
from sqlalchemy.orm import Session

from ellar_sql.services import EllarSQLService

from .cli import DBCommands
from .schemas import MigrationOption, SQLAlchemyConfig


def _invalid_configuration(message: str) -> t.Callable:
    def _raise_exception():
        raise RuntimeError(message)

    return _raise_exception


@Module(commands=[DBCommands])
class EllarSQLModule(ModuleBase, IModuleSetup):
    @middleware()
    async def session_middleware(
        cls, context: IExecutionContext, call_next: t.Callable[..., t.Coroutine]
    ):
        connection = context.switch_to_http_connection().get_client()

        db_session = connection.service_provider.get(EllarSQLService)
        session = db_session.session_factory()

        connection.state.session = session

        try:
            await call_next()
        except Exception as ex:
            res = session.rollback()
            if isinstance(res, t.Coroutine):
                await res
            raise ex

    @classmethod
    async def _on_application_tear_down(cls, db_service: EllarSQLService) -> None:
        res = db_service.session_factory.remove()
        if isinstance(res, t.Coroutine):
            await res

    @classmethod
    def setup(
        cls,
        *,
        databases: t.Union[str, t.Dict[str, t.Any]],
        migration_options: t.Union[t.Dict[str, t.Any], MigrationOption],
        session_options: t.Optional[t.Dict[str, t.Any]] = None,
        engine_options: t.Optional[t.Dict[str, t.Any]] = None,
        models: t.Optional[t.List[str]] = None,
        echo: bool = False,
        root_path: t.Optional[str] = None,
    ) -> "DynamicModule":
        """
        Configures EllarSQLModule and setup required providers.
        """
        root_path = root_path or get_main_directory_by_stack("__main__", stack_level=2)
        if isinstance(migration_options, MigrationOption):
            migration_options = migration_options.dict()
        migration_options.setdefault("directory", "migrations")
        schema = SQLAlchemyConfig.model_validate(
            {
                "databases": databases,
                "engine_options": engine_options,
                "echo": echo,
                "models": models,
                "session_options": session_options,
                "migration_options": migration_options,
                "root_path": root_path,
            },
            from_attributes=True,
        )
        return cls.__setup_module(schema)

    @classmethod
    def __setup_module(cls, sql_alchemy_config: SQLAlchemyConfig) -> DynamicModule:
        db_service = EllarSQLService(
            databases=sql_alchemy_config.databases,
            common_engine_options=sql_alchemy_config.engine_options,
            common_session_options=sql_alchemy_config.session_options,
            echo=sql_alchemy_config.echo,
            models=sql_alchemy_config.models,
            root_path=sql_alchemy_config.root_path,
            migration_options=sql_alchemy_config.migration_options,
        )
        providers: t.List[t.Any] = []

        if db_service.has_async_engine_driver:
            providers.append(ProviderConfig(AsyncEngine, use_value=db_service.engine))
            providers.append(
                ProviderConfig(
                    AsyncSession,
                    use_value=lambda: db_service.session_factory(),
                    scope=request_or_transient_scope,
                )
            )
            providers.append(
                ProviderConfig(
                    Session,
                    use_value=_invalid_configuration(
                        f"{Session} is not configured based on your database options. Please use {AsyncSession}"
                    ),
                    scope=request_or_transient_scope,
                )
            )
            providers.append(
                ProviderConfig(
                    sa.Engine,
                    use_value=_invalid_configuration(
                        f"{sa.Engine} is not configured based on your database options. Please use {AsyncEngine}"
                    ),
                    scope=request_or_transient_scope,
                )
            )
        else:
            providers.append(ProviderConfig(sa.Engine, use_value=db_service.engine))
            providers.append(
                ProviderConfig(
                    Session,
                    use_value=lambda: db_service.session_factory(),
                    scope=request_or_transient_scope,
                )
            )
            providers.append(
                ProviderConfig(
                    AsyncSession,
                    use_value=_invalid_configuration(
                        f"{AsyncSession} is not configured based on your database options. Please use {Session}"
                    ),
                    scope=request_or_transient_scope,
                )
            )
            providers.append(
                ProviderConfig(
                    AsyncEngine,
                    use_value=_invalid_configuration(
                        f"{AsyncEngine} is not configured based on your database options. Please use {sa.Engine}"
                    ),
                    scope=request_or_transient_scope,
                )
            )

        providers.append(ProviderConfig(EllarSQLService, use_value=db_service))
        app_context_teardown.connect(
            functools.partial(cls._on_application_tear_down, db_service=db_service)
        )

        return DynamicModule(
            cls,
            providers=providers,
        )

    @classmethod
    def register_setup(cls, **override_config: t.Any) -> ModuleSetup:
        """
        Register Module to be configured through `SQLALCHEMY_CONFIG` variable in Application Config
        """
        root_path = get_main_directory_by_stack("__main__", stack_level=2)
        return ModuleSetup(
            cls,
            inject=[Config],
            factory=functools.partial(
                cls.__register_setup_factory,
                root_path=root_path,
                override_config=override_config,
            ),
        )

    @staticmethod
    def __register_setup_factory(
        module: t.Type["EllarSQLModule"],
        config: Config,
        root_path: str,
        override_config: t.Dict[str, t.Any],
    ) -> DynamicModule:
        if config.get("ELLAR_SQL") and isinstance(config.ELLAR_SQL, dict):
            defined_config = dict(config.ELLAR_SQL)
            defined_config.update(override_config)
            defined_config.setdefault("root_path", root_path)

            schema = SQLAlchemyConfig.model_validate(
                defined_config, from_attributes=True
            )

            schema.migration_options.directory = get_main_directory_by_stack(
                schema.migration_options.directory,
                stack_level=0,
                from_dir=defined_config["root_path"],
            )

            return module.__setup_module(schema)
        raise RuntimeError("Could not find `ELLAR_SQL` in application config.")