import re
from copy import deepcopy
from dataclasses import dataclass
from typing import Any


_MODELICA_TYPES = {
    "real": "Real",
    "integer": "Integer",
    "boolean": "Boolean",
    "string": "String",
}
_BEVALUE_PATTERN = re.compile(r"\$([A-Za-z_][A-Za-z0-9_]*)\$")


@dataclass(frozen=True)
class CompilationResult:
    """The intermediate data and generated source produced by compilation."""

    graph: dict[str, Any]
    objects_by_id: dict[str, Any]
    modelica: str


def compile_graph(graph: dict[str, Any]) -> CompilationResult:
    """Clean, index, and emit a serialized component tree."""

    cleaned_graph = cleanup_graph(deepcopy(graph))
    objects_by_id = build_id_list(cleaned_graph)
    return CompilationResult(
        graph=cleaned_graph,
        objects_by_id=objects_by_id,
        modelica=emit_model(cleaned_graph, objects_by_id),
    )


def build_id_list(graph: dict[str, Any]) -> dict[str, Any]:
    """Index all addressable objects in a component tree by their stable ID."""

    id_list: dict[str, Any] = {}
    id_kinds: dict[str, str] = {}

    def add(item: dict[str, Any], description: str) -> None:
        item_id = item.get("id")
        if not item_id:
            raise ValueError(f"{description} has no ID")
        # Port IDs identify a port on a component definition and may therefore
        # recur in cloned component instances. Component and junction IDs are
        # document objects and must remain globally unique.
        if item_id in id_list and not (
            description == "port" and id_kinds[item_id] == "port"
        ):
            raise ValueError(f"Duplicate simulation object ID: {item_id}")
        id_list[item_id] = item
        id_kinds[item_id] = description

    def visit(component: dict[str, Any]) -> None:
        add(component, "component")
        interface = component.get("interface", {})
        for port in interface.get("inputs", []):
            add(port, "port")
        for port in interface.get("outputs", []):
            add(port, "port")

        implementation = component.get("implementation", {})
        if implementation.get("kind") != "graph":
            return
        nested_graph = implementation.get("graph", {})
        for junction in nested_graph.get("junctions", []):
            add(junction, "junction")
        for block in nested_graph.get("blocks", []):
            visit(block)

    visit(graph)
    return id_list


def emit_model(
    graph: dict[str, Any],
    id_list: dict[str, Any],
    indent: int = 0,
    connection_counts: dict[str, int] | None = None,
) -> str:
    """Emit a component and its nested definitions as Modelica source."""

    del id_list  # Kept in the public API for compiler extensions and inspection.
    indentation = "\t" * indent
    body_indent = "\t" * (indent + 1)
    model_name = identifier(graph["id"])
    lines = [f"{indentation}model {model_name}"]
    implementation = graph.get("implementation", {})
    implementation_kind = implementation.get("kind")
    nested_graph = implementation.get("graph", {})
    port_counts = connection_counts or _interface_connection_counts(graph)
    macros = _port_count_macros(graph, port_counts)

    if implementation_kind == "graph":
        for block in nested_graph.get("blocks", []):
            lines.extend(
                emit_model(
                    block,
                    {},
                    indent + 1,
                    _block_connection_counts(nested_graph, block["id"]),
                )
                .rstrip()
                .splitlines()
            )

    interface = graph.get("interface", {})
    for port in interface.get("inputs", []):
        lines.append(_port_declaration(port, "input", indent + 1, macros))
    for port in interface.get("outputs", []):
        lines.append(_port_declaration(port, "output", indent + 1, macros))

    for parameter in graph.get("parameters", []):
        parameter_type = modelica_type(parameter.get("type", "real"))
        parameter_name = identifier(parameter["name"])
        value = expand_bevalues(str(parameter.get("value", "0")), macros)
        lines.append(
            f"{body_indent}parameter {parameter_type} {parameter_name} = {value};"
        )

    if implementation_kind == "graph":
        for block in nested_graph.get("blocks", []):
            block_type = identifier(block["id"])
            block_name = identifier(block["name"])
            lines.append(f"{body_indent}{block_type} {block_name};")
        for junction in nested_graph.get("junctions", []):
            junction_type = modelica_type(junction.get("type", "signal"))
            lines.append(
                f"{body_indent}{junction_type} {_junction_name(junction['id'])};"
            )

    lines.append(f"{indentation}equation")
    if implementation_kind == "graph":
        blocks = {block["id"]: block for block in nested_graph.get("blocks", [])}
        junctions = {
            junction["id"]: junction for junction in nested_graph.get("junctions", [])
        }
        endpoint_indices: dict[tuple[str, str, str], int] = {}
        for connection in nested_graph.get("connections", []):
            source = _endpoint_expression(
                connection["source"], graph, blocks, junctions, endpoint_indices
            )
            target = _endpoint_expression(
                connection["target"], graph, blocks, junctions, endpoint_indices
            )
            # Connector types may require different equations in future.
            lines.append(f"{body_indent}{target} = {source};")
    else:
        equations = str(implementation.get("source", {}).get("equations", ""))
        equations = expand_bevalues(equations, macros)
        lines.extend(
            f"{body_indent}{line}" if line.strip() else ""
            for line in equations.splitlines()
        )

    lines.append(f"{indentation}end {model_name};")
    return "\n".join(lines) + "\n"


