tools/functions.py

"""
Agent tools (OpenAI Agents SDK).

Thin wrappers only: parse arguments, call into `helpers/`, catch exceptions and
return a JSON-serialisable dict. All real logic lives in `helpers/`.

Every tool returns a dict with a "status" key and never raises into the runner —
a raised exception aborts the turn and the model gets nothing it can act on.

`@function_tool` builds the JSON schema from the type hints and the docstring, so
the docstring IS the tool description the model reads. The leading
`RunContextWrapper` parameter is stripped from that schema automatically: it is
host state, not something the model fills in.

Keypoint selection is the AGENT'S job. No tool here picks a point. The tools give
the agent a ruler (`helpers/canvas.py`), then measure what it chose
(`helpers/validate.py`) and hand the measurements back so it can revise.
"""

from __future__ import annotations

import json
import traceback
from pathlib import Path
from typing import Any

from agents import RunContextWrapper, function_tool

from config.settings import MAX_KEYPOINT_ROUNDS
from helpers import background, meshy_client, sam3, specimen, validate, viz
from helpers.images import specimen_dir
from schemas.context import FossilContext

# segment_fossil_sam3_tool writes exactly this suffix (see helpers/sam3.py).
SAM3_CUTOUT_SUFFIX = "_sam3_cutout.png"


def _fail(step: str, exc: Exception) -> dict[str, Any]:
    traceback.print_exc()
    return {"status": "error", "step": step, "message": f"{type(exc).__name__}: {exc}"}


def _parse_points(raw: str | None, label: str) -> list[dict]:
    """Accept a JSON array of [x, y] pairs or of {"x":..,"y":..} objects."""
    if not raw or not raw.strip():
        return []
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise ValueError(f"{label} must be a JSON array, got {raw!r}") from exc

    points = []
    for item in data:
        if isinstance(item, dict):
            x, y = item["x"], item["y"]
        else:
            x, y = item[0], item[1]
        points.append({"x": int(x), "y": int(y), "label": label, "source": "agent"})
    return points


# ---------------------------------------------------------------------------
# STEP 0 — the specimen: a folder of views
# ---------------------------------------------------------------------------

@function_tool
def inspect_specimen_tool(
    ctx: RunContextWrapper[FossilContext],
    folder: str,
) -> dict[str, Any]:
    """List the views in a specimen folder and sanity-check that they show the same thing.

    A specimen is a FOLDER of photographs (different angles of one fossil), not a
    single image. Meshy's multi-image endpoint accepts 1 to 4 views; if this
    returns needs_selection, you must call select_views_tool to choose which 4.

    The consistency check is advisory and weak by design — it catches a photo of a
    different fossil dropped into the folder, but two similar specimens on the same
    desk will pass it. Read `notes` and pass any suspicion on to the analyst.

    Args:
        folder: Path to the specimen folder.
    """
    print(f"\n[Tool] inspect_specimen_tool({folder})")
    try:
        views = specimen.discover_views(folder)
        sid = specimen.specimen_id(folder)
        ctx.context.specimen_dir = str(folder)
        ctx.context.specimen_id = sid

        consistency = specimen.check_views_consistent(views)
        n = len(views)
        return {
            "status": "success",
            "specimen_id": sid,
            "folder": str(folder),
            "views": views,
            "n_views": n,
            "max_views_meshy_accepts": specimen.MAX_VIEWS,
            "needs_selection": n > specimen.MAX_VIEWS,
            "consistency": consistency,
            "next": (
                f"{n} views — more than Meshy's limit of {specimen.MAX_VIEWS}. Call "
                f"select_views_tool to let the analyst choose."
                if n > specimen.MAX_VIEWS
                else f"{n} view(s). Process each through steps 1-2, then step 3."
            ),
        }
    except Exception as exc:  # noqa: BLE001
        return _fail("inspect_specimen", exc)


@function_tool
async def select_views_tool(
    ctx: RunContextWrapper[FossilContext],
    folder: str,
) -> dict[str, Any]:
    """Ask the Vision Analyst to choose which views to reconstruct from, when a folder holds more than Meshy's limit of 4.

    Nothing here ranks or filters the views. The analyst looks at a contact sheet
    and decides — "which four views best cover this specimen" is a judgement about
    the object (the broken face is worth more than a third near-duplicate of the
    dorsal surface), not an arithmetic fact about the file listing.

    Args:
        folder: Path to the specimen folder.
    """
    print(f"\n[Tool] select_views_tool({folder})")
    try:
        from fossil_agents.vision_analyst.analyst_agent import run_view_selection

        return await run_view_selection(folder, ctx.context)
    except Exception as exc:  # noqa: BLE001
        return _fail("view_selection", exc)


# ---------------------------------------------------------------------------
# STEP 1 — background removal
# ---------------------------------------------------------------------------

