""" Protocols define how to explore an entrygraph.
"""

from __future__ import annotations

from typing import Optional, Callable

from chemrecon.entrygraph.explore_procedure import (
    ExploreProcedure, ExploreProcedureT1, ExploreProcedureT2,
    ExploreProcedureSym,
)
from chemrecon.entrygraph.filter import EntryFilter, RelationFilter, EntryFilterProcedure, RelationFilterProcedure
from chemrecon.schema import Entry, ProceduralRelation, Relation, ComposedRelation
from chemrecon.schema.direction import Direction


# Protocol specification
class ExplorationProtocol():
    """
    Represents a protocol for exploring relationships and entries within a graph structure.

    This class is designed to define and manage the exploration of entries and relationships in a graph-like
    data model. Users can specify sets of entry and relationship types, as well as define filters and
    procedures for exploration.
    """
    relation_types: set[tuple[type[Relation], Direction]]  #: List of relation types to traverse.
    entry_types: set[type[Entry]]  #: List of entry types involved in the specified relations
    relation_types_terminal: set[tuple[type[Relation], Direction]]  # Will not be explored, only added.
    entry_filters: dict[type[Entry], EntryFilter]  #: Optional filters for each type of entry.
    relation_filters: dict[type[Relation], RelationFilter]  #: Ditto, for relations.

    # Explore procedures
    explore_procedures: list[ExploreProcedure]
    post_explore_procedures: list[ExploreProcedure]
    transitive_subprocedures: dict[type[Relation], tuple[ExploreProcedure, ExploreProcedure]]

    def __init__(
            self,
            relation_types: set[type[Relation] | tuple[type[Relation], Direction]],
            relation_types_terminal: Optional[set[type[Relation]]] = None,
            entry_filters: Optional[dict[type[Entry], Callable[[Entry], bool]]] = None,
            relation_filters: Optional[dict[type[Relation], Callable[[Relation], bool]]] = None,
    ):
        """ Specify an exploration protocol.

        :param relation_types: A set of relation typess to explore. Optionally, (Relation, Direction) tuples can be
                passed to specify in which direction each relation can be traversed. By default, only the forwards
                direction is traversed. For symmetric relations, only the Direction.SYMMETRIC value is allowed.
        :type relation_types: set[type[Relation] | tuple[type[Relation], Direction]],
        :param relation_types_terminal: A set of relation types which are not used for expanding the graph, but are
                added if both endpoints were already found. Here, directionality is not specified.
        :type relation_types_terminal: Optional[set[type[Relation] | tuple[type[Relation], Direction]]]
        :param entry_filters: An optional dictionary of filters for each entry type. Filters should  be a function which
                accepts an entry, and returns False if the entry should not be included by the protocol.
        :type entry_filters: Optional[dict[type[Entry], Callable[[Entry], bool]]]
        :param relation_filters: Ditto, for relations.
        :type relation_filters: Optional[dict[type[Relation], Callable[[Relation], bool]]]
        """
        # Process relation types (assign forward/symmetric direction by default)
        self.relation_types = set()
        for item in relation_types:
            match item:
                case (reltype, direction):
                    # Given with direction
                    self.relation_types.add((reltype, direction))
                case reltype:
                    # Set forwards/symmetric direction by default
                    if reltype.symmetric:
                        self.relation_types.add((reltype, Direction.SYMMETRIC))
                    else:
                        self.relation_types.add((reltype, Direction.FORWARDS))

        self.relation_types_terminal = set()
        if relation_types_terminal is not None:
            for reltype in relation_types_terminal:
                # Set forwards/symmetric direction by default
                if reltype.symmetric:
                    self.relation_types_terminal.add((reltype, Direction.SYMMETRIC))
                else:
                    self.relation_types_terminal.add((reltype, Direction.BOTH))

        # Process filters
        self.entry_filters = dict()
        self.relation_filters = dict()
        if entry_filters is not None:
            for e_type, e_filter in entry_filters.items():
                self.entry_filters[e_type] = EntryFilterProcedure(
                    filter_proc = e_filter
                )
        if relation_filters is not None:
            for r_type, r_filter in relation_filters.items():
                self.relation_filters[r_type] = RelationFilterProcedure(
                    relation_type = r_type,
                    filter_proc = r_filter
                )

        self.entry_filters = (entry_filters) if entry_filters else dict()
        self.relation_filters: dict[type[Relation], Callable[[Relation], bool]] \
            = relation_filters if relation_filters else dict()

        # Set entry types based on the relations
        self.entry_types = set()
        for r, direction in self.relation_types:
            self.entry_types.add(r.source_entrytype)
            self.entry_types.add(r.target_entrytype)

        # List of finished ExplorationProcedure objects to call when exploring
        self.explore_procedures: list[ExploreProcedure] = list()
        self.post_explore_procedures: list[ExploreProcedure] = list()
        self.transitive_subprocedures: dict[type[Relation], tuple[ExploreProcedure, ExploreProcedure]] = dict()

        # Sanity checks of the given input
        for rtype, direction in self.relation_types:
            if rtype.source_entrytype not in self.entry_types:
                raise ValueError(f'Relation {rtype} refers to entry types not allowed in the graph.')
            if rtype.target_entrytype not in self.entry_types:
                raise ValueError(f'Relation {rtype} refers to entry types not allowed in the graph.')

        if len(self.relation_types_terminal.intersection(self.relation_types)) > 0:
            raise ValueError('Terminal relation types must not intersect relation types.')

        # Add 'None' for filters when not specified
        for entrytype in (
                self.entry_types.difference(self.entry_filters.keys())
        ):
            self.entry_filters[entrytype] = None
        for reltype in ({rel for rel, _ in self.relation_types}
                .union({rel for rel, _ in self.relation_types_terminal})
                .difference(self.relation_filters.keys())
        ):
            self.relation_filters[reltype] = None

        # Establish a series of exploration procedures
        # TODO accept no arguments, and assume (forwards) or (symmetric)

        for given_rels, proc_list in [
            [self.relation_types, self.explore_procedures],
            [self.relation_types_terminal, self.post_explore_procedures]
        ]:
            for rel_type, direction in given_rels:
                rel_type: type[Relation]
                if issubclass(rel_type, ProceduralRelation):
                    if direction not in {Direction.FORWARDS, Direction.SYMMETRIC}:
                        raise ValueError('Procedural relations can only be explored forwards.')

                match direction:
                    case Direction.FORWARDS:
                        # Forwards, explore from T1
                        assert not rel_type.symmetric
                        proc_list.append(ExploreProcedureT1(
                            relationtype = rel_type,
                            relation_filter = self.relation_filters[rel_type],
                            entry_filter = self.entry_filters[rel_type.target_entrytype]
                        ))
                    case Direction.BACKWARDS:
                        # Backwards, explore from T2
                        assert not rel_type.symmetric
                        proc_list.append(ExploreProcedureT2(
                            relationtype = rel_type,
                            relation_filter = self.relation_filters[rel_type],
                            entry_filter = self.entry_filters[rel_type.source_entrytype]
                        ))
                    case Direction.BOTH:
                        # Both
                        assert not rel_type.symmetric
                        proc_list.append(ExploreProcedureT1(
                            relationtype = rel_type,
                            relation_filter = self.relation_filters[rel_type],
                            entry_filter = self.entry_filters[rel_type.target_entrytype]
                        ))
                        proc_list.append(ExploreProcedureT2(
                            relationtype = rel_type,
                            relation_filter = self.relation_filters[rel_type],
                            entry_filter = self.entry_filters[rel_type.source_entrytype]
                        ))
                    case Direction.SYMMETRIC:
                        # Symmetric
                        assert rel_type.symmetric
                        proc_list.append(ExploreProcedureSym(
                            relationtype = rel_type,
                            relation_filter = self.relation_filters[rel_type],
                            entry_filter = self.entry_filters[rel_type.target_entrytype]
                        ))

                # If transitive, add sub procedures
                if issubclass(rel_type, ComposedRelation):
                    self.transitive_subprocedures[rel_type] = (
                        ExploreProcedureT1(
                            relationtype = rel_type.rel_type_1,
                            relation_filter = None, entry_filter = None
                        ),
                        ExploreProcedureT1(
                            relationtype = rel_type.rel_type_2,
                            relation_filter = None, entry_filter = None
                        )
                    )

        # done
        pass
