from __future__ import annotations

from typing import Optional, ClassVar, Generator

import rustworkx as rx

from chemrecon.entrygraph.draw import defaultdrawer
from chemrecon.entrygraph.explore_procedure import ExploreProcedure
from chemrecon.schema import Entry, Relation

from chemrecon import connection as connection

type VertexIndex = int
type EdgeIndex = int
type ReconID = int


# Vertices and Edges
# --------------------------------------------------------------------------------------------------------------
# Vertex and Edge objects are used as the payload for the RustWorkX graph.

class VertexBase:
    pass


class Vertex[E: Entry](VertexBase):
    """ A database entry in a graph context.
        None is used as the type parameter for artificial vertices (those not corresponding to a DB object).
        A vertex is initial if its generation is 0.
    """
    entry: E  #: The entry represented by this vertex
    recon_id: ReconID  #:
    vertex_index: VertexIndex  #: Index in the graph representation.
    generation: int  #: Distance from initial entries.

    def __init__(self, entry: E, generation: int):
        assert entry.recon_id is not None
        self.entry = entry
        self.recon_id = entry.recon_id
        self.generation = generation


class EdgeBase:
    pass


class Edge[R: Relation | None](EdgeBase):
    """ A database relation in a graph context.
        None is used as the type parameter for artificial edges (those not corresponding to a DB object).
    """
    relation: R  #:
    recon_id_1: ReconID  #:
    recon_id_2: ReconID  #:
    edge_index: EdgeIndex  #: Index in the graph representation

    def __init__(self, relation: R):
        self.relation = relation
        self.recon_id_1 = relation.recon_id_1
        self.recon_id_2 = relation.recon_id_2


# Special case graph vertices and edges which do not represent database entries or relations
class SourceVertexArtificial(VertexBase):
    """ Artificial vertex used as a source"""

    def __init__(self):
        pass


class SourceEdgeArtificial(EdgeBase):
    """ Artificial edge from the SourceVertex to initial vertices in the graph. """

    def __init__(self):
        pass


class ReturnEdgeArtificial(EdgeBase):
    """ Temporary edge added from sink vertices back to the SourceVertex for pagerank purposes. """

    def __init__(self):
        pass