def cleanup_graph(graph: dict[str, Any]) -> dict[str, Any]:
    """Remove annotations and UI-only data from a serialized component tree."""

    def remove_key_with_lists(data: Any, target_key: str) -> None:
        if isinstance(data, dict):
            for key in list(data):
                if key == target_key:
                    del data[key]
                else:
                    remove_key_with_lists(data[key], target_key)
        elif isinstance(data, list):
            for item in data:
                remove_key_with_lists(item, target_key)

    for key in (
        "position",
        "rotation",
        "iconPosition",
        "icon",
        "annotations",
        "library",
        "properties",
    ):
        remove_key_with_lists(graph, key)
    return graph


def _port_declaration(
    port: dict[str, Any], direction: str, indent: int, macros: dict[str, str]
) -> str:
    port_type = modelica_type(port.get("type", "signal"))
    indentation = "\t" * indent
    port_name = identifier(port["name"])
    dimension = f"[${port_name}_N$]" if port.get("multipleConnections", False) else ""
    declaration = f"{indentation}{direction} {port_type} {port_name}{dimension};"
    return expand_bevalues(declaration, macros)


def _endpoint_expression(
    endpoint: dict[str, Any],
    owner: dict[str, Any],
    blocks: dict[str, dict[str, Any]],
    junctions: dict[str, dict[str, Any]],
    endpoint_indices: dict[tuple[str, str, str], int],
) -> str:
    if "junction" in endpoint:
        junction_id = endpoint["junction"]
        if junction_id not in junctions:
            raise ValueError(f"Connection references unknown junction {junction_id!r}")
        return _junction_name(junction_id)

    if "interface" in endpoint:
        port = _find_port(owner, endpoint["interface"])
        expression = identifier(port["name"])
        return _index_array_endpoint(
            expression, port, ("interface", owner["id"], port["id"]), endpoint_indices
        )

    block_id = endpoint.get("block")
    port_id = endpoint.get("port")
    block = blocks.get(block_id)
    if block is None:
        raise ValueError(f"Connection references unknown block {block_id!r}")
    port = _find_port(block, port_id)
    expression = f"{identifier(block['name'])}.{identifier(port['name'])}"
    return _index_array_endpoint(
        expression, port, ("block", block_id, port_id), endpoint_indices
    )


def _index_array_endpoint(
    expression: str,
    port: dict[str, Any],
    key: tuple[str, str, str],
    endpoint_indices: dict[tuple[str, str, str], int],
) -> str:
    if not port.get("multipleConnections", False):
        return expression
    endpoint_indices[key] = endpoint_indices.get(key, 0) + 1
    return f"{expression}[{endpoint_indices[key]}]"


def _block_connection_counts(
    graph: dict[str, Any], block_id: str
) -> dict[str, int]:
    counts: dict[str, int] = {}
    for connection in graph.get("connections", []):
        for endpoint in (connection.get("source", {}), connection.get("target", {})):
            if endpoint.get("block") == block_id and endpoint.get("port"):
                port_id = endpoint["port"]
                counts[port_id] = counts.get(port_id, 0) + 1
    return counts


def _interface_connection_counts(component: dict[str, Any]) -> dict[str, int]:
    implementation = component.get("implementation", {})
    if implementation.get("kind") != "graph":
        return {}
    counts: dict[str, int] = {}
    for connection in implementation.get("graph", {}).get("connections", []):
        for endpoint in (connection.get("source", {}), connection.get("target", {})):
            if endpoint.get("interface"):
                port_id = endpoint["interface"]
                counts[port_id] = counts.get(port_id, 0) + 1
    return counts


def _port_count_macros(
    component: dict[str, Any], connection_counts: dict[str, int]
) -> dict[str, str]:
    macros: dict[str, str] = {}
    interface = component.get("interface", {})
    for port in (*interface.get("inputs", []), *interface.get("outputs", [])):
        if port.get("multipleConnections", False):
            macros[f"{identifier(port['name'])}_N"] = str(
                connection_counts.get(port["id"], 0)
            )
    return macros


def expand_bevalues(text: str, values: dict[str, str]) -> str:
    """Replace BEdit ``$name$`` macros and reject unresolved compiler values."""

    def replace(match: re.Match[str]) -> str:
        name = match.group(1)
        if name not in values:
            raise ValueError(f"Unknown BEdit value ${name}$")
        return values[name]

    return _BEVALUE_PATTERN.sub(replace, text)


def _find_port(component: dict[str, Any], port_id: str) -> dict[str, Any]:
    interface = component.get("interface", {})
    ports = [*interface.get("inputs", []), *interface.get("outputs", [])]
    for port in ports:
        if port.get("id") == port_id:
            return port
    raise ValueError(
        f"Component {component.get('name', component.get('id', '?'))!r} "
        f"has no port {port_id!r}"
    )


def modelica_type(value: str) -> str:
    normalized = str(value).strip().lower()
    if normalized in {"signal", "signal array"}:
        return "Real"
    return _MODELICA_TYPES.get(normalized, identifier(str(value)))


def identifier(value: str) -> str:
    """Return a safe unquoted Modelica identifier."""

    normalized = re.sub(r"[^A-Za-z0-9_]", "_", str(value).strip())
    if not normalized:
        raise ValueError("Modelica names cannot be empty")
    if normalized[0].isdigit():
        normalized = f"model_{normalized}"
    return normalized


def _junction_name(junction_id: str) -> str:
    return identifier(f"junction_{junction_id}")