"""Assembly interface resolution + geometric connection validation (v0.1).

Resolves Assembly IR interface ``topology_refs`` against per-part / package topology maps,
transforms the resolved geometry from part-local into assembly/world coordinates, and validates
whether each simplified connection is GEOMETRICALLY PLAUSIBLE (centroid distance, bbox overlap,
normal alignment, semantic-role fit).

This is GEOMETRY VALIDATION ONLY. It does NOT model contact, bolt preload, or run a solver.
Connections remain proxies. Unresolved topology refs are reported honestly, never invented.

Outputs:
  - assembly/interface_resolution.json
  - diagnostics/assembly_connection_geometry.json
  - updated assembly/connection_graph.json (per-edge geometry_status fields)
  - updated simulation/assembly_cae_setup_draft.json (geometry_status; invalid -> needs_user_input)
"""
from __future__ import annotations

import json
import math
import zipfile
from pathlib import Path
from typing import Any

from aieng import FORMAT_VERSION
from aieng.converters.assembly_ir import (
    ASSEMBLY_CAE_DRAFT_PATH,
    ASSEMBLY_IR_PATH,
    CONNECTION_GRAPH_PATH,
    PROXY_CONNECTION_TYPES,
    _as_list,
    _collect_interfaces,
    build_assembly_cae_setup_draft,
    build_connection_graph,
)

INTERFACE_RESOLUTION_PATH = "assembly/interface_resolution.json"
ASSEMBLY_CONNECTION_GEOMETRY_PATH = "diagnostics/assembly_connection_geometry.json"
PACKAGE_TOPOLOGY_PATH = "geometry/topology_map.json"

# proximity thresholds expressed as a fraction of the larger interface bbox diagonal
_TOUCH_FRAC = 0.02     # gap <= this*scale -> touching/overlapping
_NEAR_FRAC = 0.5       # gap >  this*scale -> far apart
_NORMAL_SAME_DOT = 0.5  # mating faces expected anti-parallel; dot > this is suspicious

# semantic roles preferred by each proxy type (for evidence checks)
_PREFERRED_ROLES = {
    "bolted_proxy": {"bolt_hole", "mounting_face"},
    "welded_proxy": {"weld_face", "contact_face", "mounting_face"},
}


# ── small geometry helpers (numpy-free where possible) ────────────────────────

def _bbox_center(bb: list[float]) -> list[float]:
    return [(float(bb[0]) + float(bb[3])) / 2.0,
            (float(bb[1]) + float(bb[4])) / 2.0,
            (float(bb[2]) + float(bb[5])) / 2.0]


def _bbox_corners(bb: list[float]) -> list[list[float]]:
    xs, ys, zs = (bb[0], bb[3]), (bb[1], bb[4]), (bb[2], bb[5])
    return [[float(x), float(y), float(z)] for x in xs for y in ys for z in zs]


def _bbox_union(boxes: list[list[float]]) -> list[float] | None:
    boxes = [b for b in boxes if isinstance(b, (list, tuple)) and len(b) == 6]
    if not boxes:
        return None
    lo = [min(float(b[i]) for b in boxes) for i in range(3)]
    hi = [max(float(b[i + 3]) for b in boxes) for i in range(3)]
    return lo + hi


def _bbox_from_points(pts: list[list[float]]) -> list[float]:
    lo = [min(p[i] for p in pts) for i in range(3)]
    hi = [max(p[i] for p in pts) for i in range(3)]
    return lo + hi


def _bbox_diag(bb: list[float]) -> float:
    return math.sqrt(sum((float(bb[i + 3]) - float(bb[i])) ** 2 for i in range(3)))


def _bbox_gap(a: list[float], b: list[float]) -> float:
    """Euclidean gap between two AABBs (0 if they overlap on every axis)."""
    sep = []
    for i in range(3):
        d = max(a[i] - b[i + 3], b[i] - a[i + 3], 0.0)
        sep.append(d)
    return math.sqrt(sum(d * d for d in sep))


def _bbox_overlap(a: list[float], b: list[float]) -> bool:
    return all(a[i] <= b[i + 3] and b[i] <= a[i + 3] for i in range(3))


