from string.templatelib import Template
from typing import Any, Dict, List, Optional, Type
from psycopg_pool import AsyncConnectionPool
from psycopg.rows import dict_row
from nc_db.metadata import nc_config
from nc_db.handlers.common import get_registered_signature as _get_registered_signature, table_exists as _table_exists, get_config_version as _get_config_version
from nc_db.identifiable import Identifiable
from nc_db.metadata.class_metadata import ClassMetadata


class QueryHandler():
    """
    Handles executing queries against the database
    """

    def __init__(self, connection_pool: AsyncConnectionPool):
        self._pool = connection_pool

    def _create_objects[T:Identifiable](self, cls: Type[T], metadata: ClassMetadata, rows: List[Dict[str, Any]] | None) -> List[T]:
        objects: List[T] = []
        if rows:
            for row in rows:
                # Create a blank dictionary and have each column mapper populate it
                obj_dict: Dict[str, Any] = {}
                for column in metadata.columns:
                    column.db_row_to_object(row, obj_dict)
                # Use pydantic's model_validate to create a class instance from the dictionary.
                objects.append(cls.model_validate(obj_dict))
        return objects

    async def table_exists(self, table: str) -> bool:
        async with self._pool.connection() as conn:
            async with conn.cursor() as cursor:
                return await _table_exists(cursor, table)

    async def get_db_config_version(self) -> int:
        async with self._pool.connection() as conn:
            async with conn.cursor() as cursor:
                return await _get_config_version(cursor, nc_config.VERSION_TABLE)

    async def get_registered_signature(self, metadata: ClassMetadata) -> str:
        async with self._pool.connection() as conn:
            async with conn.cursor() as cursor:
                return await _get_registered_signature(cursor, metadata.table_name, nc_config.CLASS_REGISTRATION_TABLE)

    async def get[T: Identifiable](self, cls: Type[T], metadata: ClassMetadata, ids: Optional[List[int]] = None) -> List[T]:
        async with self._pool.connection() as conn:
            async with conn.cursor(row_factory=dict_row) as cursor:
                # autopep8: off
                if ids:
                    query = t"SELECT * FROM {metadata.table_name:i} WHERE id=ANY({ids})"
                else:
                    query = t"SELECT * FROM {metadata.table_name:i}"
                # autopep8: on
                await cursor.execute(query)
                rows = await cursor.fetchall()
                result = self._create_objects(cls, metadata, rows)
                return result

    async def search[T:Identifiable](self, cls: Type[T], metadata: ClassMetadata, search_condition_template: Template) -> List[T]:
        async with self._pool.connection() as conn:
            async with conn.cursor(row_factory=dict_row) as cursor:
                # autopep8: off
                query = t"SELECT * FROM {metadata.table_name:i} WHERE "
                # autopep8: on
                query += search_condition_template
                await cursor.execute(query)
                rows = await cursor.fetchall()
                result = self._create_objects(cls, metadata, rows)
                return result

    async def select[T:Identifiable](self, cls: Type[T], metadata: ClassMetadata, select_sql: Template) -> List[T]:
        async with self._pool.connection() as conn:
            async with conn.cursor(row_factory=dict_row) as cursor:
                await cursor.execute(select_sql)
                rows = await cursor.fetchall()
                result = self._create_objects(cls, metadata, rows)
                return result
