# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
#
# SPDX-License-Identifier: MPL-2.0

from enum import StrEnum
from typing import Literal

from power_grid_model import ComponentType

from power_grid_model_ds._core.visualizer.typing import ComponentTypeAppliance, ComponentTypeBranch
from power_grid_model_ds.arrays import BranchArray, NodeArray

_LARGE_NODE_ID_THRESHOLD = 10_000_000


class StyleClass(StrEnum):
    """Styling classes used in the visualizer."""

    NODE = "node"
    SUBSTATION_NODE = "substation_node"
    LARGE_ID_NODE = "large_id_node"
    BRANCH = "branch"
    OPEN_BRANCH = "open_branch"
    OPEN_BRANCH_FROM = "open_branch_from"
    OPEN_BRANCH_TO = "open_branch_to"
    LINE = "line"
    TRANSFORMER = "transformer"
    LINK = "link"
    GENERIC_BRANCH = "generic_branch"
    ASYM_LINE = "asym_line"
    SOURCE = "source"
    GENERATING_APPLIANCE = "generating_appliance"
    LOADING_APPLIANCE = "loading_appliance"
    OPEN_GENERATING_APPLIANCE = "open_generating_appliance"
    OPEN_LOADING_APPLIANCE = "open_loading_appliance"
    APPLIANCE_GHOST_NODE = "appliance_ghost_node"


def get_node_classification(node_arr: NodeArray) -> str:
    """Get the space separated string of styling classes for a node.

    Note:
        Large ID nodes are those with an ID greater than 10,000,000.
        The style class decreases the font size for better readability.

    """
    classes = [StyleClass.NODE]
    if node_arr.id > _LARGE_NODE_ID_THRESHOLD:
        classes.append(StyleClass.LARGE_ID_NODE)
    if node_arr.node_type == 1:
        classes.append(StyleClass.SUBSTATION_NODE)
    return " ".join((entry.value for entry in classes))


def get_branch_classification(
    branch_arr: BranchArray,
    component_type: ComponentTypeBranch | Literal[ComponentType.three_winding_transformer],
) -> str:
    """Get the space separated string of styling classes for a branch."""
    classes = [StyleClass.BRANCH]

    if component_type != ComponentType.line:
        type_to_vizclass = {
            ComponentType.transformer: StyleClass.TRANSFORMER,
            ComponentType.three_winding_transformer: StyleClass.TRANSFORMER,
            ComponentType.link: StyleClass.LINK,
            ComponentType.generic_branch: StyleClass.GENERIC_BRANCH,
            ComponentType.asym_line: StyleClass.ASYM_LINE,
        }
        classes.append(type_to_vizclass[component_type])

    if branch_arr.from_status == 0 or branch_arr.to_status == 0:
        classes.append(StyleClass.OPEN_BRANCH)
    if branch_arr.from_status == 0:
        classes.extend([StyleClass.OPEN_BRANCH_FROM])
    if branch_arr.to_status == 0:
        classes.extend([StyleClass.OPEN_BRANCH_TO])

    return " ".join((entry.value for entry in classes))


def get_appliance_edge_classification(
    appliance_arr,
    component_type: ComponentTypeAppliance,
) -> str:
    """Get the space separated string of styling classes for an appliance edge."""
    type_to_vizclass = {
        ComponentType.asym_load: StyleClass.LOADING_APPLIANCE,
        ComponentType.asym_gen: StyleClass.GENERATING_APPLIANCE,
        ComponentType.sym_load: StyleClass.LOADING_APPLIANCE,
        ComponentType.sym_gen: StyleClass.GENERATING_APPLIANCE,
        ComponentType.source: StyleClass.GENERATING_APPLIANCE,
        ComponentType.shunt: StyleClass.LOADING_APPLIANCE,
    }

    classes = [type_to_vizclass[component_type]]
    if appliance_arr.status == 0:
        if type_to_vizclass[component_type] == StyleClass.LOADING_APPLIANCE:
            classes.append(StyleClass.OPEN_LOADING_APPLIANCE)
        else:
            classes.append(StyleClass.OPEN_GENERATING_APPLIANCE)

    return " ".join((entry.value for entry in classes))