def _dist(a: list[float], b: list[float]) -> float:
    return math.sqrt(sum((float(a[i]) - float(b[i])) ** 2 for i in range(3)))


def _norm(v: list[float]) -> float:
    return math.sqrt(sum(float(c) ** 2 for c in v))


# ── transforms (translation + euler / matrix) ─────────────────────────────────

def _rotation_from_euler_deg(euler: list[float]) -> list[list[float]]:
    rx, ry, rz = (math.radians(float(a)) for a in euler)
    cx, sx, cy, sy, cz, sz = (math.cos(rx), math.sin(rx), math.cos(ry),
                              math.sin(ry), math.cos(rz), math.sin(rz))
    # R = Rz @ Ry @ Rx
    return [
        [cz * cy, cz * sy * sx - sz * cx, cz * sy * cx + sz * sx],
        [sz * cy, sz * sy * sx + cz * cx, sz * sy * cx - cz * sx],
        [-sy, cy * sx, cy * cx],
    ]


def _resolve_transform(transform: Any) -> tuple[list[list[float]], list[float], bool, str | None]:
    """Return (R 3x3, t 3, ok, note). Identity + note when missing/invalid."""
    identity = [[1.0, 0, 0], [0, 1.0, 0], [0, 0, 1.0]]
    if transform is None:
        return identity, [0.0, 0.0, 0.0], False, "missing_transform"
    if not isinstance(transform, dict):
        return identity, [0.0, 0.0, 0.0], False, "invalid_transform"
    t = transform.get("translation") or [0.0, 0.0, 0.0]
    try:
        tvec = [float(t[0]), float(t[1]), float(t[2])]
    except Exception:
        return identity, [0.0, 0.0, 0.0], False, "invalid_transform"
    matrix = transform.get("matrix")
    if matrix is not None:
        try:
            if len(matrix) == 4:  # 4x4 row-major
                R = [[float(matrix[r][c]) for c in range(3)] for r in range(3)]
                tvec = [float(matrix[r][3]) for r in range(3)]
            elif len(matrix) == 3:
                R = [[float(matrix[r][c]) for c in range(3)] for r in range(3)]
            else:
                return identity, tvec, False, "invalid_transform"
            return R, tvec, True, None
        except Exception:
            return identity, [0.0, 0.0, 0.0], False, "invalid_transform"
    euler = transform.get("rotation_euler_deg")
    if euler is not None:
        try:
            return _rotation_from_euler_deg(euler), tvec, True, None
        except Exception:
            return identity, tvec, False, "invalid_transform"
    return identity, tvec, True, None  # translation only


def _apply_point(R: list[list[float]], t: list[float], p: list[float]) -> list[float]:
    return [R[r][0] * p[0] + R[r][1] * p[1] + R[r][2] * p[2] + t[r] for r in range(3)]


def _apply_vector(R: list[list[float]], v: list[float]) -> list[float]:
    return [R[r][0] * v[0] + R[r][1] * v[1] + R[r][2] * v[2] for r in range(3)]


# ── interface resolution ──────────────────────────────────────────────────────

def _ref_ids(iface: dict[str, Any]) -> dict[str, list[str]]:
    refs = iface.get("topology_refs") or {}
    return {k: [str(x) for x in _as_list(refs.get(k))]
            for k in ("face_ids", "edge_ids", "vertex_ids")}


