snap2sim/model_io.py

"""Parsing helpers for model responses."""

from __future__ import annotations

import json
import re
from typing import Any

from snap2sim.schema import validate_analysis


_FENCE_RE = re.compile(r"^```(?:json|html)?\s*|\s*```$", re.IGNORECASE)
_THINK_RE = re.compile(r"<think\b[^>]*>.*?</think>", re.IGNORECASE | re.DOTALL)


def parse_analysis_response(text: str) -> dict[str, Any]:
    """Extract and validate a JSON object from a model response."""
    raw = _strip_fences(text)
    errors: list[str] = []
    for start, json_text in reversed(_json_object_candidates(raw)):
        try:
            payload = json.loads(json_text)
            return validate_analysis(payload)
        except (json.JSONDecodeError, ValueError) as exc:
            errors.append(f"object at {start}: {exc}")

    if errors:
        raise ValueError("Model response did not contain a valid analysis JSON object. " + errors[-1])
    raise ValueError("Model response did not contain a complete JSON object.")


def coerce_analysis_response(text: str) -> dict[str, Any]:
    """Best-effort conversion of partial model output into a valid analysis."""
    raw = _strip_fences(text)
    fallback_component = _infer_component(raw)
    for _, json_text in reversed(_json_object_candidates(raw)):
        try:
            payload = json.loads(json_text)
        except json.JSONDecodeError:
            continue
        if not isinstance(payload, dict):
            continue
        if not _looks_like_analysis_payload(payload):
            continue
        try:
            return validate_analysis(_coerce_analysis_payload(payload, fallback_component))
        except ValueError:
            continue
    return validate_analysis(_generic_analysis(fallback_component))


def _strip_fences(text: str) -> str:
    raw = _FENCE_RE.sub("", text.strip()).strip()
    return _strip_reasoning(raw)


def _strip_reasoning(text: str) -> str:
    cleaned = _THINK_RE.sub("", text)
    unclosed = re.search(r"<think\b[^>]*>", cleaned, re.IGNORECASE)
    if unclosed:
        after = cleaned[unclosed.end() :]
        json_start = after.find("{")
        if json_start >= 0:
            cleaned = cleaned[: unclosed.start()] + after[json_start:]
        else:
            cleaned = cleaned[: unclosed.start()]
    return cleaned.strip()


def _json_object_candidates(text: str) -> list[tuple[int, str]]:
    candidates = []
    for index, char in enumerate(text):
        if char != "{":
            continue
        try:
            candidates.append((index, _balanced_json_object(text, index)))
        except ValueError:
            continue
    return candidates


def _balanced_json_object(text: str, start: int) -> str:
    depth = 0
    in_string = False
    escaped = False
    for index in range(start, len(text)):
        char = text[index]
        if in_string:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
            continue

        if char == '"':
            in_string = True
        elif char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                return text[start : index + 1]

    raise ValueError("Model response contained an unterminated JSON object.")


def _looks_like_analysis_payload(payload: dict[str, Any]) -> bool:
    return any(key in payload for key in ["component", "summary", "trigger", "motion_sequence", "parts"])


def _coerce_analysis_payload(payload: dict[str, Any], fallback_component: str) -> dict[str, Any]:
    parts = payload.get("parts")
    if not isinstance(parts, list):
        parts = []
    coerced_parts = [_coerce_part(part, index) for index, part in enumerate(parts[:6]) if isinstance(part, dict)]
    coerced_parts = [part for part in coerced_parts if part is not None]
    if not coerced_parts:
        coerced_parts = _generic_analysis(fallback_component)["parts"]

    confidence = payload.get("confidence", 0.55)
    if not isinstance(confidence, (int, float)) or isinstance(confidence, bool):
        confidence = 0.55

    result = {
        "component": _non_empty_string(payload.get("component"), fallback_component),
        "confidence": max(0.0, min(1.0, float(confidence))),
        "summary": _non_empty_string(
            payload.get("summary"),
            f"{fallback_component.title()} approximated as primitive cutaway parts.",
        ),
        "trigger": _non_empty_string(payload.get("trigger"), "manual input"),
        "motion_sequence": _string_list(
            payload.get("motion_sequence"),
            ["input is applied", "internal parts move through the inferred mechanism"],
        ),
        "parts": coerced_parts,
    }
    if payload.get("render_mode") in {"three", "annotate", "unavailable"}:
        result["render_mode"] = payload["render_mode"]
    return result


