import json
import pathlib
from pathlib import Path
from typing import Any, Self, TypeAlias

import networkx as nx
import yaml
from graphviz import Digraph
from pydantic import BaseModel, ConfigDict, Field, model_validator

import gdsfactory as gf
from gdsfactory.component import Component
from gdsfactory.config import PATH
from gdsfactory.typings import Anchor, Delta, Port, Ports
from gdsfactory.utils import is_component_spec


class OrthogonalGridArray(BaseModel):
    """Orthogonal grid array config.

    Parameters:
        columns: number of columns.
        rows: number of rows.
        column_pitch: column pitch.
        row_pitch: row pitch.
    """

    columns: int = 1
    rows: int = 1
    column_pitch: float = 0
    row_pitch: float = 0


class GridArray(BaseModel):
    """Non-orthogonal grid array config.

    Parameters:
        num_a: number of elements in the a-dimension.
        num_b: number of elements in the b-dimension.
        pitch_a: x-y pitch in the a-dimension.
        pitch_b: x-y pitch in the b-dimension.
    """

    num_a: int = 1
    num_b: int = 1
    pitch_a: tuple[float, float] = (1.0, 0.0)
    pitch_b: tuple[float, float] = (0.0, 1.0)


Array: TypeAlias = OrthogonalGridArray | GridArray


class Instance(BaseModel):
    """Instance of a component.

    Parameters:
        component: component name.
        settings: input variables.
        info: information (polarization, wavelength ...).
        array: array config to make create an array reference for this instance
        virtual: if True, creates a virtual instance without actual geometry
    """

    component: str
    settings: dict[str, Any] = Field(default_factory=dict)
    info: dict[str, Any] = Field(default_factory=dict, exclude=True)
    array: Array | None = None
    virtual: bool = False

    model_config = ConfigDict(extra="forbid")

    @model_validator(mode="before")
    @classmethod
    def update_settings_and_info(cls, values: dict[str, Any]) -> dict[str, Any]:
        """Validator to update component, settings and info based on the component."""
        if not isinstance(values, dict):
            raise ValueError(
                f"Expected a dictionary of the format {{'component': 'mycomponent', 'settings': {{...}}}}. Got a {type(values)} ({values})"
            )
        try:
            component = values.get("component")
        except KeyError:
            raise KeyError(
                f"'component' must be defined in instance definition. Got keys: {list(values.keys())}"
            )
        settings = values.get("settings", {})
        info = values.get("info", {})

        assert is_component_spec(component)

        import gdsfactory as gf

        c = gf.get_component(component, settings=settings)
        component_info = c.info.model_dump(exclude_none=True)
        component_settings = c.settings.model_dump(exclude_none=True)

        # Create a new dict to avoid modifying the input in place
        result = values.copy()
        result["info"] = {**component_info, **info}
        result["settings"] = {**component_settings, **settings}
        result["component"] = c.function_name or component
        return result

    @model_validator(mode="after")
    def validate_array(self) -> Self:
        if isinstance(self.array, OrthogonalGridArray):
            if self.array.columns < 2:
                self.array.columns = 1
            if self.array.rows < 2:
                self.array.rows = 1
            if self.array.columns == 1 and self.array.rows == 1:
                self.array = None
        elif isinstance(self.array, GridArray):
            if self.array.num_a < 2:
                self.array.num_a = 1
            if self.array.num_b < 2:
                self.array.num_b = 1
            if self.array.num_a == 1 and self.array.num_b == 1:
                self.array = None
        return self


class Placement(BaseModel):
    x: str | float | None = None
    y: str | float | None = None
    xmin: str | float | None = None
    ymin: str | float | None = None
    xmax: str | float | None = None
    ymax: str | float | None = None
    dx: Delta = 0
    dy: Delta = 0
    port: str | Anchor | None = None
    rotation: float = 0
    mirror: bool | str | float = False

    def __getitem__(self, key: str) -> Any:
        """Allows to access the placement attributes as a dictionary."""
        return getattr(self, key, 0)

    model_config = ConfigDict(extra="forbid")


class Bundle(BaseModel):
    links: dict[str, str]
    settings: dict[str, Any] = Field(default_factory=dict)
    routing_strategy: str = "route_bundle"

    model_config = ConfigDict(extra="forbid")


class Net(BaseModel):
    """Net between two ports.

    Parameters:
        p1: instance_name,port 1.
        p2: instance_name,port 2.
        name: route name.
    """

    p1: str
    p2: str
    settings: dict[str, Any] = Field(default_factory=dict)
    name: str | None = None

    def __init__(self, **data: Any) -> None:
        """Initialize the net."""
        global _route_counter
        super().__init__(**data)
        # If route name is not provided, generate one automatically
        if self.name is None:
            self.name = f"route_{_route_counter}"
            _route_counter += 1