def resolve_interface(iface: dict[str, Any], topo_index: dict[str, dict[str, Any]] | None,
                      transform: Any) -> dict[str, Any]:
    """Resolve one interface against a part topology index, then transform to world coords."""
    refs = _ref_ids(iface)
    all_ids = refs["face_ids"] + refs["edge_ids"] + refs["vertex_ids"]
    topo_index = topo_index or {}
    R, tvec, tok, tnote = _resolve_transform(transform)

    found, unresolved = [], []
    for rid in all_ids:
        ent = topo_index.get(rid)
        if ent is not None:
            found.append(ent)
        else:
            unresolved.append(rid)

    boxes, centers, normals, areas, face_count = [], [], [], [], 0
    for ent in found:
        bb = ent.get("bounding_box")
        if isinstance(bb, (list, tuple)) and len(bb) == 6:
            boxes.append([float(x) for x in bb])
            c = ent.get("centroid")
            centers.append([float(x) for x in c] if isinstance(c, (list, tuple)) and len(c) == 3
                           else _bbox_center(bb))
        if ent.get("type") == "face":
            face_count += 1
            n = ent.get("normal")
            a = float(ent.get("area") or 0.0)
            if isinstance(n, (list, tuple)) and len(n) == 3 and _norm(n) > 1e-9:
                normals.append(([float(x) for x in n], a if a > 0 else 1.0))
            if a > 0:
                areas.append(a)

    local: dict[str, Any] = {}
    world: dict[str, Any] = {}
    if boxes:
        lbb = _bbox_union(boxes)
        # area-weighted centroid where areas known, else mean of centers
        if areas and len(areas) == len(centers):
            tot = sum(areas) or 1.0
            lcen = [sum(centers[i][k] * areas[i] for i in range(len(centers))) / tot for k in range(3)]
        else:
            lcen = [sum(c[k] for c in centers) / len(centers) for k in range(3)]
        lnorm = None
        if normals:
            wsum = [sum(n[0][k] * n[1] for n in normals) for k in range(3)]
            mag = _norm(wsum)
            if mag > 1e-9:
                lnorm = [wsum[k] / mag for k in range(3)]
        local = {"bbox": lbb, "centroid": lcen, "normal": lnorm,
                 "area": round(sum(areas), 6) if areas else None}
        # world geometry: transform corners, recompute AABB; transform centroid; rotate normal
        wcorners = [_apply_point(R, tvec, c) for c in _bbox_corners(lbb)]
        wbb = _bbox_from_points(wcorners)
        wcen = _apply_point(R, tvec, lcen)
        wnorm = None
        if lnorm is not None:
            wn = _apply_vector(R, lnorm)
            m = _norm(wn)
            wnorm = [wn[k] / m for k in range(3)] if m > 1e-9 else None
        world = {"bbox": [round(x, 6) for x in wbb], "centroid": [round(x, 6) for x in wcen],
                 "normal": [round(x, 6) for x in wnorm] if wnorm else None,
                 "area": local["area"]}

    if not all_ids:
        status = "unresolved"
    elif not found:
        status = "unresolved"
    elif unresolved:
        status = "partially_resolved"
    else:
        status = "resolved"

    return {
        "interface_id": iface.get("id"),
        "part_id": iface.get("part_id"),
        "semantic_role": iface.get("semantic_role"),
        "requested_ref_count": len(all_ids),
        "resolved_entity_count": len(found),
        "topology_entity_count": len(found),
        "face_count": face_count,
        "face_like": face_count > 0 and (world.get("normal") is not None),
        "unresolved_refs": unresolved,
        "resolution_status": status,
        "transform_applied": tok,
        "transform_note": tnote,
        "local": local,
        "world": world,
    }


def resolve_assembly_interfaces(assembly: Any,
                                topology_by_part: dict[str, dict[str, dict[str, Any]]] | None) -> dict[str, Any]:
    """Resolve every interface in the Assembly IR. ``topology_by_part`` maps part_id ->
    {entity_id: entity}. Missing topology degrades to unresolved (reported honestly)."""
    assembly = assembly if isinstance(assembly, dict) else {}
    topology_by_part = topology_by_part or {}
    transforms = {p.get("id"): p.get("transform") for p in _as_list(assembly.get("parts"))
                  if isinstance(p, dict)}
    interfaces = _collect_interfaces(assembly)

    records = {}
    for iid, iface in interfaces.items():
        pid = iface.get("part_id")
        records[iid] = resolve_interface(iface, topology_by_part.get(pid), transforms.get(pid))

    counts = {"resolved": 0, "partially_resolved": 0, "unresolved": 0}
    for r in records.values():
        counts[r["resolution_status"]] = counts.get(r["resolution_status"], 0) + 1

    return {
        "format": "aieng.assembly_interface_resolution", "format_version": FORMAT_VERSION,
        "schema_version": "0.1",
        "interfaces": records,
        "summary": {"interface_count": len(records), **counts,
                    "parts_with_topology": sorted(topology_by_part.keys())},
        "provenance": {"created_by": "aieng.assembly_interface_resolution",
                       "geometry_validation_only": True, "solver_executed": False},
    }