@function_tool
def remove_background_tool(
    ctx: RunContextWrapper[FossilContext],
    image_path: str,
) -> dict[str, Any]:
    """Remove the background from one view of a fossil specimen.

    Produces a transparent cutout, a white-composited cutout, and an 8-bit alpha
    matte. The matte is a HINT for later steps, not a decision — the agent that
    picks the keypoints is free to disagree with it, and often should.

    Args:
        image_path: Path to the original photograph of this view.
    """
    print(f"\n[Tool] remove_background_tool({image_path})")
    try:
        result = background.execute_background_removal(image_path)
        if result.get("alpha_mask_path"):
            ctx.context.alpha_masks[image_path] = result["alpha_mask_path"]
        return result
    except Exception as exc:  # noqa: BLE001
        return _fail("background_removal", exc)


# ---------------------------------------------------------------------------
# STEP 2 — hand off to the vision analyst, which CHOOSES the keypoints
# ---------------------------------------------------------------------------

@function_tool
async def select_keypoints_tool(
    ctx: RunContextWrapper[FossilContext],
    image_path: str,
    alpha_mask_path: str = "",
) -> dict[str, Any]:
    """Hand one view to the Vision Analyst, which looks at it, CHOOSES the SAM3 prompt points itself, checks its own work, and runs the segmentation.

    Nothing here computes a keypoint. The analyst decides where the points go. Call
    this once per view, after background removal. You cannot see the image, so you
    have no basis for suggesting coordinates — do not.

    Args:
        image_path: Path to the ORIGINAL photograph of this view.
        alpha_mask_path: Alpha matte from remove_background_tool. Shown to the
            analyst as a boundary hint it may overrule.
    """
    print(f"\n[Tool] select_keypoints_tool({image_path})")
    try:
        from fossil_agents.vision_analyst.analyst_agent import run_selection

        return await run_selection(
            image_path, alpha_mask_path or None, ctx.context)
    except Exception as exc:  # noqa: BLE001
        return _fail("keypoint_selection", exc)


# ---------------------------------------------------------------------------
# STEP 2b — the analyst's own tools
# ---------------------------------------------------------------------------

@function_tool
def check_keypoints_tool(
    ctx: RunContextWrapper[FossilContext],
    image_path: str,
    positive_points: str,
    negative_points: str = "",
    alpha_mask_path: str = "",
) -> dict[str, Any]:
    """Check the points YOU chose, before committing them to SAM3.

    This does not pick or move any point. It measures the ones you gave it and
    reports back: whether each is in bounds, what is underneath it (specimen,
    backdrop, or cast shadow), how far it sits from the nearest edge, and whether
    you left a cast shadow unmarked. It also renders a preview of where your points
    actually landed.

    Warnings are advice. You may overrule them — you can see the photograph and the
    background remover cannot. Errors are blocking: fix them and call this again.

    Args:
        image_path: Path to the ORIGINAL photograph.
        positive_points: JSON array of the points YOU chose ON the fossil,
            e.g. "[[417,779],[300,465]]".
        negative_points: JSON array of the points YOU chose OFF the fossil —
            backdrop, cast shadow, mount, scale bar.
        alpha_mask_path: Alpha matte from step 1, if you were given one.
    """
    print(f"\n[Tool] check_keypoints_tool({image_path})")
    try:
        pos = _parse_points(positive_points, "positive")
        neg = _parse_points(negative_points, "negative")

        report = validate.validate_points(image_path, alpha_mask_path or None, pos, neg)

        c = ctx.context
        c.keypoint_round += 1
        n = c.keypoint_round

        stem = Path(image_path).stem
        preview = specimen_dir(image_path) / f"{stem}_keypoints_round{n}.png"
        report["preview_path"] = viz.render_keypoint_preview(
            image_path, alpha_mask_path or None, pos, neg, preview)
        report["round"] = n
        report["rounds_remaining"] = max(0, MAX_KEYPOINT_ROUNDS - n)
        report["status"] = "success"

        if report["ok"] and not report["warnings"]:
            report["next"] = "Points look sound. Call segment_fossil_sam3_tool with this set."
        elif report["ok"]:
            report["next"] = (
                "No blocking errors. Weigh the warnings, revise if you agree with them, "
                "then call segment_fossil_sam3_tool with your final set.")
        else:
            report["next"] = "Blocking errors. Fix the points and call this tool again."

        if report["rounds_remaining"] <= 0:
            report["next"] = (
                "You have used your revision budget. Commit your best point set to "
                "segment_fossil_sam3_tool now, or stop and explain why you cannot.")

        print(f"  [Check r{n}] {report['summary']}")
        for msg in report["errors"] + report["warnings"]:
            print(f"    - {msg}")
        return report
    except Exception as exc:  # noqa: BLE001
        return _fail("keypoint_check", exc)


