agent_core/outputs.py

import json
import os
from datetime import datetime
from pathlib import Path
from uuid import uuid4

from agent_core.config import ARTIFACT_ROOT, DEFAULT_LUX3D_OUTPUT_PATH, STEP_SUFFIXES, WORKDIR


def safe_path(p: str) -> Path:
    raw_path = Path(p)
    path = (raw_path if raw_path.is_absolute() else WORKDIR / raw_path).resolve()
    if not path.is_relative_to(WORKDIR) and not path.is_relative_to(ARTIFACT_ROOT):
        raise ValueError(f"Path escapes workspace: {p}")
    return path


def workspace_relative(path: Path) -> str:
    absolute = path if path.is_absolute() else WORKDIR / path
    normalized = Path(os.path.normpath(str(absolute)))
    if normalized.is_relative_to(WORKDIR):
        return str(normalized.relative_to(WORKDIR))
    return str(normalized)


def new_run_id() -> str:
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    return f"{timestamp}_{uuid4().hex[:6]}"


def allocate_run_dir(output_root: Path) -> tuple[Path, str]:
    for _ in range(20):
        run_id = new_run_id()
        run_dir = output_root / "runs" / run_id
        if not run_dir.exists():
            return run_dir, run_id

    raise RuntimeError("Failed to allocate a unique output run directory.")


def resolve_run_output(requested_output_path: str) -> tuple[Path, Path, str]:
    requested_path = safe_path(requested_output_path)
    suffix = requested_path.suffix.lower()

    if suffix in STEP_SUFFIXES:
        output_root = requested_path.parent
        output_name = requested_path.name
    elif suffix:
        raise ValueError("Only STEP output is supported in this version. Use a .step/.stp file path or a directory path.")
    else:
        output_root = requested_path
        output_name = "model.step"

    run_dir, run_id = allocate_run_dir(output_root)
    return run_dir, run_dir / output_name, run_id


def resolve_lux3d_run(output_path: str | None) -> tuple[Path, Path, str, str | None, str]:
    requested_output_path = output_path or DEFAULT_LUX3D_OUTPUT_PATH
    requested_path = safe_path(requested_output_path)
    if requested_path.suffix:
        output_root = requested_path.parent
        output_name = requested_path.name
    else:
        output_root = requested_path
        output_name = None

    run_dir, run_id = allocate_run_dir(output_root)
    return output_root, run_dir, run_id, output_name, requested_output_path


def update_latest_link(output_root: Path, run_dir: Path) -> str | None:
    latest = output_root / "latest"
    try:
        if latest.is_symlink() or latest.is_file():
            latest.unlink()
        elif latest.exists():
            return f"Skipped latest link because {workspace_relative(latest)} already exists and is not a symlink."

        latest.symlink_to(Path("runs") / run_dir.name, target_is_directory=True)
    except OSError as exc:
        return f"Failed to update latest link: {exc}"

    return None


def write_manifest(
    run_dir: Path,
    run_id: str,
    requested_output_path: str,
    output_path: Path,
    code: str,
    prompt: str | None,
    payload: dict,
) -> Path:
    manifest_path = run_dir / "manifest.json"
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest = {
        "run_id": run_id,
        "created_at": datetime.now().astimezone().isoformat(timespec="seconds"),
        "generator": "cadquery",
        "prompt": prompt,
        "requested_output_path": requested_output_path,
        "output_path": workspace_relative(output_path),
        "preview_path": payload.get("preview_path"),
        "format": output_path.suffix.lstrip(".").lower(),
        "code": code,
        "stdout": payload.get("stdout", ""),
        "stderr": payload.get("stderr", ""),
        "warning": payload.get("warning"),
    }
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    return manifest_path


def write_lux3d_manifest(
    run_dir: Path,
    run_id: str,
    prompt: str | None,
    image_path: Path,
    requested_output_path: str,
    output_path: Path | None,
    preview_path: Path | None,
    busid: str | int | None,
    status: int | None,
    result_url: str | None,
    poll_count: int,
    error: str | None = None,
) -> Path:
    manifest_path = run_dir / "manifest.json"
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest = {
        "run_id": run_id,
        "created_at": datetime.now().astimezone().isoformat(timespec="seconds"),
        "generator": "lux3d",
        "prompt": prompt,
        "image_path": workspace_relative(image_path),
        "requested_output_path": requested_output_path,
        "output_path": workspace_relative(output_path) if output_path else None,
        "preview_path": workspace_relative(preview_path) if preview_path else None,
        "busid": busid,
        "status": status,
        "result_url": result_url,
        "poll_count": poll_count,
        "error": error,
    }
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    return manifest_path