""" Contains functions to find relations (including procedural) related to an entry.
"""
from typing import DefaultDict

import chemrecon.schema

from chemrecon import Entry, Relation, ProceduralRelation, ComposedRelation
from chemrecon.schema.direction import Direction
from chemrecon.entrygraph.entrygraph import EntryGraph, Edge
from chemrecon.entrygraph.explore import explore
from chemrecon.schema.procedural_relation_entrygraph import ProceduralRelationEG
from chemrecon.entrygraph.explorationprotocol import ExplorationProtocol

import chemrecon.connection as globals

# Relation lists for each entry type
# ----------------------------------------------------------------------------------------------------------------------
_entrytype_relation_types: dict[
    type[Entry],
    list[tuple[type[Relation], Direction, ExplorationProtocol]]
] = DefaultDict(list)

for reltype in chemrecon.schema.relationtypes:
    if reltype.symmetric:
        # Symmetric
        protocol_spec = ExplorationProtocol(
            relation_types = {(reltype, Direction.SYMMETRIC)}
        )
        _entrytype_relation_types[reltype.source_entrytype].append(
            (reltype, Direction.SYMMETRIC, protocol_spec)
        )
    else:
        # Not symmetric
        protocol_spec_forward = ExplorationProtocol(
            relation_types = {(reltype, Direction.FORWARDS)},
        )
        _entrytype_relation_types[reltype.source_entrytype].append(
            (reltype, Direction.FORWARDS, protocol_spec_forward)
        )
        if not issubclass(reltype, ProceduralRelation):
            # Add backwards explorer only if not procedural
            protocol_spec_backward = ExplorationProtocol(
                relation_types = {(reltype, Direction.BACKWARDS)},
            )
            _entrytype_relation_types[reltype.target_entrytype].append(
                (reltype, Direction.BACKWARDS, protocol_spec_backward)
            )

# Getter functions
# ----------------------------------------------------------------------------------------------------------------------
def get_relations_from_entry[R: Relation](
        entry: Entry,
        relation_type: type[R],
        get_backward_relations: bool = True
) -> list[tuple[R, Entry]]:
    """ Get relations of a given type with the specified entry as the source.
    """
    # TODO support for get_backward_relations !
    # TODO entrygraph-based method will save procedural relations to the DB!
    # TODO better generics for this!

    if entry.recon_id is None:
        # Get recon_id first
        index_result = globals.handler.get_entry_by_index(entry)
        if not index_result:
            raise ValueError(f'Given entry: {entry} not found in database.')

    res = globals.handler.get_relations_with_entries_by_recon_ids(
        entry_type = type(entry),
        recon_ids = [entry.recon_id],
        relation_type = relation_type
    )
    out_list: list[tuple[Relation, Entry]] = list()
    for target, relations in res[0].items():
        for rel in relations:
            out_list.append((rel, target))
    return out_list

def get_all_relations(
        entry: Entry,
        get_backward_relations: bool = True,
        get_transitive_relations: bool = False,
        get_procedural_relations: bool = False,
        get_procedural_eg_relations: bool = False
) -> list[tuple[Relation, Entry]]:
    """ Get all relations of a given entry.
        Note: Procedural relations cannot be explored backwards.
    """
    # TODO change to entrygraph-based method to save procedural relations to the DB!
    # TODO better generics for this!

    if entry.recon_id is None:
        # Get recon_id first
        index_result = globals.handler.get_entry_by_index(entry)
        if not index_result:
            raise ValueError(f'Given entry: {entry} not found in database.')

    output_list: list[tuple[Relation, Entry]]

    # Create EG spec
    entrytypes: set[type[Entry]] = set()
    relation_types: set[tuple[type[Relation], Direction]] = set()
    for r_type, direction, _ in _entrytype_relation_types[type(entry)]:

        # Skip if not in defined relation types
        if issubclass(r_type, ComposedRelation) and not get_transitive_relations:
            continue
        if issubclass(r_type, ProceduralRelation):
            if not get_procedural_eg_relations: continue
            if issubclass(r_type, ProceduralRelationEG):
                if not get_procedural_eg_relations:
                    continue
        if (not get_backward_relations) and (type(entry) is not r_type.source_entrytype):
            continue

        # Add to specification
        entrytypes.add(r_type.source_entrytype)
        entrytypes.add(r_type.target_entrytype)

        if r_type.symmetric:
            relation_types.add((r_type, Direction.SYMMETRIC))
        else:
            if r_type.source_entrytype == type(entry):
                relation_types.add((r_type, Direction.FORWARDS))
            if get_backward_relations and r_type.target_entrytype == type(entry):
                relation_types.add((r_type, Direction.BACKWARDS))

    protocol_all = ExplorationProtocol(
        relation_types = relation_types,
    )

    # Create an EntryGraph and explore to depth 1
    eg = EntryGraph(
        initial_entries = {entry}
    )
    explore(eg, protocol = protocol_all, steps = 1)

    # Get out-edges of a node (vertex with index 0 is the starting node)
    res = eg.get_out_edges_of_vertex(0)  # list of (Edge, Vertex)
    return [
        (edge.relation, vertex.entry)
        for edge, vertex in res if isinstance(edge, Edge)  # Filter out artificial edges
    ]