# ---------------------------------------------------------------------------
# STEP 3 — SAM3
# ---------------------------------------------------------------------------

@function_tool
def segment_fossil_sam3_tool(
    ctx: RunContextWrapper[FossilContext],
    image_path: str,
    positive_points: str,
    negative_points: str = "",
) -> dict[str, Any]:
    """Segment the fossil with SAM3 using the point prompts YOU selected.

    Run this on the ORIGINAL photograph, not the background-removed image — SAM3
    needs the real texture gradients to find a crisp boundary.

    Call check_keypoints_tool first. Call this ONCE, with your final point set.

    Args:
        image_path: Path to the ORIGINAL photograph.
        positive_points: JSON array of points INSIDE the fossil,
            e.g. "[[417,779],[300,465]]".
        negative_points: JSON array of points in the backdrop, cast shadow, or
            adhering matrix. Strongly recommended — this is what keeps the shadow out.
    """
    print(f"\n[Tool] segment_fossil_sam3_tool({image_path})")
    try:
        pos = _parse_points(positive_points, "positive")
        neg = _parse_points(negative_points, "negative")
        if not pos:
            return {"status": "error", "message": "positive_points is empty."}

        result = sam3.execute_sam3_segmentation(image_path, pos, neg)

        if result.get("status") == "success" and result.get("cutout_path"):
            c = ctx.context
            c.cutouts[image_path] = result["cutout_path"]
            c.keypoints[image_path] = {
                "positive": pos, "negative": neg,
                "coverage": result.get("coverage"),
            }
        return result
    except Exception as exc:  # noqa: BLE001
        return _fail("sam3_segmentation", exc)


# ---------------------------------------------------------------------------
# The view-selector agent's one tool
# ---------------------------------------------------------------------------

@function_tool
def record_view_choice_tool(
    ctx: RunContextWrapper[FossilContext],
    chosen_views: str,
    reasoning: str = "",
) -> dict[str, Any]:
    """Record the views YOU chose to reconstruct from, by their V-numbers.

    This does not rank or filter anything. It writes down your decision and checks
    only that it is legal: 1 to 4 views, every number a real view.

    Args:
        chosen_views: JSON array of 1-4 view numbers as shown on the contact sheet,
            e.g. "[1, 3, 4]". Order does not matter.
        reasoning: Why these views. What each contributes that the others do not.
    """
    print(f"\n[Tool] record_view_choice_tool({chosen_views})")
    try:
        picked = [int(v) for v in json.loads(chosen_views)]
        available = list(getattr(ctx.context, "_candidate_views", []) or [])
        if not available:
            return {"status": "error", "message": "No candidate views in context."}

        if not 1 <= len(picked) <= specimen.MAX_VIEWS:
            return {"status": "error", "message": (
                f"Choose between 1 and {specimen.MAX_VIEWS} views; you chose "
                f"{len(picked)}. Meshy rejects more than {specimen.MAX_VIEWS}.")}
        bad = [v for v in picked if not 1 <= v <= len(available)]
        if bad:
            return {"status": "error",
                    "message": f"No such view(s): {bad}. Valid: 1..{len(available)}."}
        if len(set(picked)) != len(picked):
            return {"status": "error", "message": "Duplicate views chosen."}

        paths = [available[v - 1] for v in sorted(picked)]
        ctx.context._chosen_views = paths          # noqa: SLF001 — host-side handoff
        ctx.context._view_reasoning = reasoning    # noqa: SLF001

        print(f"  [Views] kept {sorted(picked)} of {len(available)}")
        return {
            "status": "success",
            "chosen_views": paths,
            "n_chosen": len(paths),
            "reasoning": reasoning,
            "next": "Views recorded. Summarise your choice and stop.",
        }
    except Exception as exc:  # noqa: BLE001
        return _fail("record_view_choice", exc)


# ---------------------------------------------------------------------------
# STEP 4 — Meshy, multi-image-to-3d
# ---------------------------------------------------------------------------

