from __future__ import annotations

import re
import tarfile
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class TestMqlComponentContact:
    component_a: str
    port_a: str
    component_b: str
    port_b: str

    def other_endpoint(self, component_alias: str) -> tuple[str, str]:
        if component_alias == self.component_a:
            return self.component_b, self.port_b
        if component_alias == self.component_b:
            return self.component_a, self.port_a
        raise KeyError(component_alias)

    def port_for(self, component_alias: str) -> str:
        if component_alias == self.component_a:
            return self.port_a
        if component_alias == self.component_b:
            return self.port_b
        raise KeyError(component_alias)


@dataclass(frozen=True)
class TestMqlCirTopology:
    component_contacts: tuple[TestMqlComponentContact, ...]

    def contacts_for(self, component_alias: str) -> tuple[TestMqlComponentContact, ...]:
        return tuple(
            contact
            for contact in self.component_contacts
            if component_alias in (contact.component_a, contact.component_b)
        )


def load_test_mql_cir_topology(
    archive_path: str | Path,
    *,
    cir_member: str = "test_mql_.cir",
) -> TestMqlCirTopology:
    with tarfile.open(archive_path) as archive:
        cir_file = archive.extractfile(cir_member)
        if cir_file is None:
            raise ValueError(f"Missing AMESim circuit member: {cir_member}")
        cir_text = cir_file.read().decode("latin1")

    root = ET.fromstring(_topology_only_xml(cir_text))
    components = root.findall(".//COMPS_LIST/COMP")
    aliases = tuple(_required_text(component, "ALIAS") for component in components)
    contacts: dict[
        tuple[tuple[int, int], tuple[int, int]],
        TestMqlComponentContact,
    ] = {}
    directed_contacts: set[tuple[tuple[int, int], tuple[int, int]]] = set()

    for component_index, component in enumerate(components):
        ports = component.findall("./COMP_PORTS_LIST/COMP_PORT")
        for port_index, port in enumerate(ports):
            if port.findtext("PORT_CONNECT") != "1":
                continue
            for connection in port.findall("./CONNECT_LIST/CONNECT"):
                target_index = int(_required_text(connection, "CONNECT_ENTITY_NUM"))
                target_port_index = int(_required_text(connection, "CONNECT_ENTITY_PORT"))
                if target_index < 0 or target_index >= len(components):
                    raise ValueError(f"Component contact references unknown entity {target_index}")
                target_ports = components[target_index].findall("./COMP_PORTS_LIST/COMP_PORT")
                if target_port_index < 0 or target_port_index >= len(target_ports):
                    raise ValueError(
                        f"Component contact references unknown port {target_port_index} "
                        f"on {aliases[target_index]}"
                    )

                endpoint = (component_index, port_index)
                target_endpoint = (target_index, target_port_index)
                directed_contacts.add((endpoint, target_endpoint))
                key = tuple(sorted((endpoint, target_endpoint)))
                first, second = key
                contacts[key] = TestMqlComponentContact(
                    component_a=aliases[first[0]],
                    port_a=f"port_{first[1] + 1}",
                    component_b=aliases[second[0]],
                    port_b=f"port_{second[1] + 1}",
                )

    for endpoint, target_endpoint in directed_contacts:
        if (target_endpoint, endpoint) not in directed_contacts:
            raise ValueError(
                "AMESim component contact is not reciprocal: "
                f"{endpoint} -> {target_endpoint}"
            )

    return TestMqlCirTopology(component_contacts=tuple(contacts.values()))


def _topology_only_xml(cir_text: str) -> str:
    # AMESim expressions inside SUBMODEL contain unescaped && and <= operators.
    # Topology lives outside those blocks, so omit them before XML parsing.
    return re.sub(
        r"<SUBMODEL>.*?</SUBMODEL>",
        "<SUBMODEL />",
        cir_text,
        flags=re.DOTALL,
    )


def _required_text(element: ET.Element, child_name: str) -> str:
    value = element.findtext(child_name)
    if value is None:
        raise ValueError(f"Missing AMESim circuit element: {child_name}")
    return value