# Generic EntryGraph
# --------------------------------------------------------------------------------------------------------------
class EntryGraph:
    """
    Represents a directed graph structure of entries and relations.

    :ivar g: RustWorkX directed graph object used to store vertices and edges.
    :type g: rx.PyDiGraph
    :ivar index_entry: Maps Entry objects to corresponding vertex indices.
    :type index_entry: dict[Entry, VertexIndex]
    :ivar index_reconid: Maps (Entry type, ReconID) tuples to vertex indices.
    :type index_reconid: dict[tuple[type[Entry], ReconID], VertexIndex]
    :ivar index_relation: Maps Relation objects to corresponding edge indices.
    :type index_relation: dict[Relation, EdgeIndex]
    :ivar initial_entries: Set of Entry objects that serve as the initial vertices in the graph.
    :type initial_entries: set[Entry]
    :ivar initial_vertices: Set of vertex indices corresponding to the initial entries.
    :type initial_vertices: set[VertexIndex]
    :ivar parent_entrygraph: Reference to the parent EntryGraph if this graph
        is a subgraph view.
    :type parent_entrygraph: Optional[EntryGraph]
    """
    g: rx.PyDiGraph

    # Indices for looking up vertices / edges based on the underlying entries
    index_entry: dict[Entry, VertexIndex]
    index_reconid: dict[tuple[type[Entry], ReconID], VertexIndex]
    index_relation: dict[Relation, EdgeIndex]

    # Misc
    initial_entries: set[Entry]
    initial_vertices: set[VertexIndex]

    # Parentage
    parent_entrygraph: Optional[EntryGraph]  # If created as a subgraph view, gives the parent graph

    # List of finished ExplorationProcedure objects to call when exploring
    explore_procedures: ClassVar[list[ExploreProcedure]]
    post_explore_procedures: ClassVar[list[ExploreProcedure]]
    transitive_subprocedures: ClassVar[dict[type[Relation], tuple[ExploreProcedure, ExploreProcedure]]]

    # TODO scoring as abstract class variables


    def __init__(
        self,  # TODO specification refactor
        initial_entries: set[Entry]
    ):
        """ Initialize an entrygraph with a set of initial vertices.
        """
        # Consistency check
        # TODO init vertices in allowed types

        # Create
        self.index_entry = dict()
        self.index_reconid = dict()
        self.index_relation = dict()

        # Create graph
        self.g = rx.PyDiGraph(multigraph = True)

        # Add initial vertices
        # The initial vertices may not have recon-ids attached, so we add those by DB lookup
        for entry in initial_entries:
            if entry.recon_id is None:
                res = connection.handler.get_entry_by_index(entry)
                if res is None:
                    continue
                entry.recon_id = res.recon_id

        # If no initial vertices have recon ids, abort
        if all(e.recon_id is None for e in initial_entries):
            raise ValueError('No initial entries have recon ids.')

        self.initial_entries = initial_entries
        self.initial_vertices = {
            self.add_vertex(e, generation = 0)
            for e in initial_entries
        }

    # Getters
    # ----------------------------------------------------------------------------------------------------------
    def get_vertex_index_by_recon_id(self, entry_type: type[Entry], recon_id: ReconID) -> Optional[VertexIndex]:
        """
        Retrieve the vertex index associated with the given entry type and recon ID.
        If no matching vertex index is found, the method returns None.

        :param entry_type: The class type of the entry being queried.
        :type entry_type: type[Entry]
        :param recon_id: The recon ID that uniquely identifies the entry.
        :type recon_id: ReconID
        :return: The vertex index corresponding to the given entry type and recon ID, or None
                 if not found.
        :rtype: Optional[VertexIndex]
        """
        return self.index_reconid.get((entry_type, recon_id), None)

    def get_vertex_by_entry(self, entry: Entry) -> Optional[VertexIndex]:
        """
        Retrieves the vertex index associated with a given entry and returns None if it does not exist.

        :param entry: The entry object to search for in the index map.
        :type entry: Entry
        :return: The vertex index associated with the given entry, or None if the entry
                 does not exist in the mapping.
        :rtype: Optional[VertexIndex]
        """
        return self.index_entry.get(entry, None)

    def get_vertex_by_vertex_index(self, vertex_index: VertexIndex) -> Optional[Vertex]:
        """
        Retrieve the vertex corresponding to the given vertex index, returns None if it does not exist.
        :param vertex_index: Index of the vertex to retrieve.
        :type vertex_index: VertexIndex
        :return: The vertex associated with the given index or None if it does not exist.
        :rtype: Optional[Vertex]
        """
        return self.g.get_node_data(vertex_index)

    def get_out_edges_of_vertex(self, vertex_index: VertexIndex) -> list[tuple[Edge, Vertex]]:
        """
        Retrieve all outgoing edges of a vertex along with the corresponding target
        vertices.

        This method provides a list of tuples where each tuple contains the edge data
        and the target vertex data for the edges originating from the specified
        vertex. This function is useful for analysing the outgoing connections
        and relationships in a graph data structure.

        :param vertex_index: Index of the vertex whose outgoing edges are to be queried.
        :type vertex_index: VertexIndex
        :return: A list of tuples with each tuple containing edge data and the target
                 vertex data.
        :rtype: list[tuple[Edge, Vertex]]
        """
        return [
            (edge_data, self.g.get_node_data(target_index))
            for _, target_index, edge_data in self.g.out_edges(vertex_index)
        ]

    # Iterators
    def vertices(self) -> Generator[Vertex, None, None]:
        """ Generator over vertices (entries). """
        yield from self.g.nodes()

    def edges(self) -> Generator[Edge, None, None]:
        """ Generator over edges (relations). """
        yield from self.g.edges()

    # Adders
    # ----------------------------------------------------------------------------------------------------------
    def add_vertex(self, entry: Entry, generation: int = -1) -> VertexIndex:
        """ Add a new entry to the graph. If it already exists, return the index.
        """
        if entry.recon_id is None:
            raise ValueError('Cannot construct vertex from entry without reconid.')
        lookup = self.get_vertex_index_by_recon_id(type(entry), entry.recon_id)
        if lookup is not None:
            return lookup

        # Add the vertex
        vertex = Vertex(entry, generation = generation)
        v_index: VertexIndex = self.g.add_node(vertex)
        self.index_reconid[(type(entry), entry.recon_id)] = v_index
        self.index_entry[entry] = v_index
        vertex.vertex_index = v_index
        return v_index

    def add_edge(
        self,
        source_v_index: VertexIndex,
        target_v_index: VertexIndex,
        relation: Relation
    ) -> EdgeIndex:
        """ Add a relation between the specified source and target vertices
        """
        # Try returning existing edge
        try:
            return self.index_relation[relation]
        except KeyError:
            # Does not already exist
            edge = Edge(relation)
            edge_index: EdgeIndex = self.g.add_edge(source_v_index, target_v_index, edge)
            edge.edge_index = edge_index
            self.index_relation[relation] = edge_index
            return edge_index

    def add_vertex_from[T1: Entry, T2: Entry](
        self,
        from_index: VertexIndex,
        relation: Relation[T1, T2],
        entry: T2,
        generation: int
    ) -> VertexIndex:
        """ Adds a new entry vertex along with a relation. Returns vertex_index of the newly creted vertex.
            If the vertex already exists, return add the edge and return the vertex index.
        """
        target_v_index = self.add_vertex(entry, generation = generation)
        self.add_edge(from_index, target_v_index, relation)
        return target_v_index

    # Draw
    # ----------------------------------------------------------------------------------------------------------
    def draw(
        self,
        filename: str = None,
        filetype: str = 'jpg',
        scores: Optional[dict[Entry, float]] = None
    ):
        """ Draw using default settings, returning an image, or write to the disk if the filename is specified.
            For more settings, create an EntryGraphDrawer instance.
        """
        return defaultdrawer.draw(self, filename = filename, filetype = filetype, scores = scores)

    def show(self, scores: Optional[dict[Entry, float]] = None):
        """ Same as draw, but displays the image in a window instead of saving to disk.
        """
        return self.draw(scores = scores).show()
