#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from torch_geometric.typing import EdgeType

from .abc import Abstract_ADBPyG_Controller
from .typings import Json


class ADBPyG_Controller(Abstract_ADBPyG_Controller):
    """ArangoDB-PyG controller.

    Responsible for controlling how nodes & edges are handled when
    transitioning from PyG to ArangoDB.

    You can derive your own custom ADBPyG_Controller.
    """

    def _prepare_pyg_node(self, pyg_node: Json, node_type: str) -> Json:
        """Prepare a PyG node object before it gets inserted into its
        designated ArangoDB collection.

        Given a JSON representation of a PyG node, you can modify it
        before it gets inserted into its ArangoDB collection,
        and/or derive a custom vertex id by updating the "_key" attribute
        of the vertex (otherwise the vertex's current "_key" value will be used).

        :param pyg_node: The PyG node object to (optionally) modify.
        :type pyg_node: adbnx_adapter.typings.Json
        :param node_type: The PyG Node Type of the node
        :type node_type: str
        :return: The PyG Node object
        :rtype: Dict[str, Any]
        """
        return pyg_node  # pragma: no cover

    def _prepare_pyg_edge(self, pyg_edge: Json, edge_type: EdgeType) -> Json:
        """Prepare a PyG edge object before it gets inserted into its
        designated ArangoDB collection.

        Given a JSON representation of a PyG edge, you can modify it
        before it gets inserted into its ArangoDB edge collection,
        and/or derive a custom edge id by setting the "_key" attribute
        of the edge (otherwise the "_key" will be randomly generated by ArangoDB).

        :param pyg_edge: The PyG edge object to (optionally) modify.
        :type pyg_edge: adbnx_adapter.typings.Json
        :param edge_type: The Edge Type of the PyG edge. Formatted
            as (from_collection, edge_collection, to_collection)
        :type edge_type: Tuple[str, str, str]
        :return: The PyG Edge object
        :rtype: Dict[str, Any]
        """
        return pyg_edge  # pragma: no cover
