#!/usr/bin/env python3
"""Extract a private, backend-neutral case record from an analytic STEP model."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any


def _rounded(value: float) -> float:
    return round(float(value), 6)


def _component(value: Any, name: str) -> float:
    component = getattr(value, name)
    return float(component() if callable(component) else component)


def _point(value: Any) -> list[float]:
    return [_rounded(_component(value, name)) for name in ("X", "Y", "Z")]


def _direction(value: Any) -> list[float]:
    return _point(value)


def _canonical_axis(direction: list[float]) -> str:
    index = max(range(3), key=lambda item: abs(direction[item]))
    return ("x", "y", "z")[index]


def _radial_offset(location: list[float], center: list[float], axis: str) -> float:
    axes = {"x": (1, 2), "y": (0, 2), "z": (0, 1)}
    first, second = axes[axis]
    return math.hypot(
        location[first] - center[first],
        location[second] - center[second],
    )


def _surface_record(face: Any, index: int) -> dict[str, Any]:
    from OCP.BRepAdaptor import BRepAdaptor_Surface
    from OCP.GeomAbs import (
        GeomAbs_Cone,
        GeomAbs_Cylinder,
        GeomAbs_Plane,
        GeomAbs_Sphere,
        GeomAbs_Torus,
    )

    adaptor = BRepAdaptor_Surface(face.wrapped)
    kind = adaptor.GetType()
    record: dict[str, Any] = {
        "surface_id": f"face_{index}",
        "area": _rounded(face.area),
        "center": _point(face.center()),
    }
    if kind == GeomAbs_Plane:
        plane = adaptor.Plane()
        record.update(
            {
                "type": "plane",
                "location": _point(plane.Location()),
                "axis": _direction(plane.Axis().Direction()),
            }
        )
    elif kind == GeomAbs_Cylinder:
        cylinder = adaptor.Cylinder()
        record.update(
            {
                "type": "cylinder",
                "location": _point(cylinder.Location()),
                "axis": _direction(cylinder.Axis().Direction()),
                "radius": _rounded(cylinder.Radius()),
            }
        )
    elif kind == GeomAbs_Cone:
        cone = adaptor.Cone()
        record.update(
            {
                "type": "cone",
                "location": _point(cone.Location()),
                "axis": _direction(cone.Axis().Direction()),
                "reference_radius": _rounded(cone.RefRadius()),
                "semi_angle_radians": _rounded(cone.SemiAngle()),
            }
        )
    elif kind == GeomAbs_Sphere:
        sphere = adaptor.Sphere()
        record.update(
            {
                "type": "sphere",
                "location": _point(sphere.Location()),
                "radius": _rounded(sphere.Radius()),
            }
        )
    elif kind == GeomAbs_Torus:
        torus = adaptor.Torus()
        record.update(
            {
                "type": "torus",
                "location": _point(torus.Location()),
                "axis": _direction(torus.Axis().Direction()),
                "major_radius": _rounded(torus.MajorRadius()),
                "minor_radius": _rounded(torus.MinorRadius()),
            }
        )
    else:
        record["type"] = "other"
    return record


def _semantic_summary(
    surfaces: list[dict[str, Any]], bbox_center: list[float], bbox_size: list[float]
) -> tuple[str, list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]]:
    cylinders = [item for item in surfaces if item["type"] == "cylinder"]
    cones = [item for item in surfaces if item["type"] == "cone"]
    axis_votes: Counter[str] = Counter()
    for surface in cylinders + cones:
        axis_votes[_canonical_axis(surface["axis"])] += max(surface["area"], 1.0)
    dominant_axis = axis_votes.most_common(1)[0][0] if axis_votes else "z"
    scale = max(bbox_size) if bbox_size else 1.0
    coaxial_tolerance = max(scale * 0.005, 1e-4)

    parallel_cylinders: list[dict[str, Any]] = []
    coaxial_cylinders: list[dict[str, Any]] = []
    off_axis_cylinders: list[dict[str, Any]] = []
    for surface in cylinders:
        if _canonical_axis(surface["axis"]) != dominant_axis:
            continue
        enriched = dict(surface)
        enriched["radial_offset"] = _rounded(
            _radial_offset(surface["location"], bbox_center, dominant_axis)
        )
        parallel_cylinders.append(enriched)
        if enriched["radial_offset"] <= coaxial_tolerance:
            coaxial_cylinders.append(enriched)
        else:
            off_axis_cylinders.append(enriched)

    distinct_coaxial_radii = {
        round(surface["radius"], 4) for surface in coaxial_cylinders
    }
    repeated_groups: dict[float, list[dict[str, Any]]] = defaultdict(list)
    for surface in off_axis_cylinders:
        repeated_groups[round(surface["radius"], 4)].append(surface)
    repeated_group_sizes = sorted(
        (len(rows) for rows in repeated_groups.values()), reverse=True
    )
    parallel_cones = [
        item for item in cones if _canonical_axis(item["axis"]) == dominant_axis
    ]
    has_rotational_stack = len(distinct_coaxial_radii) >= 2
    has_repeated_axial_holes = bool(repeated_group_sizes and repeated_group_sizes[0] >= 3)
    has_multi_axis_passages = any(
        _canonical_axis(item["axis"]) != dominant_axis for item in cylinders
    )

    features: list[dict[str, Any]] = []
    if has_rotational_stack:
        features.extend(
            [
                {"type": "rotational_body"},
                {"type": "coaxial_cylindrical_stack"},
            ]
        )
    if len(distinct_coaxial_radii) >= 3:
        features.append({"type": "central_passage_candidate"})
    if parallel_cones:
        features.append({"type": "conical_transition"})
    if has_repeated_axial_holes:
        features.append({"type": "repeated_axial_hole_pattern"})
    if has_multi_axis_passages:
        features.append({"type": "multi_axis_passage"})

    constraints: list[dict[str, Any]] = []
    if has_rotational_stack:
        constraints.extend(
            [
                {"id": "dominant_axis_alignment", "type": "axis_alignment"},
                {"id": "coaxial_stack", "type": "coaxial"},
            ]
        )
    if has_repeated_axial_holes:
        constraints.append(
            {"id": "repeated_radius_group", "type": "repeated_feature"}
        )
    if parallel_cones:
        constraints.append(
            {"id": "transition_axis_continuity", "type": "axis_alignment"}
        )

    if has_rotational_stack and parallel_cones and has_repeated_axial_holes:
        family = "flanged_hub_adapter_candidate"
    elif has_rotational_stack:
        family = "rotational_part_candidate"
    elif has_repeated_axial_holes:
        family = "patterned_prismatic_part_candidate"
    else:
        family = "unclassified_part"

    summary = {
        "dominant_axis": dominant_axis,
        "coaxial_cylinder_surface_count": len(coaxial_cylinders),
        "distinct_coaxial_radius_count": len(distinct_coaxial_radii),
        "off_axis_parallel_cylinder_surface_count": len(off_axis_cylinders),
        "largest_repeated_radius_group": repeated_group_sizes[0]
        if repeated_group_sizes
        else 0,
        "parallel_cone_surface_count": len(parallel_cones),
        "has_multi_axis_passages": has_multi_axis_passages,
    }
    return family, features, constraints, summary


def extract_step_case(source: Path) -> dict[str, Any]:
    from build123d import import_step

    resolved = source.expanduser().resolve()
    source_bytes = resolved.read_bytes()
    shape = import_step(str(resolved))
    bbox = shape.bounding_box()
    bbox_min = _point(bbox.min)
    bbox_max = _point(bbox.max)
    bbox_size = _point(bbox.size)
    bbox_center = [
        _rounded((lower + upper) / 2.0)
        for lower, upper in zip(bbox_min, bbox_max)
    ]
    surfaces = [
        _surface_record(face, index) for index, face in enumerate(shape.faces())
    ]
    surface_counts = Counter(surface["type"] for surface in surfaces)
    family, features, constraints, semantic_summary = _semantic_summary(
        surfaces, bbox_center, bbox_size
    )
    digest = hashlib.sha256(source_bytes).hexdigest()
    return {
        "schema_version": "1.0",
        "case_kind": "private_step_evidence",
        "case_id": digest,
        "provenance": {
            "source_path": str(resolved),
            "source_sha256": digest,
            "source_format": "step",
        },
        "geometry_evidence": {
            "valid": bool(shape.is_valid),
            "solid_count": len(shape.solids()),
            "face_count": len(shape.faces()),
            "volume": _rounded(shape.volume),
            "bounding_box": {
                "min": bbox_min,
                "max": bbox_max,
                "size": bbox_size,
                "center": bbox_center,
            },
            "analytic_surface_counts": dict(sorted(surface_counts.items())),
            "analytic_surfaces": surfaces,
        },
        "design_ir": {
            "part_family": family,
            "features": features,
            "constraints": constraints,
            "normalized_observations": [],
            "semantic_summary": semantic_summary,
        },
        "experience": {
            "rules": [],
            "validation_targets": [],
            "promotion_state": "evidence_only",
        },
    }


def write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Extract a private case JSON from a STEP file."
    )
    parser.add_argument("source", type=Path)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args(argv)
    write_json(args.output, extract_step_case(args.source))
    print(args.output)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())