from collections import defaultdict
from typing import Any, DefaultDict, Generator, Optional, Sequence

import psycopg as pg
import psycopg.sql as sql
from psycopg.rows import class_row

from chemrecon.schema.db_object import Column, DatabaseObject, Entry, Relation
import chemrecon.schema as schema


class QueryHandler():
    """ Context manager to wrap database queries.
    """
    conn: pg.Connection
    c: pg.Cursor

    # The handler also keeps track of a parallel 'database' of procedurally generated entries
    procedural_entries: dict[
        type[Entry], set[Entry]
    ]
    procedural_entries_index: dict[
        tuple[type[Entry], tuple],
        Entry
    ]
    procedural_entry_recon_id: dict[
        tuple[type[Entry], int],
        Entry
    ]
    procedural_id_table: dict[type[Entry], int]  # For assigning procedural ids to entries created locally


    # Context manager
    # ----------------------------------------------------------------------------------------------------------
    def __init__(self, connection: pg.Connection):
        self.conn = connection
        self.c = self.conn.cursor()

        self.procedural_entries = DefaultDict(set)
        self.procedural_entries_index = dict()
        self.procedural_entry_recon_id = dict()
        self.procedural_id_table = {
            e_type: -1 # Start at -1 and decrement
            for e_type in schema.entrytypes
        }

    def __enter__(self):
        # Enter, returning the handler as the interface
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # On leaving the context, execute the commands.
        if exc_type is not None:
            # Exception was raised in context
            raise exc_val
        self.conn.commit()
        self.c.close()

    # Procedural entries
    # ------------------------------------------------------------------------------------------------------------------
    def add_procedural_entry[T: Entry](self, entry: T) -> int:
        """ Adds a procedural entry, and returns its virtual recon id. If already exists,
            return the virtual recon id.
        """
        # TODO does the fact that index is not an ordereddict cause problems?
        index = tuple(val for val in entry.get_index_columns_with_values().values())
        try:
            # Try to return existing
            existing_entry = self.procedural_entries_index[
                type(entry), index
            ]
            entry.recon_id = existing_entry.recon_id
            return existing_entry.recon_id
        except KeyError:
            # Not found, create and insert
            new_id = self._assign_procedural_id(entry)
            self.procedural_entries[type(entry)].add(entry)
            self.procedural_entries_index[(type(entry), index)] = entry
            self.procedural_entry_recon_id[(type(entry), new_id)] = entry
            return new_id

    def _assign_procedural_id(self, entry: Entry) -> int:
        """ Assign a local id to an object, returning the id.
        """
        new_id = self.procedural_id_table[type(entry)]
        entry.recon_id = new_id
        self.procedural_id_table[type(entry)] -= 1
        return new_id

    # Row factories for Psycopg3
    # ------------------------------------------------------------------------------------------------------------------
    # Row factory for relation views
    def make_relation_entry_view_row_factory[T1: Entry, T2: Entry](
        self,
        relationtype: type[Relation[T1, T2]]
    ) -> pg.rows.RowFactory[tuple[Relation[T1, T2], T1, T2]]:

        class RelationIteratorRowFactory[T1_: Entry, T2_: Entry](pg.rows.RowFactory):
            """ Creates a row factory which constructs rows based on the relation view with attached entries.
            """

            def __init__(self, cursor: pg.Cursor[Any]):
                self.fields = [c.name for c in cursor.description]

                # Create fields dict (for fields of each table to the fields of the view)
                self.rel_fields: dict[str, str] = dict()
                self.t1_fields = dict()
                self.t2_fields = dict()

                for field in self.fields:
                    if field == 'recon_id_1':
                        self.rel_fields[field] = 'recon_id_1'
                        self.t1_fields[field] = 'recon_id'
                    elif field == 'recon_id_2':
                        self.rel_fields[field] = 'recon_id_2'
                        self.t2_fields[field] = 'recon_id'
                    elif field.startswith('rel_'):
                        self.rel_fields[field] = field.removeprefix('rel_')
                    elif field.startswith('t1_'):
                        self.t1_fields[field] = field.removeprefix('t1_')
                    elif field.startswith('t2_'):
                        self.t2_fields[field] = field.removeprefix('t2_')

            def __call__(self, values: Sequence[Any]) -> tuple[Relation[T1_, T2_], T1_, T2_]:
                # Get fields of each object
                rel_args: dict[str, Any] = dict()
                t1_args: dict[str, Any] = dict()
                t2_args: dict[str, Any] = dict()
                for fieldname, val in zip(self.fields, values):
                    # Assign values to fields
                    if fieldname in self.rel_fields:
                        rel_args[self.rel_fields[fieldname]] = val
                    if fieldname in self.t1_fields:
                        t1_args[self.t1_fields[fieldname]] = val
                    if fieldname in self.t2_fields:
                        t2_args[self.t2_fields[fieldname]] = val

                # Call constructors with the the given args
                return (
                    relationtype(**rel_args),
                    relationtype.source_entrytype(**t1_args),
                    relationtype.target_entrytype(**t2_args)
                )

        return RelationIteratorRowFactory

    # Getters
    # ----------------------------------------------------------------------------------------------------------
    def get_entry_by_recon_id[T: Entry](self, table: type[T], recon_id: int) -> T:
        """ Raises KeyError on failure. """
        res = self.get_entries_by_recon_ids(table, [recon_id])
        return res[recon_id]


    def get_entries_by_recon_ids[T: Entry](self, table: type[T], recon_ids: list[int]) -> dict[int, T]:
        """ Get the entries associated with the given recon ids. Raises KeyError on failure. """
        c = self.conn.cursor(row_factory = class_row(table))

        q = sql.SQL("""
                    SELECT * 
                    FROM {table} 
                    WHERE recon_id = %s;
        """).format(
            table = sql.Identifier(table.get_table_name())
        )
        c.executemany(
            q,
            params_seq = [[i] for i in recon_ids],
            returning = True
        )
        cursor_results = self._extract_results_from_cursor(c)
        return {
            i: e
            for i, e in zip(recon_ids, cursor_results)
        }


    def get_entry_by_index[T: Entry](self, entry_index: T) -> T:
        """ Given an entry with a valid index, get the corresponding full entry from the database.
            Raise KeyError in failure.
        """
        return self.get_entries_by_indices([entry_index])[0]


    def get_entries_by_indices[T: Entry](self, entry_indices: list[T]) -> list[T]:
        """ Given an entry with a valid index, get the corresponding full entry from the database.
            Raise KeyError in failure.
        """
        table = type(entry_indices[0])
        index_cols: list[Column] = table.get_index_columns()
        c = self.conn.cursor(row_factory = class_row(table))

        q = sql.SQL(""" 
            SELECT * 
            FROM {table}
            WHERE ({index_cols}) = ({index_placeholders})
        """).format(
            table = sql.Identifier(table.get_table_name()),
            index_cols = sql.SQL(', ').join(
                sql.Identifier(c.name) for c in index_cols
            ),
            index_placeholders = sql.SQL(', ').join(sql.SQL('%s') for _ in index_cols)
        )

        # Execute and fetch
        c.executemany(
            q,
            params_seq = [
                list(entry_index.get_index_columns_with_values().values())
                for entry_index in entry_indices
            ],
            returning = True
        )
        cursor_results = self._extract_results_from_cursor(c)
        return cursor_results


    # Relation getters
    # ----------------------------------------------------------------------------------------------------------
    def get_relations_by_recon_ids[T1: Entry, T2: Entry](
        self,
        entry_type: type[T1],
        recon_ids: list[int],
        relation_type: type[Relation[T1, T2]] | type[Relation[T2, T1]]
    ) -> list[list[Relation[T1,T2]]] | list[list[Relation[T2, T1]]]:
        """ Given a list of entries, get all connected relations of a given type.
        """
        # TODO dispatch to the three sub-functions.
        raise NotImplementedError()

    def get_relations_by_recon_ids_symmetric[T: Entry](
        self,
        recon_ids: list[int],
        relation_type: type[Relation[T, T]]
    ) -> list[list[Relation[T, T]]]:
        # Symmetric - either recon_id 1 or 2 can be found
        c = self.conn.cursor(row_factory = self.make_relation_entry_view_row_factory(relation_type))
        q = sql.SQL("""
            SELECT *
            FROM {table}
            WHERE (recon_id_1 = {recon_id_placeholder}) OR (recon_id_2 = {recon_id_placeholder});
        """).format(
            table = sql.Identifier(relation_type.get_table_name()),
            recon_id_placeholder = sql.SQL('%s')
        )
        c.executemany(q, params_seq = [[recon_id] * 2 for recon_id in recon_ids], returning = True)

        # Each input recon_id corresponds to a result set
        relations: list[list[Relation[T, T]]] = list()
        for _ in recon_ids:
            relations.append(c.fetchall())
            c.nextset()

        return relations

    def get_relations_by_recon_ids_of_t1[T1: Entry, T2: Entry](
        self,
        recon_ids: list[int],
        relation_type: type[Relation[T1, T2]]
    ) -> list[list[Relation[T1, T2]]]:
        c = self.conn.cursor(row_factory = class_row(relation_type))
        q = sql.SQL("""
                    SELECT *
                    FROM {table}
                    WHERE (recon_id_1 = {recon_id_placeholder});
                    """).format(
            table = sql.Identifier(relation_type.get_table_name()),
            recon_id_placeholder = sql.SQL('%s')
        )
        c.executemany(q, params_seq = [[recon_id] for recon_id in recon_ids], returning = True)

        # Each input recon_id corresponds to a result set
        relations: list[list[Relation[T1, T2]]] = list()
        for _ in recon_ids:
            relations.append(c.fetchall())
            c.nextset()

        return relations

    def get_relations_by_recon_ids_of_t2[T1: Entry, T2: Entry](
        self,
        recon_ids: list[int],
        relation_type: type[Relation[T2, T1]]
    ) -> list[list[Relation[T2, T1]]]:
        c = self.conn.cursor(row_factory = class_row(relation_type))
        q = sql.SQL("""
                    SELECT *
                    FROM {table}
                    WHERE (recon_id_2 = {recon_id_placeholder});
        """).format(
            table = sql.Identifier(relation_type.get_table_name()),
            recon_id_placeholder = sql.SQL('%s')
        )
        c.executemany(q, params_seq = [[recon_id] for recon_id in recon_ids], returning = True)

        # Each input recon_id corresponds to a result set
        relations: list[list[Relation[T1, T2]]] = list()
        for _ in recon_ids:
            relations.append(c.fetchall())
            c.nextset()

        return relations


    # Relation and entry getters
    # ----------------------------------------------------------------------------------------------------------
    def get_relations_with_entries_by_recon_ids[T1: Entry, T2: Entry](
        self,
        entry_type: type[T1],
        recon_ids: list[int],
        relation_type: type[Relation[T1, T2]] | type[Relation[T2, T1]],
    ) -> list[dict[T2, list[Relation[T1, T2]] | list[Relation[T2, T1]]]]:
        """ Given a list of entries, get connected relations of the given type.
            Returns a list, containing, for each input recon_id, a dictionary mapping connected vertices
            to the relation(s) with which they are connected.
        """
        # First, establish the types
        t2_entrytype = (
            relation_type.source_entrytype if relation_type.target_entrytype is entry_type
            else relation_type.target_entrytype
        )
        is_symmetric = (entry_type is t2_entrytype)

        # Dispatch to sub-functions depending on types
        if is_symmetric:
            # Case 1: Symmetric - either recon_id 1 or 2 can be found
            return self.get_relations_with_entries_by_recon_ids_symmetric(entry_type, recon_ids, relation_type)
        else:
            if relation_type.source_entrytype is entry_type:
                # Case 2: Relation is [T1, T2]
                return self.get_relations_with_entries_by_recon_ids_of_t1(entry_type, recon_ids, relation_type)
            elif relation_type.target_entrytype is entry_type:
                # Case 3: Relation is [T2, T1]
                return self.get_relations_with_entries_by_recon_ids_of_t2(entry_type, recon_ids, relation_type)
            else:
                # Invalid shape of entrytype and relation?
                raise ValueError()

    # Specialized getters for different entry/relation cases
    def get_relations_with_entries_by_recon_ids_symmetric[T: Entry](
        self,
        entry_type: type[T],
        recon_ids: list[int],
        relation_type: type[Relation[T, T]],
    ) -> list[dict[T, list[Relation[T, T]]]]:
        """ Specialization for symmetric relations.
        """
        c = self.conn.cursor(row_factory = class_row(relation_type))

        # TODO REDO, call above functions
        # TODO get by views!

        # Symmetric - either recon_id 1 or 2 can be found
        q = sql.SQL("""
                    SELECT *
                    FROM {table}
                    WHERE (recon_id_1 = {recon_id_placeholder}) OR (recon_id_2 = {recon_id_placeholder});
                    """).format(
            table = sql.Identifier(relation_type.get_table_name()),
            recon_id_placeholder = sql.SQL('%s')
        )
        c.executemany(q, params_seq = [[recon_id] * 2 for recon_id in recon_ids], returning = True)

        # Each input recon_id corresponds to a result set
        entry_relations: dict[int, list[Relation]] = dict()
        for recon_id in recon_ids:
            entry_relations[recon_id] = c.fetchall()
            c.nextset()

        # Get entries
        results: list[dict[T, list[Relation[T, T]]]] = list()

        for input_recon_id, rels in entry_relations.items():
            # For each source recon id, run a query to get relation targets
            subresults: dict[T, list[Relation[T, T]]] = defaultdict(list)
            target_recon_ids = [  # List corresponding to the list of relations rels
                rel.recon_id_1 if rel.recon_id_2 == input_recon_id else rel.recon_id_2
                for rel in rels
            ]
            q = sql.SQL("""
                        SELECT *
                        FROM {t2_table}
                        WHERE recon_id = %s;
                        """).format(
                t2_table = sql.Identifier(entry_type.get_table_name())
            )
            c = self.conn.cursor(row_factory = class_row(entry_type))
            c.executemany(q, params_seq = [[tid] for tid in target_recon_ids], returning = True)
            target_entries: list[T] = self._extract_results_from_cursor(c)

            # Add to results
            for rel, target_entry in zip(rels, target_entries):
                subresults[target_entry].append(rel)
            results.append(subresults)

        # Return
        return results

    def get_relations_with_entries_by_recon_ids_of_t1[T1: Entry, T2: Entry](
        self,
        entry_type: type[T1],
        recon_ids: list[int],
        relation_type: type[Relation[T1, T2]],
    ) -> list[dict[T2, list[Relation[T1, T2]]]]:
        """ Given a list of entries, get connected relations of the given type.
            Returns a list, containing, for each input recon_id, a dictionary mapping connected vertices
            to the relation(s) with which they are connected.
        """
        c = self.conn.cursor(row_factory = self.make_relation_entry_view_row_factory(relation_type))
        q = sql.SQL("""
                    SELECT *
                    FROM {view}
                    WHERE (recon_id_1 = {recon_id_placeholder});
                    """).format(
            view = sql.Identifier(f'{relation_type.get_table_name()}_v'),
            recon_id_placeholder = sql.SQL('%s')
        )
        c.executemany(q, params_seq = [[recon_id] for recon_id in recon_ids], returning = True)

        # Process output
        result: list[dict[T2, list[Relation[T1, T2]]]] = list()
        cursor_result = self._extract_results_from_cursor_by_sets(c)
        for recon_id, relation_list in zip(recon_ids, cursor_result):
            subdict: dict[T2, list[Relation[T1, T2]]] = defaultdict(list)
            for relation, entry_1, entry_2 in relation_list:
                assert entry_1.recon_id == recon_id
                subdict[entry_2].append(relation)
            result.append(subdict)
        return result

    def get_relations_with_entries_by_recon_ids_of_t2[T1: Entry, T2: Entry](
        self,
        entry_type: type[T1],
        recon_ids: list[int],
        relation_type: type[Relation[T2, T1]],
    ) -> list[dict[T2, list[Relation[T2, T1]]]]:
        """ Given a list of entries, get connected relations of the given type.
            Returns a list, containing, for each input recon_id, a dictionary mapping connected vertices
            to the relation(s) with which they are connected.
        """
        c = self.conn.cursor(row_factory = self.make_relation_entry_view_row_factory(relation_type))
        q = sql.SQL("""
                    SELECT *
                    FROM {view}
                    WHERE (recon_id_2 = {recon_id_placeholder});
                    """).format(
            view = sql.Identifier(f'{relation_type.get_table_name()}_v'),
            recon_id_placeholder = sql.SQL('%s')
        )
        c.executemany(q, params_seq = [[recon_id] for recon_id in recon_ids], returning = True)

        # Process output
        result: list[dict[T2, list[Relation[T2, T1]]]] = list()
        cursor_result = self._extract_results_from_cursor_by_sets(c)
        for recon_id, relation_list in zip(recon_ids, cursor_result):
            subdict: dict[T2, list[Relation[T2, T1]]] = defaultdict(list)
            for relation, entry_2, entry_1 in relation_list:
                assert entry_1.recon_id == recon_id
                subdict[entry_2].append(relation)
            result.append(subdict)
        return result

    # Graph exploration getters
    # ----------------------------------------------------------------------------------------------------------
    # These are the main queries used to efficiently explore the graph.
    def explore_vertex[T: Entry](
        self,
        entry: T,
        relation_types: type[Relation[T, Any]]
    ) -> list[tuple[Entry, list[Relation]]]:
        """ Explores the neighbouring entries to a given entry.
            Results are returned as a list of tuples, each containing a neighbouring entry and the relations
            which point to it.
            Note that the input entry may be given in the output in the case of reflexive relations.
        """

        # TODO Run get_relations_with_entries_from_source() for each allowed relation, and collate the results

        # TODO
        raise NotImplementedError()

    # Misc
    # ----------------------------------------------------------------------------------------------------------
    def get_number_of_entries[T: DatabaseObject](self, table: type[T]) -> int:
        """ Return hte number of entries in the table"""
        table_name = table.get_table_name()
        q = sql.SQL("""SELECT count(*) FROM {table}""").format(
            table = sql.Identifier(table_name)
        )
        self.c.execute(q)
        result: int = self.c.fetchone()[0]
        return result

    def iterate_table[T: DatabaseObject](
        self,
        table: type[T],
        batch_size: int = 1000,
        up_to: int = 0
    ) -> Generator[T, None, None]:
        """ Creates an iterator which queries all entries in a table.
            Entries are requested from the database in batches, but returned as a continuous iterator.
        """
        i: int = 0
        batch: list[T] = list()

        with self.conn.cursor(row_factory = class_row(table)) as cursor:
            q = sql.SQL("""
                SELECT *
                FROM {table}
                ;
            """).format(
                table = sql.Identifier(table.get_table_name())
            )
            try:
                cursor.execute(q)
            except Exception as e:
                raise RuntimeError('!: {e}')

            while True:
                # Main loop, fetch batches and yield from the batches
                batch = cursor.fetchmany(size = batch_size)
                if not batch:
                    break
                item: T
                for item in batch:
                    i += 1
                    yield item
                    if up_to and i > up_to:
                        return


    def iterate_relation_with_entries[T1: Entry, T2: Entry](
        self,
        relationtype: type[Relation[T1, T2]],
        batch_size: int = 1000,
        up_to: int = 0
    ) -> Generator[tuple[Relation[T1, T2], T1, T2], None, None]:
        """ Iterates the view over a relation table, getting entries adjacent to each relation.
        """
        i: int = 0
        batch: list = list()

        # TODO make a relationIteratorRowFactory and use

        # With the factory created, iterate the rows
        with self.conn.cursor(row_factory = self.make_relation_entry_view_row_factory(relationtype)) as cursor:
            q = sql.SQL("""
                        SELECT *
                        FROM {view}
                        ;
            """).format(
                view = sql.Identifier(f'{relationtype.get_table_name()}_v')
            )
            try:
                cursor.execute(q)
            except Exception as e:
                raise RuntimeError('!: {e}')

            while True:
                # Main loop, fetch batches and yield from the batches
                batch = cursor.fetchmany(size = batch_size)
                if not batch:
                    break
                item: tuple[Relation[T1, T2], T1, T2]
                for item in batch:
                    i += 1
                    yield item
                    if up_to and i > up_to:
                        return

    def _extract_results_from_cursor(self, cursor: Optional[pg.Cursor] = None) -> list[Any]:
        """ Iterates over all result sets of the Psycopg3 cursor and returns.
        """
        c = cursor or self.c
        result: list[Any] = [c.fetchone()]
        while c.nextset():
            result.append(c.fetchone())
        return result

    def _extract_results_from_cursor_by_sets(self, cursor: Optional[pg.Cursor] = None) -> list[list[Any]]:
        """ Iterates over all result sets of the Psycopg3 cursor and returns.
        """
        c = cursor or self.c
        result: list[list[Any]] = [c.fetchall()]
        while c.nextset():
            result.append(c.fetchall())
        return result