class Netlist(BaseModel):
    """Netlist defined component.

    Parameters:
        instances: dict of instances (name, settings, component).
        placements: dict of placements.
        connections: dict of connections.
        routes: dict of routes.
        name: component name.
        info: information (polarization, wavelength ...).
        ports: exposed component ports.
        settings: input variables.
        nets: list of nets.
        warnings: warnings.
    """

    pdk: str = ""
    instances: dict[str, Instance] = Field(default_factory=dict)
    placements: dict[str, Placement] = Field(default_factory=dict)
    connections: dict[str, str] = Field(default_factory=dict)
    routes: dict[str, Bundle] = Field(default_factory=dict)
    name: str | None = None
    info: dict[str, Any] = Field(default_factory=dict)
    ports: dict[str, str] = Field(default_factory=dict)
    settings: dict[str, Any] = Field(default_factory=dict, exclude=True)
    nets: list[Net] = Field(default_factory=list)
    warnings: dict[str, Any] = Field(default_factory=dict)

    model_config = ConfigDict(extra="forbid")

    @model_validator(mode="after")
    def validate_instance_names(self) -> Self:
        self.instances = {
            _validate_instance_name(k): v for k, v in self.instances.items()
        }
        return self


_route_counter = 0


def to_yaml_graph_networkx(
    netlist: Netlist, nets: list[Net]
) -> tuple[nx.Graph, dict[str, str], dict[str, tuple[float, float]]]:
    """Generates a netlist graph using NetworkX."""
    connections = netlist.connections
    placements = netlist.placements
    graph = nx.Graph()
    graph.add_edges_from(
        [
            (",".join(k.split(",")[:-1]), ",".join(v.split(",")[:-1]))
            for k, v in connections.items()
        ]
    )
    pos = {k: (v["x"], v["y"]) for k, v in placements.items()}
    labels = {k: ",".join(k.split(",")[:1]) for k in placements}

    for node, placement in placements.items():
        if not graph.has_node(
            node
        ):  # Check if the node is already in the graph (from connections), to avoid duplication.
            graph.add_node(node)
            pos[node] = (placement.x, placement.y)

    for net in nets:
        graph.add_edge(net.p1.split(",")[0], net.p2.split(",")[0])

    return graph, labels, pos


def to_graphviz(
    instances: dict[str, Instance],
    placements: dict[str, Placement],
    nets: list[Net],
    show_ports: bool = True,
) -> Digraph:
    """Generates a netlist graph using Graphviz."""
    from graphviz import Digraph

    # Graphviz implementation
    dot = Digraph(comment="Netlist Diagram")
    dot.attr(dpi="300", layout="neato", overlap="false")

    all_ports: list[tuple[str, Ports]] = []

    # Retrieve all the ports in the component
    for name, instance in instances.items():
        if hasattr(instance, "component"):
            instance_component = instance.component
        else:
            instance_component = instance["component"]  # type: ignore[index]
        ports_ = gf.get_component(instance_component).ports
        all_ports.append((name, ports_))

    for node, placement in placements.items():
        ports = dict(all_ports).get(node)
        assert ports is not None

        if not ports or not show_ports:
            label = node
        else:
            top_ports: list[Port] = []
            right_ports: list[Port] = []
            bottom_ports: list[Port] = []
            left_ports: list[Port] = []

            for port in ports:
                if 0 <= port.orientation < 45 or 315 <= port.orientation < 360:
                    right_ports.append(port)
                elif 45 <= port.orientation < 135:
                    bottom_ports.append(port)
                elif 135 <= port.orientation < 225:
                    left_ports.append(port)
                elif 225 <= port.orientation < 315:
                    top_ports.append(port)

            # Format ports for Graphviz record structure in anticlockwise order
            port_labels: list[str] = []

            if left_ports:
                left_ports_label = " | ".join(
                    f"<{port.name}> {port.name}" for port in reversed(left_ports)
                )
                port_labels.append(f"{{ {left_ports_label} }}")

            middle_row: list[str] = []

            if top_ports:
                top_ports_label = " | ".join(
                    f"<{port.name}> {port.name}" for port in top_ports
                )
                middle_row.append(f"{{ {top_ports_label} }}")

            middle_row.append(node)

            if bottom_ports:
                bottom_ports_label = " | ".join(
                    f"<{port.name}> {port.name}" for port in reversed(bottom_ports)
                )
                middle_row.append(f"{{ {bottom_ports_label} }}")

            port_labels.append(f"{{ {' | '.join(middle_row)} }}")

            if right_ports:
                right_ports_label = " | ".join(
                    f"<{port.name}> {port.name}" for port in right_ports
                )
                port_labels.append(f"{{ {right_ports_label} }}")

            label = " | ".join(port_labels)

        x = placement.x if hasattr(placement, "x") else placement["x"]
        y = placement.y if hasattr(placement, "y") else placement["y"]
        pos = f"{x},{y}!"
        dot.node(node, label=label, pos=pos, shape="record")

    for net in nets:
        p1 = net.p1 if hasattr(net, "p1") else net["p1"]  # type: ignore[index]
        p2 = net.p2 if hasattr(net, "p2") else net["p2"]  # type: ignore[index]

        p1_instance = p1.split(",")[0]
        p1_port = p1.split(",")[1]

        p2_instance = p2.split(",")[0]
        p2_port = p2.split(",")[1]

        dot.edge(f"{p1_instance}:{p1_port}", f"{p2_instance}:{p2_port}", dir="none")

    return dot


