import asyncio
from typing import Any, Callable, Dict, List, Optional, Type, overload
from psycopg import AsyncConnection
from psycopg_pool import AsyncConnectionPool

from nc_db.db_settings import DbSettings
from nc_db.metadata.class_metadata import ClassMetadata
from nc_db.handlers.query_handler import QueryHandler
from nc_db.handlers.command_handler import CommandHandler
from nc_db.identifiable import NEW_ID, T, Identifiable
from string.templatelib import Template

from nc_db.metadata.nc_config import CLASS_VERSION_DICT, ALLOW_SCHEMA_DRIFT


class Db:
    """
    Provides a document-oriented DB system using Postgres with b-tree indexes for better performance.
    Think mongodb for Postgres with some help from Pydantic.
    New collections (tables) are created automatically but model and schema migrations are not (use the migrate function or roll your own).
    Collections are defined by a class's signature. All objects in a collection are guaranteed to have the same shape (serialize and deserialize).
    A class's signature is unique based upon the structure of the class and the indexes it defines.
    To improve performance with large objects:
        1. Each object type has a dedicated table in the database.
        1. Class attributes annotated with DbIndex are automatically maintained in dedicated columns (B-trees, no GIN indexes)
        1. Attributes you want to search on MUST be annotated with DbIndex
            1. You can use the DbIndex type aliases for common use cases
            Examples:
                - my_class_var: Annotated[int,DbIndex('INT',unique=True, nullable=False)]
                - my_class_var: DbIntIndex

    The Db instance should share a lifecycle with your service or application. Recommended practice is to treat it like
    a singleton for each database your service connects to. It is asynchronous and uses a connection pool under the hood.
    """

    def __init__(self, config: DbSettings | AsyncConnectionPool):
        """
        Creates a Db instance that can be used for CRUD against a Postgres database.

        :param config: Either the DbSettings or a connection pool. If DbSettings, the Db class will create the connection pool.
        """
        self._connected = False
        if isinstance(config, DbSettings):
            self._pool = Db.create_connection_pool(config)
        else:
            self._pool = config
        self._class_metadata_dict: Dict[Type[Identifiable], ClassMetadata] = {}

    async def connect(self):
        """
        Connects the Db class to the database and ensures the DB is configured correctly.
        """
        if not self._connected:
            await self._pool.open()
            self._query_handler = QueryHandler(self._pool)
            self._command_handler = CommandHandler(self._pool)

            # 1. Ensure the config tables exist
            await self._command_handler.ensure_config_tables()

            self._connected = True

    async def disconnect(self):
        """
        Closes all connections to the DB
        """
        await self._pool.close()
        self._connected = False

    @classmethod
    def create_connection_pool(cls, settings: DbSettings):
        """
        A utility function that creates a connection pool for the database specified
        by the settings parameter. This is called automatically by the Db.connect() function

        settings: Specifies the information needed to create the connection pool to a specific database
        """

        #
        async def configure_connection(conn: AsyncConnection[Any]):
            # This is the only place the search path should be set.
            # We set it here to configure every connection up-front.
            await CommandHandler.set_search_path(settings.db_schema, conn)

        # Build DSN string
        dsn = (
            f"postgresql://{settings.db_user}:{settings.db_password}"
            f"@{settings.db_host}:{settings.db_port}/{settings.db_name}"
        )
        pool: AsyncConnectionPool = AsyncConnectionPool(
            conninfo=dsn,
            min_size=settings.db_pool_min,
            max_size=settings.db_pool_max,
            open=False,
            configure=configure_connection
        )
        return pool

    def _ensure_connection(self, raise_if_false: bool = True) -> bool:
        if self._connected:
            return True
        elif raise_if_false:
            raise RuntimeError("Not connected to the DB. You must await Db.connect() prior to making calls that interact with the DB")
        return False

    async def _get_class_metadata(self, cls: Type[T]) -> ClassMetadata:
        """
        Gets / creates metadata for cls and ensures the appropriate database tables exist for this class.
        Returns the metadata object once the DB is configured to support CRUD on the class represented by cls.

        :param cls: The class to be verified
        Returns metadata for the class represented by cls.
        """
        # Make sure we're connected to the database
        self._ensure_connection()

        version = CLASS_VERSION_DICT.get(cls)
        if version is None:
            raise RuntimeError(
                f"Class {cls.__name__} does not have the @schema_version decorator. You must specify a schema version using the @schema_version class decorator")

        # If we have DbClassMetadata, return it, otherwise
        # get / create it
        metadata = self._class_metadata_dict.get(cls, None)
        if metadata is None:
            # NOTE: We don't need to have a critical section here as the ClassMetadata constructor is idempotent
            # and the command_handler.ensure_class_table function uses an advisory lock to create the table if
            # it doesn't exist.
            # The metadata_dict is effectively just a cache to avoid reprocessing the cls type multiple times.

            # NOTE: The (very minor) downside to not using a critical section here is that, if multiple calls are made simultaneously
            # the very first time we see a class, we will check the DB for each

            # We don't need to check for the key as the Identifiable class throws an error at import if derived classes
            # haven't specified a schema_version with the @schema_version decorator.
            version = CLASS_VERSION_DICT[cls]
            metadata = ClassMetadata(cls, version)
            await self._command_handler.ensure_class_table(metadata)
            if cls not in ALLOW_SCHEMA_DRIFT:
                registered_signature = await self._query_handler.get_registered_signature(metadata)
                if registered_signature != metadata.signature:
                    # autopep8: off
                    msg = f"""Run-time class {metadata.class_name} version {metadata.schema_version} has a signature ({metadata.signature}) that doesn't match the registered signature in the database ({registered_signature}) """
                    # autopep8: on
                    raise ValueError(msg)
            self._class_metadata_dict.update({cls: metadata})
        return metadata

    async def insert(self, item: Identifiable) -> int:
        """
        Inserts a new instance of an object into the database.
        The object must meet the following criteria:
            1. id must be equal to NEW_ID (call update if this is an update to an existing object)
            1. The ClassConfig for this object type (class) must have already been registered once (and only once).
        NOTE: The object's id will be set to match the new id provided by the database.


        :param object: A new object to be stored in the database
        :return: The new id for the object. NOTE: After calling this function, the obect's id will be set to match this value.
        """
        if item.id != NEW_ID:
            raise ValueError("When calling insert, the object's id must be NEW_ID. Call update to update an existing object.")
        metadata = await self._get_class_metadata(item.__class__)

        result = await self._command_handler.insert(metadata, item)
        return result

    @overload
    async def search(self, cls: Type[T], search_condition: str, value: Any,
                     ) -> List[T]: ...

    @overload
    async def search(self, cls: Type[T], search_condition: Template) -> List[T]: ...

    async def search(self,
                     cls: Type[T],
                     search_condition: str | Template,
                     value: Optional[Any] = None) -> List[T]:
        """
        Description
        -----------
        Searches for objects in the database. All attributes used in the search must be declared of type DbIndex[T]

        Examples
        --------
        results = await db.search(MyClass, "age", 11)

        results = await db.search(MyClass, t"name={name} AND age={age}")

        Overloads
        ---------
        search (cls, search_condition, value)
            A basic search for all objects where the attribute specified by search_condiition matches value

        search (cls, search_condition)
            Search using a string template for more complex queries.
            The search_condition must be a valid SQL search_condition.
            DO NOT include 'WHERE' in the string template.
            NOTE: If you use string template interpolation on postgres identifiers (e.g. column names)
                you must use the 'i' format specifier. Example:  t"{column_name:i} = {value}"

        Parameters
        ----------
        search_condition: Either a SQL criteria template or an attribute name
        value: Optional -- only used when the search_condition is an attribute name.


        """
        metadata = await self._get_class_metadata(cls)
        if isinstance(search_condition, str):
            # autopep8: off
            search_template = t"{search_condition:i} = {value}"
            # autopep8: on
        else:
            search_template = search_condition
        result = await self._query_handler.search(cls, metadata, search_template)
        return result

    async def select(self, cls: Type[T], select_sql: Template) -> List[T]:
        metadata = await self._get_class_metadata(cls)
        result = await self._query_handler.select(cls, metadata, select_sql)
        return result

    async def update(self, item: Identifiable):
        """
        Completely replaces an existing object in the database with item.

        :param item: The object to be stored in the database. This will completely replace any object in the database with the same id, class, and _version_schema
        :return: A list of ids that were updated (in the order returned by the Db)
        """
        metadata = await self._get_class_metadata(item.__class__)
        await self._command_handler.update(metadata, item)

    async def get(self, cls: Type[T], ids: Optional[list[int]] = None) -> list[T]:
        """
        Returns all of the objects in the database that have an id in the ids list.
        If ids is None, all objects of the given type (and _version_schema) are returned.
        """
        metadata = await self._get_class_metadata(cls)
        return await self._query_handler.get(cls, metadata, ids)

    async def delete(self, cls: Type[T], ids: list[int]) -> list[int]:
        metadata = await self._get_class_metadata(cls)
        results = await self._command_handler.delete(metadata, ids)
        return results

    async def migrate[T: (Identifiable), S:(Identifiable)](
            self,
            old_cls: Type[T],
            migrator: Callable[[List[T]], List[S]],
            delete_after_migrate: bool = False) -> List[int]:
        """
        Uses the migrator function parameter to convert all instances of old_cls to new_cls,
        saves these to the appropriate table in the database and deletes the old instances (if delete_after_migrate is true)
        """
        # Since this function just leverages other db functions, we don't need to ensure
        # the connection of class registrations -- these are done in the other db functions.
        old_items = await self.get(old_cls, [])
        new_items = migrator(old_items)
        insert_tasks = [self.insert(item) for item in new_items]
        await asyncio.gather(*insert_tasks)
        if delete_after_migrate:
            deleted_ids = await self.delete(old_cls, [item.id for item in old_items])
            # TODO: Delete the old table and class registration
            return deleted_ids
        return []

    async def get_metadata(self, cls: Type[T]):
        metadata = await self._get_class_metadata(cls)
        return metadata
