from __future__ import annotations

from typing import Any, List, Union

from pydantic import BaseModel
from typing_extensions import Literal

from solutions.graph.models import (
    DefinitionId,
    EdgeGroupDefinitionMetadata,
    EdgeGroupId,
    FilterMetadata,
    GraphId,
    NodeGroupDefinitionMetadata,
    NodeGroupId,
)


class NodeViewConfigParams(BaseModel):
    color: str
    size: Literal["large", "medium", "normal", "small", "x-small"]
    icon: str


class FilterParams(BaseModel):
    column: str
    operator: Literal["matches", "excludes", "above", "below", "range"]
    value: Any
    max_value: Any


class NodeDefinitionParams(BaseModel):
    definition_id: DefinitionId
    source_dataset: str
    primary_col: str
    label_col: str
    property_list: List[str]

    filters_stored: List[FilterParams]
    filters_association: Literal["and", "or"]


class NodeMetadataParams(BaseModel):
    # Pydantic raises an error at import time when using |
    node_id: Union[NodeGroupId, None] = None
    node_group: str
    definitions: List[NodeDefinitionParams]


class AddOrUpdateNodeParams(BaseModel):
    graph_id: GraphId
    meta: NodeMetadataParams
    view_config: NodeViewConfigParams
    ref_version_token: str


class EdgeViewConfigParams(BaseModel):
    color: Union[str, None] = None
    size: int


class EdgeDefinitionParams(BaseModel):
    definition_id: DefinitionId
    edge_dataset: str

    source_column: str
    target_column: str

    property_list: List[str]

    filters_stored: List[FilterParams]
    filters_association: Literal["and", "or"]


class EdgeMetadataParams(BaseModel):
    # Pydantic raises an error at import time when using |
    edge_id: Union[EdgeGroupId, None] = None
    edge_group: str

    source_node_id: str
    target_node_id: str

    definitions: List[EdgeDefinitionParams]


class AddOrUpdateEdgeParams(BaseModel):
    graph_id: GraphId
    meta: EdgeMetadataParams
    view_config: EdgeViewConfigParams
    ref_version_token: str


def to_node_group_definition_metadata(new_group: NodeMetadataParams) -> List[NodeGroupDefinitionMetadata]:
    return [
        NodeGroupDefinitionMetadata(
            definition_id=definition.definition_id,
            source_dataset=definition.source_dataset,
            primary_col=definition.primary_col,
            label_col=definition.label_col,
            property_list=definition.property_list,
            filters_stored=[
                FilterMetadata(column=f.column, operator=f.operator, value=f.value, max_value=f.max_value)
                for f in definition.filters_stored
            ],
            filters_association=definition.filters_association,
        )
        for definition in new_group.definitions
    ]


def to_edge_group_definition_metadata(new_group: EdgeMetadataParams) -> List[EdgeGroupDefinitionMetadata]:
    return [
        EdgeGroupDefinitionMetadata(
            definition_id=definition.definition_id,
            edge_dataset=definition.edge_dataset,
            source_column=definition.source_column,
            target_column=definition.target_column,
            property_list=definition.property_list,
            filters_stored=[
                FilterMetadata(column=f.column, operator=f.operator, value=f.value, max_value=f.max_value)
                for f in definition.filters_stored
            ],
            filters_association=definition.filters_association,
        )
        for definition in new_group.definitions
    ]