def _coerce_part(part: dict[str, Any], index: int) -> dict[str, Any] | None:
    annotation = _coerce_annotation(part.get("annotation"))
    geometry = part.get("geometry")
    has_geometry_input = isinstance(geometry, dict) and bool(geometry)
    if not has_geometry_input:
        geometry = {}
    motion = part.get("motion")
    if not isinstance(motion, dict):
        motion = {}

    base_part: dict[str, Any] = {
        "id": _identifier(part.get("id"), f"part_{index + 1}"),
        "name": _non_empty_string(part.get("name"), f"part {index + 1}"),
        "role": _non_empty_string(part.get("role"), "inferred mechanical element"),
    }
    if annotation:
        base_part["annotation"] = annotation
    if annotation and not has_geometry_input:
        return base_part

    shape = geometry.get("shape")
    if shape not in {"box", "cylinder", "sphere", "gear", "rod", "cone", "capsule", "torus", "spring"}:
        shape = "box"
    motion_type = motion.get("type")
    if motion_type not in {"rotate", "translate", "oscillate", "static", "screw", "orbit", "pulse"}:
        motion_type = "static"

    coerced_motion: dict[str, Any] = {"type": motion_type}
    axis = _axis_vector(motion.get("axis"))
    if axis:
        coerced_motion["axis"] = axis
    for key in ["speed", "amplitude", "phase", "pitch"]:
        if isinstance(motion.get(key), (int, float)) and not isinstance(motion.get(key), bool):
            coerced_motion[key] = float(motion[key])
    values = _number_list(motion.get("range"), 2)
    if values:
        coerced_motion["range"] = values
    values = _number_list(motion.get("pivot"), 3)
    if values:
        coerced_motion["pivot"] = values

    coerced_geometry: dict[str, Any] = {
        "shape": shape,
        "size": _geometry_size(geometry, shape),
        "position": _number_list(geometry.get("position"), 3) or [float(index) - 1.0, 0.0, 0.0],
    }
    values = _number_list(geometry.get("rotation"), 3)
    if values:
        coerced_geometry["rotation"] = values
    if isinstance(geometry.get("teeth"), int):
        coerced_geometry["teeth"] = geometry["teeth"]
    if isinstance(geometry.get("coils"), int):
        coerced_geometry["coils"] = geometry["coils"]
    if isinstance(geometry.get("wire"), (int, float)) and not isinstance(geometry.get("wire"), bool):
        coerced_geometry["wire"] = float(geometry["wire"])
    if isinstance(geometry.get("color"), str) and geometry["color"].strip():
        coerced_geometry["color"] = geometry["color"].strip()

    base_part["geometry"] = coerced_geometry
    base_part["motion"] = coerced_motion
    return base_part


def _generic_analysis(component: str) -> dict[str, Any]:
    return {
        "component": component,
        "confidence": 0.45,
        "summary": f"{component.title()} rendered as a conservative generic cutaway.",
        "trigger": "manual input",
        "motion_sequence": [
            "input is applied to the housing",
            "central rotor transfers motion",
            "guide elements hold alignment",
        ],
        "parts": [
            {
                "id": "housing",
                "name": "outer housing",
                "role": "supports the internal mechanism",
                "geometry": {"shape": "box", "size": [2.4, 0.45, 1.4], "position": [0, 0, 0]},
                "motion": {"type": "static"},
            },
            {
                "id": "rotor",
                "name": "central rotor",
                "role": "transfers motion through the assembly",
                "geometry": {"shape": "cylinder", "size": [0.9, 0.35, 0.9], "position": [0, 0.18, 0]},
                "motion": {"type": "rotate", "speed": 0.55},
            },
            {
                "id": "guide",
                "name": "guide rail",
                "role": "keeps the moving part aligned",
                "geometry": {"shape": "rod", "size": [1.6, 0.12, 0.12], "position": [0, 0.42, 0.48]},
                "motion": {"type": "static"},
            },
        ],
    }


def _infer_component(text: str) -> str:
    match = re.search(r'"component"\s*:\s*"([^"]+)"', text)
    if match and match.group(1).strip():
        return match.group(1).strip()
    lowered = text.lower()
    for name in ["ratchet", "gear", "hinge", "motor", "lens", "target", "switch", "bearing"]:
        if name in lowered:
            return name
    return "observed component"


def _non_empty_string(value: Any, fallback: str) -> str:
    return value.strip() if isinstance(value, str) and value.strip() else fallback


def _string_list(value: Any, fallback: list[str]) -> list[str]:
    if isinstance(value, list):
        values = [item.strip() for item in value if isinstance(item, str) and item.strip()]
        if values:
            return values
    return fallback


def _number_list(value: Any, length: int) -> list[float] | None:
    if not isinstance(value, list) or len(value) != length:
        return None
    if not all(isinstance(item, (int, float)) and not isinstance(item, bool) for item in value):
        return None
    return [float(item) for item in value]


def _geometry_size(geometry: dict[str, Any], shape: str) -> list[float]:
    size = _number_list(geometry.get("size"), 3)
    if size:
        return size

    radius = _number_value(geometry.get("radius"))
    height = _number_value(geometry.get("height"))
    length = _number_value(geometry.get("length"))
    width = _number_value(geometry.get("width"))
    depth = _number_value(geometry.get("depth"))
    if shape in {"cylinder", "gear", "cone", "capsule", "torus", "spring"} and radius and height:
        diameter = radius * 2
        return [diameter, height, diameter]
    if shape == "rod" and (length or height) and radius:
        diameter = radius * 2
        return [diameter, diameter, length or height or 1.0]
    if width and height and depth:
        return [width, height, depth]
    return [1.0, 0.4, 1.0]


def _axis_vector(value: Any) -> list[float] | None:
    vector = _number_list(value, 3)
    if vector:
        return vector
    if isinstance(value, str):
        key = value.strip().lower()
        if key == "x":
            return [1.0, 0.0, 0.0]
        if key == "y":
            return [0.0, 1.0, 0.0]
        if key == "z":
            return [0.0, 0.0, 1.0]
    return None


def _number_value(value: Any) -> float | None:
    if isinstance(value, (int, float)) and not isinstance(value, bool):
        return float(value)
    return None


def _unit_number_list(value: Any, length: int) -> list[float] | None:
    values = _number_list(value, length)
    if not values:
        return None
    return [max(0.0, min(1.0, item)) for item in values]


def _coerce_annotation(value: Any) -> dict[str, Any] | None:
    if not isinstance(value, dict):
        return None
    point = _unit_number_list(value.get("point"), 2)
    if not point:
        return None
    annotation: dict[str, Any] = {"point": point}
    box = _unit_number_list(value.get("box"), 4)
    if box:
        annotation["box"] = box
    for key in ["label", "note"]:
        if isinstance(value.get(key), str) and value[key].strip():
            annotation[key] = value[key].strip()
    return annotation


def _identifier(value: Any, fallback: str) -> str:
    text = _non_empty_string(value, fallback).lower()
    text = re.sub(r"[^a-z0-9_]+", "_", text).strip("_")
    return text or fallback