# ── connection geometry validation ────────────────────────────────────────────

def _status_from_reasons(reasons: list[str], insufficient: bool) -> str:
    if insufficient:
        return "insufficient_data"
    if any(r in {"far_apart", "missing_transform_blocks_validation"} for r in reasons):
        return "invalid"
    if reasons:
        return "warning"
    return "plausible"


def validate_connection_geometry(assembly: Any, resolution: dict[str, Any]) -> dict[str, Any]:
    """Validate each connection's geometry from resolved interface records."""
    assembly = assembly if isinstance(assembly, dict) else {}
    recs = (resolution or {}).get("interfaces") or {}
    out = []
    for idx, conn in enumerate(_as_list(assembly.get("connections"))):
        if not isinstance(conn, dict):
            continue
        cid = conn.get("id") or f"connection_{idx:03d}"
        ctype = conn.get("type")
        ia, ib = conn.get("interface_a"), conn.get("interface_b")
        ra, rb = recs.get(ia), recs.get(ib)
        reasons: list[str] = []
        metrics: dict[str, Any] = {}

        # need both interfaces resolved with world geometry to validate
        insufficient = False
        if ia is None or ib is None:
            insufficient = True
            reasons.append("unresolved_interface")
        elif ra is None or rb is None or ra["resolution_status"] == "unresolved" \
                or rb["resolution_status"] == "unresolved":
            insufficient = True
            reasons.append("unresolved_interface")
        else:
            wa, wb = ra.get("world") or {}, rb.get("world") or {}
            if not (wa.get("bbox") and wb.get("bbox")):
                insufficient = True
                reasons.append("unresolved_interface")
            else:
                if ra.get("transform_note") in {"missing_transform", "invalid_transform"} \
                        or rb.get("transform_note") in {"missing_transform", "invalid_transform"}:
                    reasons.append("missing_transform")
                ca, cb = wa["centroid"], wb["centroid"]
                centroid_distance = _dist(ca, cb)
                gap = _bbox_gap(wa["bbox"], wb["bbox"])
                overlap = _bbox_overlap(wa["bbox"], wb["bbox"])
                scale = max(_bbox_diag(wa["bbox"]), _bbox_diag(wb["bbox"]), 1e-6)
                touch_tol, near_tol = _TOUCH_FRAC * scale, _NEAR_FRAC * scale
                na, nb = wa.get("normal"), wb.get("normal")
                normal_alignment = None
                if na and nb:
                    normal_alignment = round(sum(na[k] * nb[k] for k in range(3)), 6)
                metrics = {
                    "centroid_distance": round(centroid_distance, 6),
                    "bbox_gap": round(gap, 6), "bbox_overlap": overlap,
                    "normal_alignment": normal_alignment, "scale": round(scale, 6),
                }

                far_apart = gap > near_tol
                touching = gap <= touch_tol or overlap
                if far_apart:
                    reasons.append("far_apart")
                elif not touching:
                    reasons.append("no_overlap")
                if normal_alignment is not None and normal_alignment > _NORMAL_SAME_DOT:
                    reasons.append("normals_same_direction")

                # type-specific evidence/behavior
                roles = {ra.get("semantic_role"), rb.get("semantic_role")}
                if ctype in _PREFERRED_ROLES and not (roles & _PREFERRED_ROLES[ctype]):
                    reasons.append("no_bolt_hole_evidence" if ctype == "bolted_proxy"
                                   else "no_weld_face_evidence")
                if ctype == "spring_proxy":
                    reasons.append("centroid_based_proxy")
                if ctype == "contact_proxy":
                    reasons.append("contact_draft_only")
                if ctype not in PROXY_CONNECTION_TYPES and ctype not in {"rigid_tie", "bonded"}:
                    reasons.append("unsupported_connection_type")

        status = _status_from_reasons(reasons, insufficient)
        # spring/contact proxies never claim better than "warning"
        if ctype in {"spring_proxy", "contact_proxy"} and status == "plausible":
            status = "warning"

        out.append({
            "connection_id": cid, "type": ctype,
            "part_a": conn.get("part_a"), "part_b": conn.get("part_b"),
            "interface_a": ia, "interface_b": ib,
            "geometry_status": status, "reasons": sorted(set(reasons)),
            "metrics": metrics, "is_proxy": ctype in PROXY_CONNECTION_TYPES,
        })

    tally: dict[str, int] = {}
    for c in out:
        tally[c["geometry_status"]] = tally.get(c["geometry_status"], 0) + 1
    return {
        "format": "aieng.assembly_connection_geometry", "format_version": FORMAT_VERSION,
        "schema_version": "0.1",
        "connections": out,
        "summary": {"connection_count": len(out), **tally},
        "limitations": [
            "Geometry validation only — no contact, friction, or bolt preload is modeled.",
            "Proximity/normal checks are approximate AABB+centroid heuristics, not exact surface mating.",
        ],
        "provenance": {"created_by": "aieng.assembly_interface_resolution",
                       "geometry_validation_only": True, "contact_physics_modeled": False,
                       "bolt_preload_modeled": False, "solver_executed": False},
    }