class Link(BaseModel):
    """Link between instances.

    Parameters:
        instance1: instance name 1.
        instance2: instance name 2.
        port1: port name 1.
        port2: port name 2.
    """

    instance1: str
    instance2: str
    port1: str
    port2: str


class Schematic(BaseModel):
    """Schematic."""

    netlist: Netlist = Field(default_factory=Netlist)
    nets: list[Net] = Field(default_factory=list)
    placements: dict[str, Placement] = Field(default_factory=dict)
    links: list[Link] = Field(default_factory=list)

    def add_instance(
        self, name: str, instance: Instance, placement: Placement | None = None
    ) -> None:
        self.netlist.instances[name] = instance
        if placement:
            self.add_placement(name, placement)

    def add_placement(
        self,
        instance_name: str,
        placement: Placement,
    ) -> None:
        """Add placement to the netlist.

        Args:
            instance_name: instance name.
            placement: placement.
        """
        self.placements[instance_name] = placement
        self.netlist.placements[instance_name] = placement

    def from_component(self, component: Component) -> None:
        raise NotImplementedError
        n = component.to_yaml()
        self.netlist = Netlist.model_validate(n)

    def add_net(self, net: Net) -> None:
        """Add a net between two ports."""
        self.nets.append(net)
        if net.name not in self.netlist.routes:
            assert net.name is not None
            self.netlist.routes[net.name] = Bundle(
                links={net.p1: net.p2}, settings=net.settings
            )
        else:
            self.netlist.routes[net.name].links[net.p1] = net.p2

    def to_graphviz(self, show_ports: bool = True) -> Digraph:
        """Generates a netlist graph using Graphviz.

        Args:
            show_ports: whether to show ports or not.
        """
        return to_graphviz(
            self.netlist.instances, self.placements, self.nets, show_ports
        )

    def to_yaml_graph_networkx(
        self,
    ) -> tuple[nx.Graph, dict[str, str], dict[str, tuple[float, float]]]:
        return to_yaml_graph_networkx(self.netlist, self.nets)

    def plot_graphviz(self, interactive: bool = False, splines: str = "ortho") -> None:
        """Plots the netlist graph (Automatic fallback to networkx).

        Args:
            interactive: whether to plot the graph interactively or not.
            splines: type of splines to use for the graph.

        """
        dot = self.to_graphviz()
        plot_graphviz(dot, interactive=interactive, splines=splines)

    def write_netlist(
        self, netlist: dict[str, Any], filepath: str | pathlib.Path | None = None
    ) -> str:
        """Returns netlist as YAML string.

        Args:
            netlist: netlist to write.
            filepath: Optional file path to write to.
        """
        yaml_string = yaml.safe_dump(netlist)
        if filepath:
            filepath = pathlib.Path(filepath)
            filepath.write_text(yaml_string)
        return yaml_string


def plot_graphviz(
    graph: Digraph, interactive: bool = False, splines: str = "ortho"
) -> None:
    """Plots the netlist graph (Automatic fallback to networkx)."""
    from IPython.display import Image, display

    valid_splines = ["ortho", "spline", "line", "polyline", "curved"]
    if splines not in valid_splines:
        raise ValueError(f"Invalid splines value. Choose from {valid_splines}")

    graph.graph_attr["splines"] = splines

    if interactive:
        graph.view()
    else:
        png_data = graph.pipe(format="png")
        display(Image(data=png_data))


def write_schema(
    model: type[BaseModel] = Netlist, schema_path_json: Path = PATH.schema_netlist
) -> None:
    s = model.model_json_schema()
    schema_path_yaml = schema_path_json.with_suffix(".yaml")

    with open(schema_path_json, "w") as f:
        json.dump(s, f, indent=2)
    with open(schema_path_yaml, "w") as f:
        yaml.dump(s, f)


def _validate_instance_name(name: str) -> str:
    if "," in name:
        raise ValueError(
            f"Having a ',' in an instance name is not supported. The ',' is used for port-delineation. Got: {name!r}."
        )
    if "-" in name:
        raise ValueError(
            f"Having a '-' in an instance name is not supported. The '-' is used for bundle routing. Got: {name!r}."
        )
    if ":" in name:
        raise ValueError(
            f"Having a ':' in an instance name is not supported. The ':' is used for bundle routing. Got: {name!r}."
        )
    return name
