""" Method for creating an explored Entry Graph based on the database.
"""
from typing import Optional

import psycopg as pg

import chemrecon.connection as connection

from chemrecon.schema import Entry, Relation, ComposedRelation, ProceduralRelation, ProceduralGeneratorError

from chemrecon.entrygraph.entrygraph import EntryGraph, ReconID, VertexIndex
from chemrecon.entrygraph.explore_procedure import (
    ExploreProcedure, ExploreProcedureT1, ExploreProcedureT2, ExploreProcedureSym
)
from chemrecon.entrygraph.explorationprotocol import ExplorationProtocol


def explore(
    entrygraph: EntryGraph,
    protocol: ExplorationProtocol,
    steps: int = 4
):
    """ Expand the given entry graph by traversing the database network using the specified protocol for a given
        number of steps.
    """
    Explorer(entrygraph, protocol, steps).explore()


class Explorer:
    entrygraph: EntryGraph
    generations: Optional[int] = None
    protocol: ExplorationProtocol

    def __init__(
        self,
        entrygraph: EntryGraph,
        protocol: ExplorationProtocol,
        generations: int = None
    ):
        self.entrygraph = entrygraph
        self.generations = generations
        self.protocol = protocol

        # Generate cursors for each procedure
        for proc in {
            *protocol.explore_procedures,
            *protocol.post_explore_procedures,
            *[p for p, _ in protocol.transitive_subprocedures.values()],
            *[p for _, p in protocol.transitive_subprocedures.values()]
        }:
            proc: ExploreProcedure
            proc.cursor = pg.Cursor(
                connection = connection.handler.conn,
                row_factory = connection.handler.make_relation_entry_view_row_factory(proc.relation_type)
            )

    # Main exploration alg.
    # --------------------------------------------------------------------------------------------------------------
    def explore(self) -> None:
        """ Explores the graph based on the database connection.
            After execution, the graph will be fully populated.
        """
        import chemrecon.core.populate_query_handler

        # Sets of already explored vertices (updated at end of every iteration)
        vs_new: set[VertexIndex] = set()
        vs_explored: set[VertexIndex] = set()

        # Add initial vertices
        for v_idx in self.entrygraph.initial_vertices:
            vs_new.add(v_idx)

        # Iteratively explore
        i: int = 0  # Generation
        while True:
            # Main iteration
            i += 1

            # Compute new set of vertices to explore
            to_explore: dict[type[Entry], set[tuple[VertexIndex, ReconID]]] = {
                etype: set() for etype in self.protocol.entry_types
            }
            for v_idx in vs_new:
                v = self.entrygraph.get_vertex_by_vertex_index(v_idx)
                if v is None:
                    assert False, 'unreachable'
                to_explore[type(v.entry)].add((v_idx, v.recon_id))

            vs_explored.update(vs_new)

            # For each type of relation, call the explore procedure, and sort out entries according to the filter
            for explr_proc in self.protocol.explore_procedures:

                # If past generation limit, explore only relations which explicitly ignore this
                if i > self.generations:
                    if not explr_proc.relation_type.ignore_generation_limit:
                        continue

                # Run procedure, producing relations and entry endpoints
                input_entries = list(to_explore[explr_proc.takes_entrytype])
                if len(input_entries) > 0:
                    try:
                        res = self._run_procedure(
                            explr_proc,
                            [self.entrygraph.g.get_node_data(v).entry for v, _ in input_entries]
                        )
                        if res is None:
                            # Should not happen (?)
                            continue
                    except ProceduralGeneratorError:
                        # Could not generate, so do not add
                        continue

                    # res: list of (t_takes, relation[t_takes, t_gives], t_gives)
                    for (takes_vertex_index, takes_recon_id), subresult in zip(input_entries, res):
                        # TODO error here, sym relations continually gives the e_takes instead of the e_gives
                        for rel, e_gives in subresult:

                            # Filter
                            if explr_proc.relation_filter is not None:
                                if not explr_proc.relation_filter(rel):
                                    continue
                            if explr_proc.entry_filter is not None:
                                if not explr_proc.entry_filter(e_gives):
                                    continue

                            # Add relation to entrygraph
                            v_new_idx = self.entrygraph.add_vertex_from(
                                from_index = takes_vertex_index,
                                relation = rel,
                                entry = e_gives,
                                generation = i
                            )
                            vs_new.add(v_new_idx)  # Add new vertices to explore

            # End of iteration, update lists
            vs_new = vs_new - vs_explored
            if len(vs_new) == 0:
                break

        # Explore terminal relation (does not lead to any new vertices to add)
        # TODO should not add new vertices, only relations between vertices already in the graph
        # TODO !important
        for explr_proc in self.protocol.post_explore_procedures:
            input_entries: list[tuple[VertexIndex, ReconID]] = [
                (v.vertex_index, v.recon_id)
                for v in self.entrygraph.g.nodes()
                if type(v.entry) is explr_proc.takes_entrytype
            ]
            if len(input_entries) > 0:
                res = self._run_procedure(
                    explr_proc,
                    [self.entrygraph.g.get_node_data(v).entry for v, _ in input_entries]
                )
                for (takes_vertex_index, takes_recon_id), subresult in zip(input_entries, res):
                    for rel, e_gives in subresult:
                        # Add relation to entrygraph only if exists
                        lookup = self.entrygraph.get_vertex_by_entry(e_gives)
                        if lookup is not None:
                            self.entrygraph.add_edge(
                                source_v_index = takes_vertex_index,
                                target_v_index = lookup,
                                relation = rel
                            )

                        # TODO add setting to create new entries from terminal relations?
                        # Old, added new entries unnecessarily
                        # v_new_idx = self.entrygraph.add_vertex_from(
                        #     from_index = takes_vertex_index,
                        #     relation = rel,
                        #     entry = e_gives,
                        #     generation = -1
                        # )

        # Commit changes made if handler can populate the backend
        if isinstance(connection.handler, chemrecon.core.populate_query_handler.PopulateQueryHandler):
            connection.handler.conn.commit()

        # Done, finalise
        pass


    # Dispatchers for individual procedures
    # ---------------------------------------------------------------------------------------------------------------
    def _run_procedure[T_takes: Entry, T_gives: Entry](
        self,
        proc: ExploreProcedure[T_takes, T_gives],
        take_entries: list[T_takes]
    ) -> list[list[tuple[Relation[T_takes, T_gives], T_gives]]]:
        if len(take_entries) == 0:
            return []
        # Dispatch depending on whether procedural or database-backed relation
        if issubclass(proc.relation_type, ComposedRelation):
            return self._run_procedure_transitive(proc, take_entries)
        elif issubclass(proc.relation_type, ProceduralRelation):
            return self._run_procedure_procedural(proc, take_entries)
        else:
            return self._run_procedure_database(proc, take_entries)


    def _run_procedure_database[T_takes: Entry, T_gives: Entry](
        self,
        proc: ExploreProcedure[T_takes, T_gives],
        take_entries: list[T_takes],
    ) -> list[list[tuple[Relation[T_takes, T_gives], T_gives]]]:
        # The result of the procedure relies only on information in the database
        proc.cursor.executemany(
            query = proc.q,
            params_seq = [[e.recon_id] * proc.n_params for e in take_entries],
            returning = True
        )

        # Fetch results and return
        result: list[list[tuple[Relation, Entry, Entry]]] = [proc.cursor.fetchall()]
        while proc.cursor.nextset():
            result.append(proc.cursor.fetchall())

        # How to interpret and add depends on T1, T2, or sym relation
        match proc:
            case ExploreProcedureT1():
                # Result is relation and T2
                return [
                    [(r, e2) for r, e1, e2 in subresult]
                    for subresult in result
                ]
            case ExploreProcedureT2():
                # Result is relation and T1
                return [
                    [(r, e1) for r, e1, e2 in subresult]
                    for subresult in result
                ]
            case ExploreProcedureSym():
                # Result entry is T1 or T2, depending on which is different from the input
                return [
                    [
                        (r, e2) if e1.recon_id == takes_recon_id else (r, e1)
                        for r, e1, e2 in subresult
                    ]
                    for takes_recon_id, subresult in zip((
                        e.recon_id for e in take_entries
                    ), result)
                ]
            case _:
                assert False, 'unreachable'


    def _run_procedure_procedural[T_takes: Entry, T_gives: Entry](
        self,
        proc: ExploreProcedure[T_takes, T_gives],
        take_entries: list[T_takes],
    ) -> list[list[tuple[Relation[T_takes, T_gives], T_gives]]]:
        import chemrecon.core.populate_query_handler

        # TODO handle catching the ProceduralGeneratorError()!

        # Run the generate() method of the relation, and possibly update the DB
        # if the handler supports insertion
        assert issubclass(proc.relation_type, ProceduralRelation)

        # First, check that the relation already exists in the DB.
        db_result = self._run_procedure_database(proc, take_entries)

        # If we found results for all queries
        if all(len(subresult) > 0 for subresult in db_result):
            return db_result

        # If not , use generator
        procedural_result: list[list[tuple[Relation[T_takes, T_gives], T_gives]]] = list()

        for i, take_entry in enumerate(take_entries):
            if len(db_result[i]) > 0:
                procedural_result.append(db_result[i])
                continue
            res = proc.relation_type.generate(
                take_entry
            )
            procedural_result.append([
                (rel, e_gives)
                for rel, e_gives in res
            ])

        # Update the database based on the db_result if possible
        if isinstance(connection.handler, chemrecon.core.populate_query_handler.PopulateQueryHandler):
            for take_reconid, subresult in zip((e.recon_id for e in take_entries), procedural_result):
                assert take_reconid is not None
                if len(subresult) == 0:
                    continue

                assigned_ids = connection.handler.add_relations_to_entry_with_reconid(
                    recon_id = take_reconid,
                    entry_table = proc.relation_type.source_entrytype,
                    relations = [
                        (rel, e_gives)
                        for rel, e_gives in subresult
                    ]
                )

                # Assign recon_ids to the created entries
                for assign_recon_id, (new_rel, new_entry) in zip(assigned_ids, subresult):
                    new_entry.recon_id = assign_recon_id

        else:
            # Assign virtual Recon IDs to the generated entries (do not necessarily correspond with db)
            # These will be negative numbers
            for subresult in procedural_result:
                for rel, e in subresult:
                    connection.handler.add_procedural_entry(e)

        # Return result
        return procedural_result


    def _run_procedure_transitive[T_takes: Entry, T_gives: Entry, T_intermediate: Entry](
        self,
        proc: ExploreProcedure[T_takes, T_gives],
        take_entries: list[T_takes],
    ) -> list[list[tuple[Relation[T_takes, T_gives], T_gives]]]:
        import chemrecon.core.populate_query_handler

        assert issubclass(proc.relation_type, ComposedRelation)

        # First, check that the relation already exists in the DB.
        db_result = self._run_procedure_database(proc, take_entries)

        # If we found results for all queries
        if all(len(subresult) > 0 for subresult in db_result):
            return db_result

        # Get rel1 results
        rel_1_proc, rel_2_proc = self.protocol.transitive_subprocedures[proc.relation_type]
        rel_1_result = self._run_procedure(rel_1_proc, take_entries)

        # Filter by rel1 and e_inter
        rel_1_result_filtered: list[list[tuple[Relation[T_takes, T_intermediate], T_intermediate]]] = list()
        for subresult_1 in rel_1_result:
            subresult_1_filtered: list[tuple[Relation[T_takes, T_intermediate], T_intermediate]] = list()
            for r1, e_inter in subresult_1:
                if proc.relation_type.filter_rel_1(r1) and proc.relation_type.filter_intermediate(e_inter):
                    subresult_1_filtered.append((r1, e_inter))
            rel_1_result_filtered.append(subresult_1_filtered)

        # Results of rel1 are passed as take_entries to rel_2
        intermediate_map: dict[T_takes, list[tuple[Relation[T_takes, T_intermediate], T_intermediate]]] = {
            e_take: subresult
            for e_take, subresult in zip(take_entries, rel_1_result_filtered)
        }
        intermediate_entries: set[T_intermediate] = set()
        for rels in intermediate_map.values():
            intermediate_entries.update(e for _, e in rels)
        intermediate_entries_list = [e for e in intermediate_entries if proc.relation_type.filter_intermediate(e)]

        # Get rel2 results
        rel_2_result = self._run_procedure(rel_2_proc, intermediate_entries_list)
        rel_2_result_dict: dict[T_intermediate, list[tuple[Relation[T_intermediate, T_gives], T_gives]]] = {
            e_inter: subresult
            for e_inter, subresult in zip(intermediate_entries_list, rel_2_result)
        }

        # TODO filter rel2 results (not currently used)

        # Make list of entries to return
        result: list[list[tuple[Relation[T_takes, T_gives], T_gives]]] = list()
        for i, take_entry in enumerate(take_entries):
            # Consider routes through all intermediates
            subresult: list[tuple[Relation[T_takes, T_gives], T_gives]] = list()
            for rel_1, e_inter in intermediate_map[take_entry]:
                # Add all from this intermediate
                subresult.extend(
                    (proc.relation_type(rel_1 = rel_1, rel_2 = rel_2, intermediate = e_inter), e_gives)
                    for rel_2, e_gives in rel_2_result_dict[e_inter]
                )
            result.append(subresult)


        if isinstance(connection.handler, chemrecon.core.populate_query_handler.PopulateQueryHandler):
            # Update the database from rel_1 if procedural
            for take_entry, subresult in zip(take_entries, rel_1_result_filtered):
                assert take_entry.recon_id is not None
                assigned_ids = connection.handler.add_relations_to_entry_with_reconid(
                    recon_id = take_entry.recon_id,
                    entry_table = proc.relation_type.source_entrytype,
                    relations = [
                        (rel, e_gives)
                        for rel, e_gives in subresult
                    ]
                )
                for assign_recon_id, (new_rel, new_entry) in zip(assigned_ids, subresult):
                    new_entry.recon_id = assign_recon_id

            # Update the database from rel_2 if procedural
            for intermediate_entry, subresult in zip(intermediate_entries_list, rel_2_result):
                assert intermediate_entry.recon_id is not None
                assigned_ids = connection.handler.add_relations_to_entry_with_reconid(
                    recon_id = intermediate_entry.recon_id,
                    entry_table = proc.relation_type.intermediate_entrytype,
                    relations = [
                        (rel, e_gives)
                        for rel, e_gives in subresult
                    ]
                )
                for assign_recon_id, (new_rel, new_entry) in zip(assigned_ids, subresult):
                    new_entry.recon_id = assign_recon_id

            # Update the database for the transitive relation
            for take_entry, subresult in zip(take_entries, result):
                connection.handler.add_relations_to_entry_with_reconid(
                    recon_id = take_entry.recon_id,
                    entry_table = proc.relation_type.source_entrytype,
                    relations = [
                        (rel, e_gives)
                        for rel, e_gives in subresult
                    ]
                )

        else:
            # Assign virtual reconids to the generated entries
            for subresult in result:
                for rel, e in subresult:
                    connection.handler.add_procedural_entry(e)

        return result