# ── package integration ──────────────────────────────────────────────────────

def _topo_index(topo_doc: Any) -> dict[str, dict[str, Any]]:
    ents = (topo_doc or {}).get("entities") if isinstance(topo_doc, dict) else None
    return {e["id"]: e for e in _as_list(ents) if isinstance(e, dict) and e.get("id")}


def build_topology_by_part(zf: zipfile.ZipFile, names: set[str],
                           assembly: dict[str, Any]) -> dict[str, dict[str, dict[str, Any]]]:
    """Build part_id -> {entity_id: entity}. Preference: explicit part topology_ref ->
    parts/<id>/topology_map.json -> geometry/parts/<id>/topology_map.json -> shared package map."""
    def _read(name: str) -> Any:
        try:
            return json.loads(zf.read(name)) if name in names else None
        except Exception:
            return None

    shared = _topo_index(_read(PACKAGE_TOPOLOGY_PATH))
    out: dict[str, dict[str, dict[str, Any]]] = {}
    for part in _as_list(assembly.get("parts")):
        if not isinstance(part, dict) or not part.get("id"):
            continue
        pid = part["id"]
        candidates = [part.get("topology_ref"), f"parts/{pid}/topology_map.json",
                      f"geometry/parts/{pid}/topology_map.json"]
        idx: dict[str, dict[str, Any]] = {}
        for cand in candidates:
            if cand and cand in names:
                idx = _topo_index(_read(cand))
                if idx:
                    break
        if not idx and shared:
            # shared single-package topology: filter by body_id == part id/name when present,
            # else expose the whole shared map (refs are globally unique).
            scoped = {eid: e for eid, e in shared.items()
                      if e.get("body_id") in {pid, part.get("name")}}
            idx = scoped or shared
        if idx:
            out[pid] = idx
    return out