@function_tool
def generate_3d_model_tool(
    ctx: RunContextWrapper[FossilContext],
    cutout_paths: str,
    specimen_id: str = "",
    include_texture: str = "",
) -> dict[str, Any]:
    """Reconstruct a 3D mesh from the SAM3 cutouts — one per view — via Meshy's multi-image-to-3d endpoint.

    Pass the SAM3 CUTOUTS, not the original photographs. Feeding raw photos hands
    Meshy the backdrop and the cast shadow to reconstruct along with the fossil,
    which is the entire reason steps 1-2 exist.

    1 to 4 views. More is a hard API error, not a degradation.

    THIS SPENDS CREDITS. When 3D reconstruction is switched off for the run it
    returns the request it *would* have sent and nothing is spent. Report that plan
    and stop; do not try to work around the switch.

    Args:
        cutout_paths: JSON array of SAM3 cutout paths, one per view,
            e.g. '["output/UMMP-1/v1_sam3_cutout.png"]'.
        specimen_id: Folder name of the specimen, used to name the output dir.
        include_texture: "true" for a textured mesh, "false" for geometry only.
            Leave EMPTY to use the run's configured default — only set it if the
            user actually asked one way or the other.
    """
    print(f"\n[Tool] generate_3d_model_tool({cutout_paths})")
    try:
        c = ctx.context
        paths = json.loads(cutout_paths) if isinstance(cutout_paths, str) else cutout_paths
        paths = [str(p) for p in paths]

        if not paths:
            return {"status": "error", "message": "No cutouts supplied."}

        missing = [p for p in paths if not Path(p).exists()]
        if missing:
            return {"status": "error", "message": f"Cutouts not found: {missing}"}

        # ---- The one thing that must never be got wrong -------------------
        # Meshy reconstructs whatever it is given. Hand it a background-removal
        # cutout and it will faithfully build a mesh of fossil-plus-cast-shadow and
        # return it looking like a specimen. The artefacts have distinct names for
        # exactly this reason; enforce it rather than trust it.
        wrong = [p for p in paths if not p.endswith(SAM3_CUTOUT_SUFFIX)]
        if wrong:
            return {"status": "error", "message": (
                f"These are not SAM3 cutouts: {wrong}\n"
                f"Every view passed to Meshy must be a *{SAM3_CUTOUT_SUFFIX} file "
                f"produced by segment_fossil_sam3_tool.\n"
                f"  *_cutout_white.png / *_cutout_rgba.png = background-removal output. "
                f"That is the fuzzy matte WITH THE CAST SHADOW IN IT — the exact thing "
                f"SAM3 exists to replace. Never send it.\n"
                f"  *.jpg / *.jpeg = the raw photograph, backdrop and all.\n"
                f"If a view has no SAM3 cutout, drop that view. Do not substitute.")}

        if len(paths) > meshy_client.MAX_VIEWS:
            return {"status": "error", "message": (
                f"Meshy accepts at most {meshy_client.MAX_VIEWS} views, got "
                f"{len(paths)}. Use select_views_tool to choose.")}

        raw = (include_texture or "").strip().lower()
        if raw in {"true", "1", "yes", "on"}:
            textured = True
        elif raw in {"false", "0", "no", "off"}:
            textured = False
        elif raw == "":
            textured = c.include_texture   # the run's default; NOT a global env read
        else:
            return {"status": "error", "message": (
                f"include_texture must be 'true', 'false', or empty; got {include_texture!r}")}

        out_dir = (c.output_dir() if c.specimen_id else
                   Path("output") / (specimen_id or "specimen")) / "mesh"

        if not c.meshy_enabled:
            return {
                "status": "not_executed",
                "reason": "3D reconstruction is switched off for this run.",
                "plan": {
                    "endpoint": meshy_client.ENDPOINT,
                    "n_views": len(paths),
                    "views": paths,
                    "views_verified_as_sam3_cutouts": True,
                    "should_texture": textured,
                    "image_enhancement": False,
                    "should_remesh": False,
                    "output_dir": str(out_dir),
                },
                "message": (
                    f"Ready to reconstruct from {len(paths)} view(s), but 3D is switched "
                    f"off, so no credits were spent. Report this plan to the user and stop."
                ),
            }

        print(f"  [Meshy] sending {len(paths)} SAM3 cutout(s), "
              f"texture={'on' if textured else 'OFF (geometry only)'}:")
        for p in paths:
            print(f"    - {p}  ({Path(p).stat().st_size / 1024:.0f} KB)")

        task_id = meshy_client.create_task(paths, should_texture=textured)
        c.meshy_task_id = task_id
        c.meshy_status = "PENDING"
        c.meshy_percent = 0

        def _publish(status: str, percent: int) -> None:
            # Called from the polling loop in this worker thread. The UI watcher reads
            # these on the event loop; plain attribute writes are fine for that.
            c.meshy_status = status or ""
            c.meshy_percent = int(percent or 0)
        print(f"  [Meshy] task {task_id} — polling")
        task = meshy_client.wait_for_task(
            task_id, on_progress=lambda s, p: print(f"  [Meshy] {s} {p}%", flush=True))
        c.meshy_status = "SUCCEEDED"
        c.meshy_percent = 100
        models = meshy_client.download_models(task, out_dir)
        c.mesh_paths.update(models)

        return {
            "status": "success",
            "task_id": task_id,
            "n_views": len(paths),
            "textured": textured,
            "model_paths": models,
            "consumed_credits": task.get("consumed_credits"),
            "message": f"Mesh reconstructed from {len(paths)} view(s) -> {out_dir}",
        }
    except Exception as exc:  # noqa: BLE001
        return _fail("meshy_reconstruction", exc)