#!/usr/bin/env python3
"""Validate and semantically patch coordinate-free Assembly DesignIR."""

from __future__ import annotations

import argparse
import copy
import json
from pathlib import Path
from typing import Any

import jsonschema


FORBIDDEN_PLACEMENT_KEYS = {
    "transform",
    "transforms",
    "matrix",
    "matrices",
    "placement",
    "placements",
    "position",
    "positions",
    "location",
    "locations",
    "xyz",
    "rpy",
    "quaternion",
    "translation",
}


def walk_forbidden(value: Any, path: str = "$") -> list[str]:
    failures: list[str] = []
    if isinstance(value, dict):
        for key, child in value.items():
            child_path = f"{path}.{key}"
            if key.lower() in FORBIDDEN_PLACEMENT_KEYS:
                failures.append(child_path)
            failures.extend(walk_forbidden(child, child_path))
    elif isinstance(value, list):
        for index, child in enumerate(value):
            failures.extend(walk_forbidden(child, f"{path}[{index}]"))
    return failures


def component_interfaces(payload: dict[str, Any]) -> dict[str, set[str]]:
    return {
        component["id"]: {interface["id"] for interface in component["interfaces"]}
        for component in payload["components"]
    }


def validate_references(payload: dict[str, Any]) -> list[str]:
    failures: list[str] = []
    interfaces = component_interfaces(payload)
    component_ids = set(interfaces)
    if len(component_ids) != len(payload["components"]):
        failures.append("component ids are not unique")

    relation_ids: set[str] = set()
    adjacency = {component_id: set() for component_id in component_ids}
    for relation in payload["relations"]:
        if relation["id"] in relation_ids:
            failures.append(f"duplicate relation id: {relation['id']}")
        relation_ids.add(relation["id"])
        endpoints: list[tuple[str, str]] = []
        for side in ("a", "b"):
            component_id, interface_id = relation[side].split(".", 1)
            endpoints.append((component_id, interface_id))
            if component_id not in interfaces:
                failures.append(
                    f"{relation['id']}.{side}: unknown component {component_id}"
                )
            elif interface_id not in interfaces[component_id]:
                failures.append(
                    f"{relation['id']}.{side}: unknown interface "
                    f"{component_id}.{interface_id}"
                )
        if all(component_id in component_ids for component_id, _ in endpoints):
            adjacency[endpoints[0][0]].add(endpoints[1][0])
            adjacency[endpoints[1][0]].add(endpoints[0][0])

    anchor = payload["kinematics"]["anchor_component"]
    if anchor not in component_ids:
        failures.append(f"unknown anchor component: {anchor}")
    elif component_ids:
        visited: set[str] = set()
        frontier = [anchor]
        while frontier:
            current = frontier.pop()
            if current in visited:
                continue
            visited.add(current)
            frontier.extend(adjacency[current] - visited)
        disconnected = sorted(component_ids - visited)
        if disconnected:
            failures.append(
                "constraint graph disconnected: " + ", ".join(disconnected)
            )

    group_ids: set[str] = set()
    for group in payload["kinematics"]["rigid_groups"]:
        if group["id"] in group_ids:
            failures.append(f"duplicate rigid group id: {group['id']}")
        group_ids.add(group["id"])
        for member in group["members"]:
            if member not in component_ids:
                failures.append(
                    f"rigid group {group['id']} has unknown member {member}"
                )

    joint_ids: set[str] = set()
    for joint in payload["kinematics"]["joints"]:
        joint_ids.add(joint["id"])
        for side in ("parent_group", "child_group"):
            if joint[side] not in group_ids:
                failures.append(
                    f"joint {joint['id']} has unknown {side} {joint[side]}"
                )
        for ref_key in ("axis_ref", "station_ref"):
            component_id, interface_id = joint[ref_key].split(".", 1)
            if (
                component_id not in interfaces
                or interface_id not in interfaces[component_id]
            ):
                failures.append(
                    f"joint {joint['id']} has unknown {ref_key} "
                    f"{joint[ref_key]}"
                )

    for transmission in payload["kinematics"]["transmissions"]:
        for key in ("input_joint", "output_joint"):
            if transmission[key] not in joint_ids:
                failures.append(
                    f"transmission {transmission['id']} has unknown "
                    f"{key} {transmission[key]}"
                )
        fixed_member = transmission.get("fixed_member")
        if fixed_member and fixed_member not in component_ids:
            failures.append(
                f"transmission {transmission['id']} has unknown fixed member "
                f"{fixed_member}"
            )

    parameter_ids = set(payload["parameters"])
    for parameter_id in payload["edit_interface"]["editable_parameters"]:
        if parameter_id not in parameter_ids:
            failures.append(f"unknown editable parameter: {parameter_id}")
        elif not payload["parameters"][parameter_id]["editable"]:
            failures.append(f"parameter is not editable: {parameter_id}")

    return failures


