import inspect
import logging
from typing import Dict, List, Literal, Mapping, Optional, Union, overload

from motor.motor_asyncio import AsyncIOMotorClient
from neo4j import AsyncDriver, AsyncGraphDatabase, Driver, GraphDatabase
from neo4j import AsyncSession as Neo4jAsyncSession
from neo4j import AsyncTransaction as Neo4jAsyncTransaction
from neo4j import Session as Neo4jSession
from pymongo import IndexModel, MongoClient
from qdrant_client import AsyncQdrantClient, QdrantClient
from qdrant_client.http import models as qdrant_models
from qdrant_client.http.exceptions import UnexpectedResponse
from qdrant_client.http.models import PayloadSchemaType
from redis import ConnectionPool
from redis import Redis as R
from redis.asyncio import ConnectionPool as AsyncConnectionPool
from redis.asyncio import Redis as AsyncRedis
from sqlalchemy import Engine
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
from sqlmodel import Session, create_engine
from sqlmodel.ext.asyncio.session import AsyncSession as SQLModelAsyncSession

from keble_db.wrapper import ExtendedAsyncRedis
from keble_db.wrapper.redis import ExtendedRedis

from .schemas import DbSettingsABC

logger = logging.getLogger(__name__)


class Db:
    def __init__(self, settings: DbSettingsABC):
        self.__settings = settings
        self.__redis_pool = (
            ConnectionPool.from_url(settings.redis_uri, decode_responses=True)
            if settings.redis_uri is not None
            else None
        )
        self.__aredis_pool = (
            AsyncConnectionPool.from_url(settings.redis_uri, decode_responses=True)
            if settings.redis_uri is not None
            else None
        )
        self.__sql_read_engine = (
            create_engine(
                str(
                    settings.sql_read_uri
                    if settings.sql_read_uri is not None
                    else settings.sql_uri
                ),
                pool_pre_ping=True,
                pool_size=1000,
                max_overflow=10,
                pool_recycle=3600,
                pool_use_lifo=True,
                echo=False,
            )
            if settings.sql_read_uri is not None or settings.sql_uri is not None
            else None
        )
        self.__sql_write_engine = (
            create_engine(
                str(
                    settings.sql_write_uri
                    if settings.sql_write_uri is not None
                    else settings.sql_uri
                ),
                pool_pre_ping=True,
                pool_size=1000,
                max_overflow=10,
                pool_recycle=3600,
                pool_use_lifo=True,
                echo=False,
            )
            if settings.sql_write_uri is not None or settings.sql_uri is not None
            else None
        )

        # New async engines
        self.__async_sql_read_engine = (
            create_async_engine(
                str(
                    settings.sql_read_uri
                    if settings.sql_read_uri is not None
                    else settings.sql_uri
                ),
                pool_pre_ping=True,
                pool_size=100,
                max_overflow=10,
                pool_recycle=3600,
                echo=False,
            )
            if settings.sql_read_uri is not None or settings.sql_uri is not None
            else None
        )
        self.__async_sql_write_engine = (
            create_async_engine(
                str(
                    settings.sql_write_uri
                    if settings.sql_write_uri is not None
                    else settings.sql_uri
                ),
                pool_pre_ping=True,
                pool_size=100,
                max_overflow=10,
                pool_recycle=3600,
                echo=False,
            )
            if settings.sql_write_uri is not None or settings.sql_uri is not None
            else None
        )

        self._sql_read_instance = None
        self._sql_write_instance = None
        self._async_sql_read_instance = None
        self._async_sql_write_instance = None
        self._redis_instance = None
        self._aredis_instance = None
        self._mongo_instance = None
        self._amongo_instance: Optional[AsyncIOMotorClient] = None
        self._aqdrant_instance: Optional[AsyncQdrantClient] = None
        self._qdrant_instance = None
        self._neo4j_driver: Optional[Driver] = None
        self._neo4j_async_driver: Optional[AsyncDriver] = None
        self._neo4j_session: Optional[Neo4jSession] = None
        self._neo4j_async_session: Optional[Neo4jAsyncSession] = None

    @property
    def sql_write_engine(self) -> Optional[Engine]:
        return self.__sql_write_engine

    @property
    def sql_read_engine(self) -> Optional[Engine]:
        return self.__sql_read_engine

    @property
    def async_sql_write_engine(self) -> Optional[AsyncEngine]:
        return self.__async_sql_write_engine

    @property
    def async_sql_read_engine(self) -> Optional[AsyncEngine]:
        return self.__async_sql_read_engine

    def get_redis(self, *, force_new_instance: bool = False) -> R:
        """Return conn to redis
        You do not need to explicitly close the redis conn"""
        if self.__redis_pool is None:
            raise ValueError("[Db] redis_pool required")
        if force_new_instance:
            return R(connection_pool=self.__redis_pool, decode_responses=True)
        if not force_new_instance and self._redis_instance is not None:
            return self._redis_instance
        self._redis_instance = R(
            connection_pool=self.__redis_pool, decode_responses=True
        )
        return self._redis_instance

    def get_aredis(self, *, new_instance: bool = True) -> AsyncRedis:
        if self.__aredis_pool is None:
            raise ValueError("[Db] aredis_pool required")
        if new_instance:
            return AsyncRedis(connection_pool=self.__aredis_pool, decode_responses=True)
        if not new_instance and self._aredis_instance is not None:
            return self._aredis_instance
        self._aredis_instance = AsyncRedis(
            connection_pool=self.__aredis_pool, decode_responses=True
        )
        return self._aredis_instance

    def get_extended_aredis(
        self, *, new_instance: bool = True, namespace: Optional[str] = None
    ) -> ExtendedAsyncRedis:
        return ExtendedAsyncRedis.extend(
            aredis=self.get_aredis(new_instance=new_instance), namespace=namespace
        )

    def get_extended_redis(
        self, *, force_new_instance: bool = False, namespace: Optional[str] = None
    ) -> ExtendedRedis:
        return ExtendedRedis.extend(
            redis=self.get_redis(force_new_instance=force_new_instance),
            namespace=namespace,
        )

    def get_mongo(self, *, new_instance: bool = True) -> MongoClient:
        """Return a mongo conn
        You do not need to close PyMongo connections. Leave them open so that PyMongo connection pooling gives you the most efficient performance"""
        if self.__settings.mongo_db_uri is None:
            raise ValueError("[Db] mongo_db_uri required")

        if new_instance:
            return MongoClient(
                self.__settings.mongo_db_uri, uuidRepresentation="standard"
            )

        if self._mongo_instance is not None:
            try:
                self._mongo_instance.server_info()
                return self._mongo_instance
            except Exception as e:
                logger.critical(
                    f"[Db] Old mongo instance has been closed, now create a new one and returned. Server info exception: {e}"
                )
                pass

        self._mongo_instance = MongoClient(
            self.__settings.mongo_db_uri, uuidRepresentation="standard"
        )
        return self._mongo_instance

    async def aget_amongo(self, *, new_instance: bool = True) -> AsyncIOMotorClient:
        """Return a async mongo conn
        You do not need to close PyMongo connections. Leave them open so that PyMongo connection pooling gives you the most efficient performance"""
        if self.__settings.mongo_db_uri is None:
            raise ValueError("[Db] mongo_db_uri required")

        if new_instance:
            # check new instance first, new instance is independent and will not save as self.instance
            return AsyncIOMotorClient(
                self.__settings.mongo_db_uri, uuidRepresentation="standard"
            )

        if self._amongo_instance is not None:
            try:
                await self._amongo_instance.server_info()
                return self._amongo_instance
            except Exception as e:
                logger.critical(
                    f"[Db] Old async mongo instance has been closed, now create a new one and returned. Server info exception: {e}"
                )
                pass

        self._amongo_instance = AsyncIOMotorClient(
            self.__settings.mongo_db_uri, uuidRepresentation="standard"
        )
        return self._amongo_instance

    def get_amongo(self) -> AsyncIOMotorClient:
        """Return a async mongo conn
        without checking old instance, use with caustious to prevent creating too many session/conn to mongo at the same time.
        """
        if self.__settings.mongo_db_uri is None:
            raise ValueError("[Db] mongo_db_uri required")

        return AsyncIOMotorClient(
            self.__settings.mongo_db_uri, uuidRepresentation="standard"
        )

    def get_qdrant_client(self, *, new_instance: bool = True) -> QdrantClient:
        if self.__settings.qdrant_host is None or self.__settings.qdrant_port is None:
            raise ValueError("[Db] qdrant_host and qdrant_port required")
        if new_instance:
            # check new instance first, new instance is independent and will not save as self.instance
            return QdrantClient(
                host=self.__settings.qdrant_host,
                port=self.__settings.qdrant_port,
                timeout=120,
            )
        if not new_instance and self._qdrant_instance is not None:
            return self._qdrant_instance
        self._qdrant_instance = QdrantClient(
            host=self.__settings.qdrant_host,
            port=self.__settings.qdrant_port,
            timeout=120,
        )
        return self._qdrant_instance

    def get_aqdrant_client(self, *, new_instance: bool = True) -> AsyncQdrantClient:
        if self.__settings.qdrant_host is None or self.__settings.qdrant_port is None:
            raise ValueError("[Db] qdrant_host and qdrant_port required")
        if new_instance:
            # check new instance first, new instance is independent and will not save as self.instance
            return AsyncQdrantClient(
                host=self.__settings.qdrant_host,
                port=self.__settings.qdrant_port,
                timeout=120,
            )
        if not new_instance and self._aqdrant_instance is not None:
            return self._aqdrant_instance
        self._aqdrant_instance = AsyncQdrantClient(
            host=self.__settings.qdrant_host,
            port=self.__settings.qdrant_port,
            timeout=120,
        )
        return self._aqdrant_instance

    def get_neo4j_driver(self, *, new_instance: bool = True) -> Driver:
        if self.__settings.neo4j_uri is None:
            raise ValueError("[Db] neo4j_uri required")
        if self._neo4j_driver is None or new_instance:
            auth = None
            if (
                self.__settings.neo4j_user is not None
                and self.__settings.neo4j_password is not None
            ):
                auth = (self.__settings.neo4j_user, self.__settings.neo4j_password)
            self._neo4j_driver = GraphDatabase.driver(
                self.__settings.neo4j_uri,
                auth=auth,
            )
        return self._neo4j_driver

    def get_aneo4j_driver(self, *, new_instance: bool = True) -> AsyncDriver:
        if self.__settings.neo4j_uri is None:
            raise ValueError("[Db] neo4j_uri required")
        if self._neo4j_async_driver is None or new_instance:
            auth = None
            if (
                self.__settings.neo4j_user is not None
                and self.__settings.neo4j_password is not None
            ):
                auth = (self.__settings.neo4j_user, self.__settings.neo4j_password)
            self._neo4j_async_driver = AsyncGraphDatabase.driver(
                self.__settings.neo4j_uri,
                auth=auth,
            )
        return self._neo4j_async_driver

    def get_async_neo4j_driver(self, *, new_instance: bool = True) -> AsyncDriver:
        return self.get_aneo4j_driver(new_instance=new_instance)

    def get_neo4j_session(
        self,
        *,
        new_instance: bool = True,
        database: Optional[str] = None,
    ) -> Neo4jSession:
        driver = self.get_neo4j_driver(new_instance=new_instance)
        target_db = database if database is not None else self.__settings.neo4j_database
        if new_instance or self._neo4j_session is None:
            self._neo4j_session = driver.session(database=target_db)
        return self._neo4j_session

    @overload
    async def get_aneo4j(
        self,
        *,
        new_instance: bool = True,
        database: Optional[str] = None,
        as_transaction: Literal[True] = True,
    ) -> Neo4jAsyncTransaction: ...

    @overload
    async def get_aneo4j(
        self,
        *,
        new_instance: bool = True,
        database: Optional[str] = None,
        as_transaction: Literal[False],
    ) -> Neo4jAsyncSession: ...

    async def get_aneo4j(
        self,
        *,
        new_instance: bool = True,
        database: Optional[str] = None,
        as_transaction: bool = True,
    ) -> Neo4jAsyncTransaction | Neo4jAsyncSession:
        session = self.get_async_neo4j_session(
            new_instance=new_instance,
            database=database,
        )
        if not as_transaction:
            return session
        return await session.begin_transaction()

    def get_async_neo4j_session(
        self,
        *,
        new_instance: bool = True,
        database: Optional[str] = None,
    ) -> Neo4jAsyncSession:
        driver = self.get_aneo4j_driver(new_instance=new_instance)
        target_db = database if database is not None else self.__settings.neo4j_database
        if new_instance or self._neo4j_async_session is None:
            self._neo4j_async_session = driver.session(database=target_db)
        return self._neo4j_async_session

    def get_sql_write_client(self, *, new_instance: bool = True) -> Session:
        if self.sql_write_engine is None:
            raise ValueError("[Db] sql_write_engine required")
        if new_instance:
            return Session(self.sql_write_engine)
        if not new_instance and self._sql_write_instance is not None:
            return self._sql_write_instance
        self._sql_write_instance = Session(self.sql_write_engine)
        return self._sql_write_instance

    def get_sql_read_client(self, *, new_instance: bool = True) -> Session:
        if self.sql_read_engine is None:
            raise ValueError("[Db] sql_read_engine required")
        if new_instance:
            return Session(self.sql_read_engine)
        if not new_instance and self._sql_read_instance is not None:
            return self._sql_read_instance
        self._sql_read_instance = Session(self.sql_read_engine)
        return self._sql_read_instance

    # def get_async_sql_write_client(
    #     self, *, new_instance: bool = True
    # ) -> SQLModelAsyncSession:
    #     """Get an async SQL session for writing operations."""
    #     if self.async_sql_write_engine is None:
    #         raise ValueError("[Db] async_sql_write_engine required")
    #     if new_instance:
    #         return SQLModelAsyncSession(self.async_sql_write_engine, expire_on_commit=False)
    #     if not new_instance and self._async_sql_write_instance is not None:
    #         return self._async_sql_write_instance
    #     self._async_sql_write_instance = SQLModelAsyncSession(
    #         self.async_sql_write_engine,
    #         expire_on_commit=False  # Prevent greenlet_spawn error on attribute access after commit
    #     )
    #     return self._async_sql_write_instance

    # def get_async_sql_read_client(
    #     self, *, new_instance: bool = True
    # ) -> SQLModelAsyncSession:
    #     """Get an async SQL session for reading operations."""
    #     if self.async_sql_read_engine is None:
    #         raise ValueError("[Db] async_sql_read_engine required")
    #     if new_instance:
    #         return SQLModelAsyncSession(self.async_sql_read_engine, expire_on_commit=False)
    #     if not new_instance and self._async_sql_read_instance is not None:
    #         return self._async_sql_read_instance
    #     self._async_sql_read_instance = SQLModelAsyncSession(
    #         self.async_sql_read_engine,
    #         expire_on_commit=False  # Prevent greenlet_spawn error on attribute access after commit
    #     )
    #     return self._async_sql_read_instance

    @staticmethod
    def ready_mongo(client: MongoClient):
        # check by list
        client.list_databases()

    @staticmethod
    def ready_mongo_indexes(
        client: MongoClient,
        db_name: str,
        collection_name: str,
        indexes: List[Dict[str, Union[int, str]]],
    ) -> None:
        """
        Ensure MongoDB indexes exist. Creates them if not.

        :param client: MongoClient instance
        :param db_name: Name of the MongoDB database
        :param collection_name: Name of the MongoDB collection
        :param indexes: List of dictionaries defining the index fields and their types.
                        - For normal ascending/descending indexes, use 1 (ascending) or -1 (descending).
                        - For text search indexes, use "text".
                        Example: [{"field1": 1, "field2": -1}, {"field3": "text"}]
        """
        collection = client[db_name][collection_name]
        existing_indexes = collection.index_information()

        # Create index models based on the input
        indexes_need_to_create = []
        for index in indexes:
            compound_index_key = [(field, order) for field, order in index.items()]

            # Check if the index already exists in the collection
            if not any(
                compound_index_key == idx.get("key")
                for idx in existing_indexes.values()
            ):
                index_model = IndexModel(compound_index_key)
                indexes_need_to_create.append(index_model)
        if len(indexes_need_to_create) > 0:
            collection.create_indexes(indexes_need_to_create)

    @staticmethod
    def ready_redis(redis: R):
        # check by get key
        redis.get("k")

    @staticmethod
    def ready_qdrant(
        client: QdrantClient,
        collection: str,
        vectors_config: Union[
            qdrant_models.VectorParams, Mapping[str, qdrant_models.VectorParams]
        ],
    ):
        # need to check database exist or not
        try:
            client.get_collection(collection_name=collection)
        except (AssertionError, UnexpectedResponse):
            # collection not found, or connection error. Try to create
            client.create_collection(
                collection_name=collection,
                vectors_config=vectors_config,
            )

    @staticmethod
    def ready_neo4j(driver: Driver):
        # Verify connectivity; raises if misconfigured
        driver.verify_connectivity()

    @staticmethod
    def ready_qdrant_indexes(
        client: QdrantClient,
        collection_name: str,
        indexes: List[Dict[str, PayloadSchemaType]],
    ) -> None:
        """
        Ensure Qdrant payload indexes exist. Creates them if not.

        :param client: QdrantClient instance
        :param collection_name: Name of the Qdrant collection
        :param indexes: List of dictionaries where keys are field names and values are their corresponding types (PayloadSchemaType).
                        Example: [{"field1": PayloadSchemaType.INTEGER}, {"field2": PayloadSchemaType.KEYWORD}]
        """
        try:
            collection_info = client.get_collection(collection_name)
            existing_indexes = collection_info.payload_schema or {}

            for index in indexes:
                for field_name, schema_type in index.items():
                    if field_name not in existing_indexes:
                        client.create_payload_index(
                            collection_name=collection_name,
                            field_name=field_name,
                            field_schema=qdrant_models.PayloadSchemaType(schema_type),
                        )
        except (AssertionError, UnexpectedResponse):
            logger.critical(
                f"[Db] Collection {collection_name} does not exist or failed to retrieve collection info."
            )

    @staticmethod
    def try_close(
        *sessions: Optional[
            MongoClient
            | R
            | Session
            | QdrantClient
            | AsyncIOMotorClient
            | AsyncQdrantClient
            | AsyncRedis
            | SQLModelAsyncSession
            | AsyncEngine
            | Driver
            | AsyncDriver
            | Neo4jSession
            | Neo4jAsyncSession
            | Neo4jAsyncTransaction
        ],
    ):
        import asyncio

        for session in sessions:
            if session is None:
                continue

            try:
                # Call close() method
                close_method = getattr(session, "close", None)
                if close_method is None:
                    continue

                # Check if it's an async method
                if inspect.iscoroutinefunction(close_method):
                    # Create a coroutine
                    coro = close_method()

                    # Try to run it or create a task based on context
                    try:
                        # Check if we're in an event loop
                        loop = asyncio.get_event_loop()
                        if loop.is_running():
                            # Create a task - we can't await here
                            # Type ignores for linter errors - these are valid operations
                            asyncio.create_task(coro)  # type: ignore
                        else:
                            # No running loop, create one and run
                            asyncio.run(coro)  # type: ignore
                    except RuntimeError:
                        # No event loop exists, create one
                        try:
                            asyncio.run(coro)  # type: ignore
                        except Exception as e:
                            logger.error(f"[Db] Failed to run async close: {e}")
                else:
                    # Regular synchronous close
                    result = close_method()

                    # Check if the result is awaitable
                    if inspect.isawaitable(result):
                        try:
                            loop = asyncio.get_event_loop()
                            if loop.is_running():
                                asyncio.create_task(result)  # type: ignore
                            else:
                                asyncio.run(result)  # type: ignore
                        except RuntimeError:
                            try:
                                asyncio.run(result)  # type: ignore
                            except Exception as e:
                                logger.error(f"[Db] Failed to run awaitable: {e}")

            except Exception as e:
                # Log errors but don't raise
                logger.error(f"[Db] Failed to close a db session due to {e}")

    @staticmethod
    async def try_close_async(
        *sessions: Optional[
            MongoClient
            | R
            | Session
            | QdrantClient
            | AsyncIOMotorClient
            | AsyncQdrantClient
            | AsyncRedis
            | SQLModelAsyncSession
            | AsyncEngine
            | Driver
            | AsyncDriver
            | Neo4jSession
            | Neo4jAsyncSession
            | Neo4jAsyncTransaction
        ],
    ):
        """Async version of try_close that can properly await async close methods."""
        for session in sessions:
            if session is None:
                continue

            try:
                close_method = getattr(session, "close", None)
                if close_method is None:
                    continue

                if inspect.iscoroutinefunction(close_method):
                    await close_method()
                else:
                    result = close_method()
                    if inspect.isawaitable(result):
                        await result

            except Exception as e:
                logger.error(f"[Db] Failed to close async session due to {e}")