def resolve_and_validate_assembly_geometry(package_path: str | Path) -> dict[str, Any]:
    """Best-effort: resolve interfaces + validate connection geometry for a package carrying
    assembly/assembly_ir.json. Writes interface_resolution, connection_geometry, refreshes the
    connection_graph + CAE draft with geometry_status. Returns {assembly_present: bool, ...}."""
    package_path = Path(package_path)
    if not package_path.exists():
        return {"assembly_present": False, "reason": "package not found"}
    try:
        with zipfile.ZipFile(package_path, "r") as zf:
            names = set(zf.namelist())
            if ASSEMBLY_IR_PATH not in names:
                return {"assembly_present": False}
            try:
                assembly = json.loads(zf.read(ASSEMBLY_IR_PATH))
            except Exception:
                assembly = None
            topo_by_part = build_topology_by_part(zf, names, assembly or {})
    except Exception as exc:  # noqa: BLE001
        return {"assembly_present": False, "error": f"{type(exc).__name__}: {exc}"}

    resolution = resolve_assembly_interfaces(assembly, topo_by_part)
    geometry = validate_connection_geometry(assembly, resolution)

    # refresh connection graph with geometry status per edge
    graph = build_connection_graph(assembly)
    geo_by_conn = {c["connection_id"]: c for c in geometry["connections"]}
    for edge in graph.get("edges", []):
        g = geo_by_conn.get(edge.get("id"))
        if g:
            edge["geometry_status"] = g["geometry_status"]
            edge["geometry_reasons"] = g["reasons"]
            edge["geometry_metrics"] = g["metrics"]
    graph["geometry_validation"] = {"available": True, "summary": geometry["summary"]}

    # refresh CAE draft with geometry status; invalid connections -> needs_user_input
    draft = build_assembly_cae_setup_draft(assembly)
    draft["geometry_validation_only"] = True
    draft.setdefault("provenance", {})["geometry_validation_only"] = True
    for cdraft in draft.get("connections", []):
        g = geo_by_conn.get(cdraft.get("connection_id"))
        if not g:
            continue
        cdraft["geometry_status"] = g["geometry_status"]
        cdraft["geometry_reasons"] = g["reasons"]
        if g["geometry_status"] == "invalid":
            cdraft["disabled"] = True
            msg = f"connection '{cdraft['connection_id']}' geometry invalid: {', '.join(g['reasons'])}"
            if msg not in draft.get("needs_user_input", []):
                draft.setdefault("needs_user_input", []).append(msg)
    # attach resolved interface geometry to loads/supports where available
    for group in ("loads", "supports"):
        for ref in draft.get(group, []):
            rec = resolution["interfaces"].get(ref.get("interface_id"))
            if rec and rec.get("world", {}).get("centroid") is not None:
                ref["world_centroid"] = rec["world"]["centroid"]
                ref["topology_resolved"] = rec["resolution_status"] != "unresolved"
    if draft.get("needs_user_input"):
        draft["status"] = "needs_user_input"

    def _dumps(o: Any) -> bytes:
        return (json.dumps(o, indent=2, sort_keys=True) + "\n").encode()

    members = {
        INTERFACE_RESOLUTION_PATH: _dumps(resolution),
        ASSEMBLY_CONNECTION_GEOMETRY_PATH: _dumps(geometry),
        CONNECTION_GRAPH_PATH: _dumps(graph),
        ASSEMBLY_CAE_DRAFT_PATH: _dumps(draft),
    }
    tmp = package_path.with_suffix(".tmp.aieng")
    try:
        with (
            zipfile.ZipFile(package_path, "r") as src,
            zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED) as dst,
        ):
            for item in src.infolist():
                if item.filename not in members:
                    dst.writestr(item, src.read(item.filename))
            for name, data in members.items():
                dst.writestr(name, data)
        tmp.replace(package_path)
    except Exception:
        tmp.unlink(missing_ok=True)
        raise

    cae_result: dict[str, Any] = {}
    try:
        from aieng.converters.assembly_cae import process_assembly_cae_package
        cae_result = process_assembly_cae_package(package_path)
    except Exception as exc:  # noqa: BLE001 - assembly CAE v0 is best-effort
        cae_result = {"error": f"{type(exc).__name__}: {exc}"}

    return {
        "assembly_present": True,
        "resolution_summary": resolution["summary"],
        "geometry_summary": geometry["summary"],
        "cae_draft_status": draft["status"],
        "assembly_cae_model_status": cae_result.get("assembly_cae_model_status"),
        "solver_deck_status": cae_result.get("solver_deck_status"),
        "solver_execution_status": cae_result.get("solver_execution_status"),
        "assembly_result_mapping_status": cae_result.get("assembly_result_mapping_status"),
        "assembly_topopt_status": cae_result.get("assembly_topopt_status"),
        "assembly_topopt_standard_problem_emitted": cae_result.get("assembly_topopt_standard_problem_emitted"),
        "artifacts": sorted(members.keys()),
    }
