from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any, List, Optional, cast

import networkx
import pandas as pd
try:
    # Try IPython < 8.0 location
    from IPython.core.display import HTML, DisplayHandle, display
except ImportError:
    # IPython >= 8.0 moved DisplayHandle to IPython.display
    from IPython.display import HTML, DisplayHandle, display  # type: ignore
from pydantic import ValidationError

from outhad_edge import RETE_CONFIG
from outhad_edge.backend import JupyterServer, ServerManager
from outhad_edge.backend.callback import list_dataprocessor, list_dataprocessor_mock
from outhad_edge.backend.tracker import (
    collect_data_performance,
    time_performance,
    track,
    tracker,
)
from outhad_edge.eventstream.types import EventstreamType
from outhad_edge.exceptions.server import ServerErrorWithResponse
from outhad_edge.exceptions.widget import WidgetParseError
from outhad_edge.preprocessing_graph.interface import (
    CombineHandlerPayload,
    CreateNodeErrorDesc,
    FieldErrorDesc,
    NodeLink,
    Payload,
    PreprocessingGraphConfig,
)
from outhad_edge.preprocessing_graph.nodes import (
    EventsNode,
    MergeNode,
    Node,
    SourceNode,
    build_node,
)
from outhad_edge.templates import PreprocessingGraphRenderer


