#!/usr/bin/env python3
"""Export 3DEditFormer Edit3D-Bench predictions to official evaluator layout."""

from __future__ import annotations

import argparse
import json
import os
import shutil
from pathlib import Path
from typing import Any


def replace_link(dst: Path, src: Path, *, copy: bool = False) -> None:
    if dst.exists() or dst.is_symlink():
        if dst.is_symlink() and Path(os.readlink(dst)) == src:
            return
        dst.unlink()
    dst.parent.mkdir(parents=True, exist_ok=True)
    if copy:
        shutil.copy2(src, dst)
    else:
        os.symlink(src, dst)


def render_id(path: Path) -> str | None:
    stem = path.stem
    if stem.startswith("render_"):
        suffix = stem.removeprefix("render_")
    else:
        suffix = stem
    if suffix.isdigit():
        return f"{int(suffix):04d}"
    return None


def load_task_map(adapter_root: Path) -> dict[str, dict[str, Any]]:
    path = adapter_root / "edit3d_task_map.json"
    with path.open("r", encoding="utf-8") as f:
        task_map = json.load(f)
    if not isinstance(task_map, dict):
        raise TypeError(f"Expected dict task map in {path}")
    return task_map


def export_predictions(args: argparse.Namespace) -> dict[str, Any]:
    adapter_root = Path(args.adapter_root).resolve()
    eval_results_dir = Path(args.eval_results_dir).resolve()
    pred_root = Path(args.pred_root).resolve()
    task_map = load_task_map(adapter_root)

    exported_models = 0
    exported_images = 0
    missing: list[str] = []

    for key, meta in task_map.items():
        src_task_dir = eval_results_dir / "flux_edit" / key
        src_mesh = src_task_dir / "edit.glb"
        src_render_dir = src_task_dir / "render"

        dst_task_dir = pred_root / meta["dataset"] / meta["source_model"] / meta["prompt_key"]
        dst_mesh = dst_task_dir / "edit.glb"
        dst_image_dir = dst_task_dir / "images"

        if src_mesh.exists():
            replace_link(dst_mesh, src_mesh, copy=args.copy)
            exported_models += 1
        else:
            missing.append(f"{key}: missing mesh: {src_mesh}")

        if not src_render_dir.exists():
            missing.append(f"{key}: missing render dir: {src_render_dir}")
            continue

        render_files = sorted(
            p for p in src_render_dir.glob("*.png") if render_id(p) is not None
        )
        if not render_files:
            missing.append(f"{key}: no render pngs in {src_render_dir}")
            continue

        for src_img in render_files:
            img_id = render_id(src_img)
            if img_id is None:
                continue
            dst_img = dst_image_dir / f"render_{img_id}.png"
            replace_link(dst_img, src_img, copy=args.copy)
            exported_images += 1

    summary = {
        "pred_root": str(pred_root),
        "num_tasks": len(task_map),
        "exported_models": exported_models,
        "exported_images": exported_images,
        "missing": missing,
    }
    pred_root.mkdir(parents=True, exist_ok=True)
    (pred_root / "export_summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    return summary


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--adapter_root", required=True)
    parser.add_argument("--eval_results_dir", required=True)
    parser.add_argument("--pred_root", required=True)
    parser.add_argument("--copy", action="store_true", help="Copy files instead of symlinking")
    parser.add_argument("--strict", action="store_true")
    args = parser.parse_args()

    summary = export_predictions(args)
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    if args.strict and summary["missing"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()