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

import itertools
from pathlib import Path
from typing import TYPE_CHECKING, Generic, Type, TypeVar

from power_grid_model import ComponentType

from power_grid_model_ds._core.model.enums.nodes import NodeType

if TYPE_CHECKING:
    from power_grid_model_ds._core.model.grids.base import Grid

G = TypeVar("G", bound="Grid")


def serialize_to_str(grid: "Grid") -> str:
    """See Grid.__str__()"""
    grid_str = ""

    for transformer3 in grid.three_winding_transformer:
        nodes = [transformer3.node_1.item(), transformer3.node_2.item(), transformer3.node_3.item()]
        for combo in itertools.combinations(nodes, 2):
            grid_str += f"S{combo[0]} S{combo[1]} {transformer3.id.item()},3-transformer\n"

    for branch in grid.branches:
        from_node = grid.node.get(id=branch.from_node).record
        to_node = grid.node.get(id=branch.to_node).record

        from_node_str = f"S{from_node.id}" if from_node.node_type == NodeType.SUBSTATION_NODE else str(from_node.id)
        to_node_str = f"S{to_node.id}" if to_node.node_type == NodeType.SUBSTATION_NODE else str(to_node.id)

        suffix_str = str(branch.id.item())
        if branch.from_status.item() == 0 or branch.to_status.item() == 0:
            suffix_str = f"{suffix_str},open"

        if branch.id in grid.transformer.id:
            suffix_str = f"{suffix_str},{ComponentType.transformer.value}"
        elif branch.id in grid.link.id:
            suffix_str = f"{suffix_str},{ComponentType.link.value}"
        elif branch.id in grid.line.id:
            pass  # no suffix needed
        elif branch.id in grid.generic_branch.id:
            suffix_str = f"{suffix_str},{ComponentType.generic_branch.value}"
        elif branch.id in grid.asym_line.id:
            suffix_str = f"{suffix_str},{ComponentType.asym_line.value}"
        else:
            raise ValueError(f"Branch {branch.id} is not a transformer, link, line, generic_branch or asym_line")

        grid_str += f"{from_node_str} {to_node_str} {suffix_str}\n"
    return grid_str


def deserialize_from_str(grid_class: type[G], *args: str) -> G:
    """See Grid.from_txt()"""
    return _TextSource(grid_class).load_from_txt(*args)


def deserialize_from_txt_file(grid_class: type[G], txt_file_path: Path) -> G:
    """See Grid.from_txt_file()"""
    with open(txt_file_path, "r", encoding="utf-8") as f:
        txt_lines = f.readlines()
    return deserialize_from_str(grid_class, *txt_lines)


class _TextSource(Generic[G]):
    """Class for handling text sources.

    Text sources are only intended for test purposes so that a grid can quickly be designed from a text file.
    Moreover, these text sources are compatible with the grid editor at https://csacademy.com/app/graph_editor/

    Example of a text file:
        S1 2
        2 3
        3 4 transformer
        4 5
        S1 7

    See docs/examples/3_drawing_a_grid.md for more information.
    """

    def __init__(self, grid_class: Type[G]):
        self.grid: G = grid_class.empty()

    def load_from_txt(self, *args: str) -> G:
        """Load a grid from text"""

        text_lines = [line for arg in args for line in arg.strip().split("\n")]

        txt_nodes, txt_branches = self.read_txt(text_lines)
        self.add_nodes(txt_nodes)
        self.add_branches(txt_branches)
        self.grid.set_feeder_ids()
        return self.grid

    @staticmethod
    def read_txt(txt_lines: list[str]) -> tuple[set, dict]:
        """Extract assets from text"""

        txt_nodes = set()
        txt_branches = {}
        for text_line in txt_lines:
            if not text_line.strip() or text_line.startswith("#"):
                continue  # skip empty lines and comments
            try:
                from_node_str, to_node_str, *comments = text_line.strip().split()
            except ValueError as err:
                raise ValueError(f"Text line '{text_line}' is invalid. Skipping...") from err
            comments = comments[0].split(",") if comments else []

            txt_nodes |= {from_node_str, to_node_str}
            txt_branches[(from_node_str, to_node_str)] = comments
        return txt_nodes, txt_branches

    def add_nodes(self, nodes: set[str]):
        """Add nodes to the grid"""
        source_nodes = {int(node[1:]) for node in nodes if node.startswith("S")}
        regular_nodes = {int(node) for node in nodes if not node.startswith("S")}

        if source_nodes.intersection(regular_nodes):
            raise ValueError("Source nodes and regular nodes have overlapping ids")

        for node_id in source_nodes:
            new_node = self.grid.node.empty(1)
            new_node.id = node_id
            new_node.node_type = NodeType.SUBSTATION_NODE
            self.grid.append(new_node, check_max_id=False)

        for node_id in regular_nodes:
            new_node = self.grid.node.empty(1)
            new_node.id = node_id
            self.grid.append(new_node, check_max_id=False)

    def add_branches(self, branches: dict[tuple[str, str], list[str]]):
        """Add branches to the grid"""
        for branch, comments in branches.items():
            self.add_branch(branch, comments)

    def add_branch(self, branch: tuple[str, str], comments: list[str]):
        """Add a branch to the grid"""
        from_node_str, to_node_str = branch
        from_node = int(from_node_str.replace("S", ""))
        to_node = int(to_node_str.replace("S", ""))

        if ComponentType.transformer.value in comments:
            new_branch = self.grid.transformer.empty(1)
        elif ComponentType.link.value in comments:
            new_branch = self.grid.link.empty(1)
        elif ComponentType.generic_branch.value in comments:
            new_branch = self.grid.generic_branch.empty(1)
        elif ComponentType.asym_line.value in comments:
            new_branch = self.grid.asym_line.empty(1)
        else:  # assume it is a line
            new_branch = self.grid.line.empty(1)

        branch_ids = [branch_id for branch_id in comments if branch_id.isdigit()]
        if branch_ids:
            if len(branch_ids) > 1:
                raise ValueError(f"Multiple branch ids found in row {branch} {','.join(comments)}")
            new_branch.id = int(branch_ids[0])

        new_branch.from_node = from_node
        new_branch.to_node = to_node
        new_branch.from_status = 1
        if "open" in comments:
            new_branch.to_status = 0
        else:
            new_branch.to_status = 1
        self.grid.append(new_branch, check_max_id=False)