class PreprocessingGraph:
    """
    Collection of methods for preprocessing graph construction and calculation.

    Parameters
    ----------
    source_stream : EventstreamType
        Source eventstream.

    Notes
    -----
    See :doc:`Preprocessing user guide</user_guides/preprocessing>` for the details.

    """

    DEFAULT_GRAPH_URL = "https://static.server.outhad_edge.com/package/@rete/preprocessing-graph/version/2/dist/preprocessing-graph.umd.js"
    DEFAULT_GRAPH_STYLE_URL = (
        "https://static.server.outhad_edge.com/package/@rete/preprocessing-graph/version/2/dist/style.css"
    )
    root: SourceNode
    _combine_result: EventstreamType | None

    _ngraph: networkx.DiGraph
    __server_manager: ServerManager | None = None
    __server: JupyterServer | None = None

    @time_performance(  # type: ignore
        scope="preprocessing_graph",
        event_name="init",
    )
    def __init__(self, source_stream: EventstreamType) -> None:
        self._source_stream = source_stream
        self.root = SourceNode(source=source_stream)
        self._combine_result = None
        self._ngraph = networkx.DiGraph()
        self._ngraph.add_node(self.root)

    def add_node(self, node: Node, parents: List[Node]) -> None:
        """
        Add node to ``PreprocessingGraph`` instance.

        Parameters
        ----------
        node : Node
            An instance of either ``EventsNode`` or ``MergeNode``.
        parents : list of Nodes

            - If ``node`` is ``EventsNode`` - only 1 parent must be defined.
            - If ``node`` is ``MergeNode`` - at least 2 parents have to be defined.

        Returns
        -------
        None

        See Also
        --------
        PreprocessingGraph.combine : Start PreprocessingGraph recalculation.
        .EventsNode : Regular nodes of a preprocessing graph.
        .MergeNode : Merge nodes of a preprocessing graph.

        """
        self.__valiate_already_exists(node)
        self.__validate_not_found(parents)

        if node.events is not None:
            self.__validate_schema(node.events)

        if not isinstance(node, MergeNode) and len(parents) > 1:
            raise ValueError("multiple parents are only allowed for merge nodes!")

        self._ngraph.add_node(node)

        for parent in parents:
            self._ngraph.add_edge(parent, node)

    @time_performance(
        scope="preprocessing_graph",
        event_name="combine",
    )
    def combine(self, node: Node) -> EventstreamType:
        """
        Run calculations from the ``SourceNode`` up to the specified ``node``.

        Parameters
        ----------
        node : Node
            Instance of either ``SourceNode``, ``EventsNode`` or ``MergeNode``.

        Returns
        -------
        EventstreamType
            ``Eventstream`` with all changes applied by data processors.
        """
        self.__validate_not_found([node])

        if isinstance(node, SourceNode):
            result_eventstream = node.events.copy()

        elif isinstance(node, EventsNode):
            result_eventstream = self._combine_events_node(node)

        else:
            result_eventstream = self._combine_merge_node(node)

        collect_data_performance(
            scope="preprocessing_graph",
            called_params={"node": node},
            performance_data={
                "parent": {
                    "index": getattr(self.__get_node_eventstream(node), "_eventstream_index", None),
                    "hash": getattr(self.__get_node_eventstream(node), "_hash", None),
                },
                "child": {"index": result_eventstream._eventstream_index, "hash": result_eventstream._hash},
            },
            eventstream_index=getattr(self.__get_node_eventstream(node), "_eventstream_index", None),
            parent_eventstream_index=getattr(self.__get_node_eventstream(node), "_eventstream_index", None),
            child_eventstream_index=result_eventstream._eventstream_index,
        )
        return result_eventstream

    def __get_node_eventstream(self, node: Node) -> pd.DataFrame | pd.Series | None:
        if events := getattr(node, "events", None):
            return events
        return None

    def _combine_events_node(self, node: EventsNode) -> EventstreamType:
        parent = self._get_events_node_parent(node)
        parent_events = self.combine(parent)
        return node.processor._get_new_data(eventstream=parent_events)

    def _combine_merge_node(self, node: MergeNode) -> EventstreamType:
        parents = self._get_merge_node_parents(node)
        curr_eventstream: Optional[EventstreamType] = None

        for parent_node in parents:
            if curr_eventstream is None:
                curr_eventstream = self.combine(parent_node)
            else:
                new_eventstream = self.combine(parent_node)
                curr_eventstream.index_order = new_eventstream.index_order
                curr_eventstream.append_eventstream(new_eventstream)
        node.events = curr_eventstream

        return cast(EventstreamType, curr_eventstream)

    def get_parents(self, node: Node) -> List[Node]:
        """
        Show parents of the specified ``node``.

        Parameters
        ----------
        node : Node
            Instance of one of the classes SourceNode, EventsNode or MergeNode.

        Returns
        -------
        list of Nodes

        """
        self.__validate_not_found([node])
        parents: List[Node] = []

        for parent in self._ngraph.predecessors(node):
            parents.append(parent)
        return parents

    def _get_merge_node_parents(self, node: MergeNode) -> List[Node]:
        parents = self.get_parents(node)
        if len(parents) == 0:
            raise ValueError("orphan merge node!")

        return parents

    def _get_events_node_parent(self, node: EventsNode) -> Node:
        parents = self.get_parents(node)
        if len(parents) > 1:
            raise ValueError("invalid preprocessing graph: events node has more than 1 parent")

        return parents[0]

    def __validate_schema(self, eventstream: EventstreamType) -> bool:
        return self.root.events.schema.is_equal(eventstream.schema)

    def __valiate_already_exists(self, node: Node) -> None:
        if node in self._ngraph.nodes:
            raise ValueError("node already exists!")

    def __validate_not_found(self, nodes: List[Node]) -> None:
        for node in nodes:
            if node not in self._ngraph.nodes:
                raise ValueError("node not found!")

    @property
    def graph_url(self) -> str:
        env_url = os.getenv("RETE_PREPROCESSING_GRAPH_URL", "")
        return env_url if env_url else self.DEFAULT_GRAPH_URL

    @property
    def graph_style_url(self) -> str:
        env_style = os.getenv("RETE_PREPROCESSING_GRAPH_STYLE_URL", "")
        return env_style if env_style else self.DEFAULT_GRAPH_STYLE_URL

    @time_performance(
        scope="preprocessing_graph",
        event_name="display",
    )
    def display(self, width: int = 960, height: int = 600) -> DisplayHandle:
        """
        Show constructed ``PreprocessingGraph``.

        Parameters
        ----------
        width : int, default 960
            Width of plot in pixels.
        height : int, default 600
            Height of plot in pixels.

        Returns
        -------
            Rendered preprocessing graph.
        """
        if not self.__server_manager:
            self.__server_manager = ServerManager()

        if not self.__server:
            self.__server = self.__server_manager.create_server()
            self.__server.register_action("list-dataprocessor-mock", list_dataprocessor_mock)
            self.__server.register_action("list-dataprocessor", list_dataprocessor)
            self.__server.register_action("set-graph", self._set_graph_handler)
            self.__server.register_action("get-graph", self._export_handler)
            self.__server.register_action("combine", self._combine_handler)

        render = PreprocessingGraphRenderer()

        collect_data_performance(
            scope="preprocessing_graph",
            event_name="metadata",
            called_params={"width": width, "height": height},
            eventstream_index=self.root.events._eventstream_index,
        )

        return display(
            HTML(
                render.show(
                    server_id=self.__server.pk,
                    env=self.__server_manager.check_env(),
                    width=width,
                    height=height,
                    graph_url=self.graph_url,
                    graph_style_url=self.graph_style_url,
                    kernel_id=self.__server.kernel_id,
                    tracking_hardware_id=RETE_CONFIG.user.pk,
                    tracking_eventstream_index=self._source_stream._eventstream_index,
                    tracking_scope="preprocessing_graph",
                )
            )
        )

    def _export_handler(self, payload: dict[str, Any]) -> dict:
        """
        Show ``PreprocessingGraph`` as a dict.

        Parameters
        ----------
        payload : dict

        Returns
        -------
        dict

        """
        source, target, link = "source", "target", "links"
        graph = self._ngraph
        data = {
            "directed": graph.is_directed(),
            "nodes": [n.export() for n in graph],
            link: [{source: u.pk, target: v.pk} for u, v, d in graph.edges(data=True)],
        }
        return data

    def __resolve_path(self, relative_path: str) -> str:
        path = Path(relative_path)
        absolute_path = str(path.resolve())
        return absolute_path

    @time_performance(
        scope="preprocessing_graph",
        event_name="export_to_file",
    )
    def export_to_file(self, filename: str) -> None:
        """
        Export constructed ``PreprocessingGraph`` to a json file.

        Parameters
        ----------
        filename : str
            Path to the json file.

        Returns
        -------
        None
        """

        # validate access rights to file

        graph_path = self.__resolve_path(filename)

        if not os.access(os.path.dirname(graph_path), os.W_OK):
            raise PermissionError(f"{os.path.dirname(graph_path)} is not writable")

        # export preprocessing_graph
        data = self._export_handler({})
        with open(graph_path, "w") as f:
            json.dump(data, f)

    @time_performance(
        scope="preprocessing_graph",
        event_name="import_from_file",
    )
    def import_from_file(self, filename: str) -> None:
        """
        Import constructed ``PreprocessingGraph`` from a json file.

        Parameters
        ----------
        filename : str
            Path to the json file.

        Returns
        -------
        None
        """
        graph_path = self.__resolve_path(filename)

        # validate access rights to file
        if not os.access(os.path.dirname(graph_path), os.R_OK):
            raise PermissionError(f"{os.path.dirname(graph_path)} is not readable")

        with open(graph_path, "r") as f:
            data = json.load(f)
        if validation_error := self._validate_payload(payload=data):
            raise ValueError("Invalid json file: %s" % validation_error)

        payload: Payload = Payload(**data)  # type: ignore
        self._set_graph_handler(payload=payload)

    def _combine_handler(self, payload: CombineHandlerPayload) -> None:
        node = self._find_node(payload["node_pk"])
        if not node:
            raise ServerErrorWithResponse(message="node not found!", type="unexpected_error")
        self._combine_result = self.combine(node)

    @property
    @time_performance(
        scope="preprocessing_graph",
        event_name="combine_result",
    )
    def combine_result(self) -> Optional[EventstreamType]:
        """
        Keep and get the last combining result from preprocessing graph GUI.

        Returns
        -------
        EventstreamType or None
            Preprocessed eventstream.

        """
        return self._combine_result

    def _set_graph_handler(self, payload: Payload) -> dict:
        current_graph = self._ngraph
        current_root = self.root

        def restore_graph() -> None:
            self._ngraph = current_graph
            self.root = current_root

        try:
            self._set_graph(payload=payload)
            return self._export_handler({})
        except ServerErrorWithResponse as err:
            restore_graph()
            raise err
        except Exception as err:
            restore_graph()
            raise ServerErrorWithResponse(message=str(err), type="unexpected_error")

    def _validate_payload(self, payload: Payload) -> None | Exception:
        try:
            config = PreprocessingGraphConfig(**payload)
            return None
        except Exception as err:
            return err

    def _set_graph(self, payload: Payload) -> None:
        """
        Payload example:

        {
            "nodes": [
                {
                    "name": "SourceNode",
                    "pk": "0dc3b706-e6cc-401e-96f7-6a45d3947d5c"
                },
                {
                    "name": "EventsNode",
                    "pk": "07921cb0-60b8-45af-928d-272d1b622b25",
                    "processor": {
                        "name": "SimpleGroup",
                        "values": {"event_name": "add_to_cart", "event_type": "group_alias"},
                    },
                },
                {
                    "name": "EventsNode",
                    "pk": "114251ae-0f03-45e6-a163-af51bb02dfd5",
                    "processor": {
                        "name": "SimpleGroup",
                        "values": {"event_name": "logout", "event_type": "group_alias"},
                    },
                },
            ],
            "links": [
                {
                    'source': '0dc3b706-e6cc-401e-96f7-6a45d3947d5c',
                    'target': '07921cb0-60b8-45af-928d-272d1b622b25'
                },
                {
                    'source': '07921cb0-60b8-45af-928d-272d1b622b25',
                    'target': '114251ae-0f03-45e6-a163-af51bb02dfd5'
                }
            ]
        }

        """
        errors: List[CreateNodeErrorDesc] = []
        nodes: List[Node] = []

        # create nodes & validate params
        for node in payload["nodes"]:
            node_pk = node["pk"]
            processor = node.get("processor", {})
            processor_name = processor.get("name", None) if processor else None
            processor_params = processor.get("values", None) if processor else None
            description = node.get("description", None)

            try:
                actual_node = build_node(
                    source_stream=self.root.events,
                    pk=node_pk,
                    node_name=node["name"],
                    processor_name=processor_name,
                    processor_params=processor_params,
                    descriptionn=description,
                )
                nodes.append(actual_node)
            except Exception as error:
                error_desc = self._build_node_error_desc(node_pk=node_pk, error=error)
                errors.append(error_desc)

        if errors:
            raise ServerErrorWithResponse(
                message="set preprocessing graph error", type="create_nodes_error", errors=errors
            )

        self._ngraph = networkx.DiGraph()

        # add nodes
        for created_node in nodes:
            if isinstance(created_node, SourceNode):
                self.root = created_node
            self._ngraph.add_node(created_node)

        # add links
        # @TODO: validate links (graph structure)
        for link in payload["links"]:
            source = self._find_node(link["source"])
            target = self._find_node(link["target"])
            if not source:
                raise ServerErrorWithResponse(message="source not found", type="create_link_error")
            if not target:
                raise ServerErrorWithResponse(message="target not found", type="create_link_error")
            self._ngraph.add_edge(source, target)

    def _build_node_error_desc(self, node_pk: str, error: Exception) -> CreateNodeErrorDesc:
        if isinstance(error, ValidationError):
            return self._build_pydantic_error_desc(
                node_pk=node_pk,
                validation_error_exception=error,
            )

        if isinstance(error, WidgetParseError):
            field_errors: List[FieldErrorDesc] = (
                [{"field": error.field_name, "msg": str(error)}] if error.field_name else []
            )
            return {
                "type": "node_error",
                "msg": str(error),
                "node_pk": node_pk,
                "fields_errors": field_errors,
            }

        return {
            "type": "node_error",
            "msg": str(error),
            "node_pk": node_pk,
            "fields_errors": [],
        }

    def _build_pydantic_error_desc(
        self, node_pk: str, validation_error_exception: ValidationError
    ) -> CreateNodeErrorDesc:
        raw_errs = validation_error_exception.errors()
        result_errors: List[FieldErrorDesc] = []

        for raw_err in raw_errs:
            loc = raw_err.get("loc", ())
            field = next(iter(loc), None)
            msg = raw_err.get("msg")

            result_errors.append(
                {
                    "field": str(field),
                    "msg": msg,
                }
            )

        return {"type": "node_error", "node_pk": node_pk, "msg": "node error", "fields_errors": result_errors}

    def _find_parents_by_links(self, target_node: str, link_list: list[NodeLink]) -> list[Node]:
        parents: list[str] = []
        for node in link_list:
            if node["target"] == target_node:
                parents.append(node["source"])

        parent_nodes = [self._find_node(parent) for parent in parents]
        return parent_nodes  # type: ignore

    def _find_node(self, pk: str) -> Node | None:
        for node in self._ngraph:
            if node.pk == pk:
                return node
        else:
            return None
