from abc import ABC, abstractmethod
from typing import TypedDict

from networkx import MultiDiGraph
from pandas import DataFrame

from graph.ingestion.dataiku.data_types import GraphDataikuDataType


class NodeData(TypedDict):
    """Type definition for serialized Node data."""

    id: str

    _id: str
    _label: str

    _dku_grp_def_id: str
    _dku_reserved_user_defined_pk_prop: str

    # Additional properties from the node are added dynamically


class RelationshipData(TypedDict):
    """Type definition for serialized Relationship data."""

    _src: str
    _dst: str
    _label: str

    _dku_grp_def_id: str
    # Additional properties from the relationship are added dynamically


class TableData(TypedDict):
    df: DataFrame
    schema: dict[str, GraphDataikuDataType]


class QueryResult(ABC):
    @abstractmethod
    def get_as_df(self) -> DataFrame: ...
    @abstractmethod
    def get_as_table(self) -> TableData: ...

    @abstractmethod
    def get_as_networkx(self) -> MultiDiGraph: ...

    @abstractmethod
    def close(self) -> None: ...
