#!/usr/bin/env python3
"""Execute one uploaded STEP -> DesignIR 3.0 reconstruction task.

The teacher STEP is passed only to deterministic extraction and acceptance
commands.  The independent rebuild subprocess receives the DesignIR path and
output path only.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any


SCRIPT_PATH = Path(__file__).resolve()
SKILL_ROOT = SCRIPT_PATH.parents[1]
WORKSPACE_ROOT = SKILL_ROOT.parents[2]
SURFACEIR_PIPELINE = (
    WORKSPACE_ROOT / "designir-pipeline" / "scripts" / "surfaceir_pipeline.py"
)
EXPERIENCE_LIBRARY = (
    WORKSPACE_ROOT / "designir-pipeline" / "output" / "library.json"
)
DESIGNIR_SCHEMA = (
    WORKSPACE_ROOT
    / "designir-pipeline"
    / "contracts"
    / "designir-3.0.schema.json"
)

EXECUTABLE_EDIT_STATES = {
    "executable_enlarge_cut_binding",
    "executable_axis_affine_binding",
    "executable_uniform_scale_binding",
    "validated_executable_binding",
}
FORBIDDEN_DESIGNIR_TOKENS = (
    '"source_path"',
    '"teacher_step"',
    "teacher.step",
    "import_step",
    "stepcontrol_reader",
    '"brep"',
    '"mesh"',
)


class ReconstructionError(RuntimeError):
    """Raised when a reconstruction stage cannot produce an accepted result."""


def _runtime_python() -> Path:
    configured = os.environ.get("CAD_ROUTER_PYTHON", "").strip()
    candidates = [
        Path(configured).expanduser() if configured else None,
        WORKSPACE_ROOT / "text-to-cad" / ".venv" / "bin" / "python",
        Path(sys.executable),
    ]
    for candidate in candidates:
        if candidate is not None and candidate.is_file():
            # Preserve a virtualenv launcher symlink. Resolving it to the base
            # interpreter discards the virtualenv's package search path.
            return candidate.absolute()
    raise ReconstructionError(
        "No CAD runtime Python found; set CAD_ROUTER_PYTHON or create "
        "text-to-cad/.venv."
    )


def _safe_stem(source: Path) -> str:
    stem = re.sub(r"[^a-zA-Z0-9_.-]+", "_", source.stem).strip("._")
    return stem or "uploaded_part"


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


def _run_json(
    command: list[str],
    *,
    timeout_seconds: int,
    stage: str,
) -> dict[str, Any]:
    completed = subprocess.run(
        command,
        capture_output=True,
        text=True,
        timeout=timeout_seconds,
        check=False,
    )
    if completed.returncode != 0:
        detail = completed.stderr.strip() or completed.stdout.strip()
        raise ReconstructionError(
            f"{stage} failed with exit code {completed.returncode}: "
            f"{detail[-2000:]}"
        )
    lines = [line for line in completed.stdout.splitlines() if line.strip()]
    if not lines:
        return {}
    try:
        return json.loads(lines[-1])
    except json.JSONDecodeError as exc:
        raise ReconstructionError(
            f"{stage} returned non-JSON output: {lines[-1][-1000:]}"
        ) from exc


def _validate_designir_contract(payload: dict[str, Any]) -> dict[str, Any]:
    required = {
        "schema_version",
        "designir_kind",
        "model_id",
        "units",
        "document_status",
        "reconstruction_mode",
        "semantic_layer",
        "surface_layer",
        "edit_interface",
        "validation_contract",
    }
    missing = sorted(required - set(payload))
    if missing:
        raise ReconstructionError(
            f"DesignIR 3.0 is missing required fields: {', '.join(missing)}"
        )
    if payload["schema_version"] != "3.0":
        raise ReconstructionError("Uploaded STEP did not produce DesignIR 3.0")
    if payload["designir_kind"] != "independent_parametric_cad":
        raise ReconstructionError("DesignIR is not source-independent")
    if payload["validation_contract"].get("source_independence") is not True:
        raise ReconstructionError("DesignIR source-independence contract is false")
    serialized = json.dumps(payload, ensure_ascii=False).lower()
    violations = [
        token for token in FORBIDDEN_DESIGNIR_TOKENS if token in serialized
    ]
    if violations:
        raise ReconstructionError(
            "DesignIR anti-cheat scan failed: " + ", ".join(violations)
        )
    return {
        "schema_contract": "designir-3.0",
        "schema_path": str(DESIGNIR_SCHEMA),
        "required_fields_present": True,
        "source_independence_declared": True,
        "anti_cheat_pass": True,
        "forbidden_token_matches": 0,
    }


def _candidate_parameters(
    payload: dict[str, Any], maximum: int
) -> list[tuple[str, float]]:
    parameters = payload.get("semantic_layer", {}).get("parameters", {})
    names = payload.get("edit_interface", {}).get("semantic_parameters", [])
    feature_candidates: list[tuple[str, float]] = []
    fallback: tuple[str, float] | None = None
    for name in names:
        parameter = parameters.get(name)
        if not isinstance(parameter, dict):
            continue
        if parameter.get("edit_state") not in EXECUTABLE_EDIT_STATES:
            continue
        try:
            old_value = float(parameter["value"])
        except (KeyError, TypeError, ValueError):
            continue
        candidate = (name, old_value * 1.1)
        if name == "overall_scale":
            fallback = candidate
        else:
            feature_candidates.append(candidate)
    selected = feature_candidates[: max(0, maximum)]
    if fallback is not None:
        selected.append(fallback)
    return selected


def _surfaceir_experience_ids() -> set[str]:
    if not EXPERIENCE_LIBRARY.is_file():
        return set()
    try:
        payload = json.loads(EXPERIENCE_LIBRARY.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return set()
    policy = payload.get("policy", {})
    if (
        payload.get("library_kind") != "generalized_cad_experience"
        or policy.get("surfaceir_consumable") is not True
        or policy.get("instance_parameters_allowed") is not False
    ):
        return set()
    return {
        item["id"]
        for item in payload.get("experiences", [])
        if isinstance(item, dict)
        and isinstance(item.get("id"), str)
        and set(item.get("scope", [])).intersection(
            {"surfaceir_reconstruction", "surfaceir_semantic_editing"}
        )
    }


def _applied_experience_ids(
    available: set[str], accepted_parameters: list[dict[str, Any]]
) -> list[str]:
    applied = {
        experience_id
        for experience_id in (
            "failure_repair.surface_topology.bind_face_pcurves_and_group_shells",
            "failure_repair.surface_topology.preserve_vertex_tolerance",
        )
        if experience_id in available
    }
    names = {parameter["name"] for parameter in accepted_parameters}
    if any("hole_diameter" in name or "bore_diameter" in name for name in names):
        applied.update(
            experience_id
            for experience_id in (
                "edit_strategy.semantic_cylindrical_cut.enlarge_diameter",
                "failure_repair.semantic_edit.merge_duplicate_cylindrical_targets",
                "edit_strategy.semantic_edit.preserve_validated_geometry_strategy",
            )
            if experience_id in available
        )
    if "overall_scale" in names:
        experience_id = (
            "edit_strategy.semantic_parameter.source_independent_overall_scale"
        )
        if experience_id in available:
            applied.add(experience_id)
    if names.intersection(
        {"outer_diameter", "body_length", "plate_thickness"}
    ):
        experience_id = (
            "edit_strategy.semantic_parameter.primary_axis_affine_dimensions"
        )
        if experience_id in available:
            applied.add(experience_id)
    if any(name.startswith("overall_size_") for name in names):
        experience_id = (
            "edit_strategy.semantic_parameter.canonical_directional_envelope"
        )
        if experience_id in available:
            applied.add(experience_id)
    return sorted(applied)


def _persist_edit_results(
    designir_path: Path,
    attempts: list[dict[str, Any]],
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    payload = json.loads(designir_path.read_text(encoding="utf-8"))
    parameters = payload.get("semantic_layer", {}).get("parameters", {})
    accepted_parameters: list[dict[str, Any]] = []
    attempted_names = set()
    for attempt in attempts:
        name = attempt["parameter"]
        attempted_names.add(name)
        parameter = parameters[name]
        old_value = float(parameter["value"])
        requested_value = float(attempt["requested_value"])
        if attempt["accepted"]:
            parameter["editable"] = True
            parameter["edit_state"] = "validated_executable_binding"
            parameter["validated_range"] = {
                "direction": "increase",
                "minimum_tested_inclusive": old_value,
                "maximum_tested_inclusive": requested_value,
                "acceptance_contract": "designir_3_semantic_edit_acceptance",
            }
            accepted_parameters.append(
                {
                    "name": name,
                    "semantic_role": parameter.get("semantic_role", name),
                    "value": parameter["value"],
                    "unit": parameter.get("unit", "mm"),
                    "validated_range": parameter["validated_range"],
                    "confidence": parameter.get("confidence"),
                }
            )
        else:
            parameter["editable"] = False
            parameter["edit_state"] = "rejected_by_runtime_validation"
            parameter.pop("validated_range", None)
    payload["edit_interface"]["semantic_parameters"] = sorted(
        parameter["name"] for parameter in accepted_parameters
    )
    payload["reconstruction_strategy"] = {
        "boundary_strategy": "exact_3d_pcurve",
        "selection_status": "geometry_validated",
    }
    _write_json(designir_path, payload)
    return payload, accepted_parameters


def reconstruct_uploaded_step(
    source: Path,
    task_dir: Path,
    *,
    request: str,
    maximum_parameters: int = 8,
    timeout_seconds: int = 180,
) -> dict[str, Any]:
    source = source.expanduser().resolve()
    task_dir = task_dir.expanduser().resolve()
    if not source.is_file():
        raise ReconstructionError(f"Uploaded STEP does not exist: {source}")
    if source.suffix.lower() not in {".step", ".stp"}:
        raise ReconstructionError("Uploaded reconstruction source must be STEP/STP")
    if source == task_dir or task_dir == source.parent:
        raise ReconstructionError(
            "Task directory must be separate from the uploaded STEP directory"
        )
    if source == task_dir or task_dir in source.parents:
        raise ReconstructionError(
            "Task directory cannot contain the uploaded teacher STEP"
        )
    if not SURFACEIR_PIPELINE.is_file():
        raise ReconstructionError(
            f"DesignIR 3.0 runtime is missing: {SURFACEIR_PIPELINE}"
        )

    task_dir.mkdir(parents=True, exist_ok=True)
    stem = _safe_stem(source)
    designir_path = task_dir / f"{stem}.designir.json"
    rebuilt_path = task_dir / f"{stem}.reconstructed.step"
    validation_dir = task_dir / "validation"
    geometry_report_path = validation_dir / "geometry.json"
    edit_report_path = validation_dir / "edits.json"
    parameters_path = task_dir / "parameters.json"
    task_manifest_path = task_dir / "cad-task.json"
    collisions = [
        path
        for path in (
            designir_path,
            rebuilt_path,
            parameters_path,
            task_manifest_path,
        )
        if path.exists()
    ]
    if collisions:
        raise ReconstructionError(
            "Refusing to overwrite existing task artifacts: "
            + ", ".join(str(path) for path in collisions)
        )

    runtime = _runtime_python()
    common = [str(runtime), str(SURFACEIR_PIPELINE)]
    available_experience_ids = _surfaceir_experience_ids()
    boundary_strategy = (
        "exact_3d_pcurve"
        if (
            "failure_repair.surface_topology.bind_face_pcurves_and_group_shells"
            in available_experience_ids
        )
        else "exact_3d"
    )

    extraction = _run_json(
        common
        + [
            "extract",
            str(source),
            "--output",
            str(designir_path),
        ],
        timeout_seconds=timeout_seconds,
        stage="DesignIR 3.0 extraction",
    )
    payload = json.loads(designir_path.read_text(encoding="utf-8"))
    contract_checks = _validate_designir_contract(payload)

    # The rebuild worker receives no teacher path.  It can regenerate the model
    # from the DesignIR artifact alone.
    _run_json(
        common
        + [
            "rebuild",
            str(designir_path),
            "--output",
            str(rebuilt_path),
            "--boundary-strategy",
            boundary_strategy,
        ],
        timeout_seconds=timeout_seconds,
        stage="source-independent SurfaceIR rebuild",
    )

    manifest_path = task_dir / "manifest.json"
    _write_json(
        manifest_path,
        {
            "schema_version": "1.0",
            "input_count": 1,
            "rows": [
                {
                    "source_name": source.name,
                    "output_name": designir_path.name,
                    "state": "extracted",
                }
            ],
        },
    )
    validation_dir.mkdir(parents=True, exist_ok=True)
    geometry_work_dir = validation_dir / ".geometry-work"
    geometry_result = _run_json(
        common
        + [
            "validate-folder",
            str(source.parent),
            "--designir-dir",
            str(task_dir),
            "--rebuilt-dir",
            str(geometry_work_dir),
            "--limit",
            "1",
            "--boolean",
            "--boundary-strategy",
            boundary_strategy,
        ],
        timeout_seconds=timeout_seconds,
        stage="Acceptance-Agent geometry validation",
    )
    raw_geometry_report = geometry_work_dir / "validation-report.json"
    geometry_report = json.loads(raw_geometry_report.read_text(encoding="utf-8"))
    _write_json(geometry_report_path, geometry_report)
    geometry_pass = geometry_report.get("geometry_pass_count") == 1
    if not geometry_pass:
        raise ReconstructionError(
            "Independent reconstruction did not satisfy the geometry contract; "
            f"see {geometry_report_path}"
        )
    shutil.rmtree(geometry_work_dir)

    attempts: list[dict[str, Any]] = []
    if payload["document_status"] == "geometry_present":
        for name, requested_value in _candidate_parameters(
            payload, maximum_parameters
        ):
            safe_name = _safe_stem(Path(name))
            attempt_dir = validation_dir / "parameter-tests" / safe_name
            try:
                result = _run_json(
                    common
                    + [
                        "validate-edit",
                        str(source),
                        str(designir_path),
                        "--parameter",
                        name,
                        "--value",
                        str(requested_value),
                        "--output-dir",
                        str(attempt_dir),
                    ],
                    timeout_seconds=timeout_seconds,
                    stage=f"parameter validation ({name})",
                )
                report = json.loads(
                    (attempt_dir / "acceptance-report.json").read_text(
                        encoding="utf-8"
                    )
                )
                attempts.append(
                    {
                        "parameter": name,
                        "requested_value": requested_value,
                        "accepted": bool(report.get("accepted")),
                        "state": report.get("state"),
                        "report": str(
                            attempt_dir / "acceptance-report.json"
                        ),
                        "worker": result,
                    }
                )
            except (ReconstructionError, subprocess.TimeoutExpired) as exc:
                attempts.append(
                    {
                        "parameter": name,
                        "requested_value": requested_value,
                        "accepted": False,
                        "state": "worker_failed",
                        "error": str(exc),
                    }
                )

    payload, accepted_parameters = _persist_edit_results(
        designir_path, attempts
    )
    payload["reconstruction_strategy"] = {
        "boundary_strategy": boundary_strategy,
        "selection_status": "geometry_validated",
    }
    _write_json(designir_path, payload)
    applied_experience_ids = _applied_experience_ids(
        available_experience_ids, accepted_parameters
    )
    edit_report = {
        "schema_version": "1.0",
        "report_kind": "cad_router_runtime_parameter_acceptance",
        "attempt_count": len(attempts),
        "accepted_count": sum(row["accepted"] for row in attempts),
        "attempts": attempts,
    }
    _write_json(edit_report_path, edit_report)
    parameter_catalog = {
        "schema_version": "1.0",
        "model_id": payload["model_id"],
        "document_status": payload["document_status"],
        "parameter_count": len(accepted_parameters),
        "parameters": accepted_parameters,
        "policy": "Only parameters that passed an isolated perturbation test are exposed.",
    }
    _write_json(parameters_path, parameter_catalog)

    task_manifest = {
        "schema_version": "1.0",
        "request": request,
        "operation": "uploaded_step_reconstruction",
        "route": {
            "router": "cad-router",
            "runtime": "DesignIR 3.0 / SurfaceIR",
            "selected_backend": "build123d",
            "execution_adapter": "surfaceir_occt",
            "source_of_truth": designir_path.name,
        },
        "source": {
            "format": source.suffix.lower().lstrip("."),
            "role": "teacher_and_acceptance_truth_only",
            "embedded_in_output": False,
        },
        "artifacts": [
            {"path": rebuilt_path.name, "role": "primary_step"},
            {"path": designir_path.name, "role": "editable_designir"},
            {"path": parameters_path.name, "role": "validated_parameters"},
            {
                "path": str(geometry_report_path.relative_to(task_dir)),
                "role": "geometry_acceptance",
            },
            {
                "path": str(edit_report_path.relative_to(task_dir)),
                "role": "parameter_acceptance",
            },
        ],
        "experience": {
            "library": (
                str(EXPERIENCE_LIBRARY)
                if EXPERIENCE_LIBRARY.is_file()
                else None
            ),
            "runtime_policy": "promoted_generalized_experience_only",
            "available_method_count": len(available_experience_ids),
            "methods_applied": applied_experience_ids,
        },
        "validation": {
            "geometry_accepted": True,
            "source_independent_rebuild": True,
            "anti_cheat_pass": True,
            "validated_parameter_count": len(accepted_parameters),
        },
        "edit_command_template": [
            str(runtime),
            str(SURFACEIR_PIPELINE),
            "edit",
            designir_path.name,
            "--parameter",
            "<validated-parameter-name>",
            "--value",
            "<value-within-validated-range>",
            "--output-designir",
            "<edited.designir.json>",
            "--output-step",
            "<edited.step>",
        ],
    }
    _write_json(task_manifest_path, task_manifest)
    manifest_path.unlink(missing_ok=True)

    return {
        "schema_version": "1.0",
        "status": (
            "accepted"
            if payload["document_status"] == "geometry_present"
            else "accepted_empty_source"
        ),
        "task_dir": str(task_dir),
        "designir": str(designir_path),
        "rebuilt_step": str(rebuilt_path),
        "parameters": str(parameters_path),
        "task_manifest": str(task_manifest_path),
        "geometry_report": str(geometry_report_path),
        "edit_report": str(edit_report_path),
        "geometry_accepted": True,
        "document_status": payload["document_status"],
        "validated_parameter_count": len(accepted_parameters),
        "validated_parameter_names": [
            parameter["name"] for parameter in accepted_parameters
        ],
        "experience": {
            "available_method_count": len(available_experience_ids),
            "methods_applied": applied_experience_ids,
        },
        "source_independence": {
            "rebuild_command_received_teacher_path": False,
            **contract_checks,
        },
        "stages": {
            "extraction": extraction,
            "geometry_validation": geometry_result,
        },
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--task-dir", type=Path, required=True)
    parser.add_argument("--request", default="重建上传的 STEP")
    parser.add_argument("--maximum-parameters", type=int, default=8)
    parser.add_argument("--timeout-seconds", type=int, default=180)
    return parser


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    try:
        result = reconstruct_uploaded_step(
            args.source,
            args.task_dir,
            request=args.request,
            maximum_parameters=args.maximum_parameters,
            timeout_seconds=args.timeout_seconds,
        )
    except (ReconstructionError, subprocess.TimeoutExpired) as exc:
        print(
            json.dumps(
                {"status": "failed", "error": str(exc)},
                ensure_ascii=False,
            ),
            file=sys.stderr,
        )
        return 1
    print(json.dumps(result, ensure_ascii=False))
    return 0


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