def apply_patch(payload: dict[str, Any], patch: dict[str, Any]) -> dict[str, Any]:
    if patch["target_model"] != payload["model_id"]:
        raise ValueError("Patch target_model does not match base model")
    result = copy.deepcopy(payload)
    groups = {
        group["id"]: group for group in result["kinematics"]["rigid_groups"]
    }
    for operation in patch["operations"]:
        kind = operation["op"]
        if kind == "add_parameter":
            result["parameters"][operation["id"]] = operation["value"]
            result["edit_interface"]["editable_parameters"].append(
                operation["id"]
            )
        elif kind == "add_component":
            result["components"].append(operation["value"])
        elif kind == "add_relation":
            result["relations"].append(operation["value"])
        elif kind == "add_rigid_group_member":
            groups[operation["group"]]["members"].append(
                operation["component"]
            )
        elif kind == "add_preserved_interface":
            result["edit_interface"]["preserved_interfaces"].append(
                operation["value"]
            )
        else:
            raise ValueError(f"Unsupported semantic patch operation: {kind}")
    result["model_id"] = f"{payload['model_id']}_with_output_shaft"
    return result


def validate(
    payload: dict[str, Any], schema: dict[str, Any]
) -> dict[str, Any]:
    schema_errors = sorted(
        jsonschema.Draft202012Validator(schema).iter_errors(payload),
        key=lambda error: list(error.absolute_path),
    )
    failures = [
        f"{'.'.join(map(str, error.absolute_path))}: {error.message}"
        for error in schema_errors
    ]
    forbidden = walk_forbidden(payload)
    failures.extend(
        f"coordinate placement key is forbidden: {path}" for path in forbidden
    )
    if not schema_errors:
        failures.extend(validate_references(payload))
    return {
        "status": "pass" if not failures else "fail",
        "schema_version": payload.get("schema_version"),
        "model_id": payload.get("model_id"),
        "coordinate_free": not forbidden,
        "component_count": len(payload.get("components", [])),
        "relation_count": len(payload.get("relations", [])),
        "joint_count": len(payload.get("kinematics", {}).get("joints", [])),
        "transmission_count": len(
            payload.get("kinematics", {}).get("transmissions", [])
        ),
        "failures": failures,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("model", type=Path)
    parser.add_argument("--schema", type=Path, required=True)
    parser.add_argument("--patch", type=Path)
    parser.add_argument("--resolved-output", type=Path)
    parser.add_argument("--report", type=Path)
    args = parser.parse_args()

    payload = json.loads(args.model.read_text(encoding="utf-8"))
    schema = json.loads(args.schema.read_text(encoding="utf-8"))
    if args.patch:
        patch = json.loads(args.patch.read_text(encoding="utf-8"))
        payload = apply_patch(payload, patch)
        if args.resolved_output:
            args.resolved_output.parent.mkdir(parents=True, exist_ok=True)
            args.resolved_output.write_text(
                json.dumps(payload, ensure_ascii=False, indent=2),
                encoding="utf-8",
            )

    report = validate(payload, schema)
    if args.report:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(
            json.dumps(report, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
    print(json.dumps(report, ensure_ascii=False))
    raise SystemExit(0 if report["status"] == "pass" else 1)


if __name__ == "__main__":
    main()