"""Run SamGeo3 text segmentation for multiple prompts (one model load).

Prompt sets are Grok + Gemini vision merged lists (see prompts/*.json).

Usage (from project root, venv active):
    python scripts/multi_prompt_segment.py
    python scripts/multi_prompt_segment.py --tier compact
    python scripts/multi_prompt_segment.py --tier A_high
    python scripts/multi_prompt_segment.py --tier all
    python scripts/multi_prompt_segment.py --prompts "building,solar panel,tree"
    python scripts/multi_prompt_segment.py --image data/DJI_20260306100802_0016.JPG --tier compact
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_PROMPTS_JSON = ROOT / "prompts" / "dji_20260306_0016.json"
DEFAULT_IMAGE = ROOT / "data" / "DJI_20260306100802_0016.JPG"
# fallback to sample folder outside lab if data copy missing
SAMPLE_FALLBACK = Path(
    r"D:\MYCLAUDE_PROJECT\segment-geospatial\sample\DJI_20260306100802_0016.JPG"
)


def safe_name(prompt: str) -> str:
    return "".join(c if c.isalnum() or c in "-_" else "_" for c in prompt.strip())


def _dedupe(prompts: list[str]) -> list[str]:
    seen: set[str] = set()
    out: list[str] = []
    for p in prompts:
        p = p.strip()
        if p and p not in seen:
            seen.add(p)
            out.append(p)
    return out


def _merge_prompt_lists(*lists: list[str]) -> list[str]:
    merged: list[str] = []
    for lst in lists:
        merged.extend(lst or [])
    return _dedupe(merged)


def load_prompt_config(path: Path, _stack: list[Path] | None = None) -> dict[str, Any]:
    """Load JSON; if 'extends' is set, inherit base (image-1) and apply scene delta.

    Extension model (image N expands image 1, not a separate system):
      - base tiers / compact / gap / improved / gap2* are kept
      - child may add scene_delta.prompts and optional tier/set additions
      - child image path overrides base image
    """
    path = path.resolve()
    stack = _stack or []
    if path in stack:
        raise ValueError(f"Circular extends: {path}")
    stack = stack + [path]

    with open(path, encoding="utf-8") as f:
        cfg = json.load(f)

    extends = cfg.get("extends")
    if not extends:
        return cfg

    base_path = Path(extends)
    if not base_path.is_absolute():
        # relative to project root first, then to this file's dir
        cand = (ROOT / extends).resolve()
        if not cand.is_file():
            cand = (path.parent / extends).resolve()
        base_path = cand
    if not base_path.is_file():
        raise FileNotFoundError(f"extends not found: {extends} (from {path})")

    base = load_prompt_config(base_path, stack)

    # Start from base, then overlay child metadata / image
    merged: dict[str, Any] = json.loads(json.dumps(base))  # deep copy
    merged["id"] = cfg.get("id", base.get("id"))
    merged["extends"] = str(extends)
    merged["extends_resolved"] = str(base_path)
    for key in (
        "image",
        "image_abs_hint",
        "sources",
        "scene_notes_ko",
        "recommended",
        "expected_difficulty",
    ):
        if key in cfg:
            merged[key] = cfg[key]

    # scene-only additions (this image expands base)
    delta = list(cfg.get("scene_delta", {}).get("prompts", []))
    merged["scene_delta"] = cfg.get("scene_delta", {"prompts": delta})

    # Merge tiers: base list + child additions (if child redefines tier fully, use add only
    # unless replace_tiers=true)
    child_tiers = cfg.get("tiers", {})
    if cfg.get("replace_tiers"):
        merged["tiers"] = child_tiers
    else:
        base_tiers = dict(merged.get("tiers", {}))
        for tname, tblock in child_tiers.items():
            add = list(tblock.get("prompts", []))
            if tname in base_tiers:
                base_tiers[tname] = {
                    **base_tiers[tname],
                    "description": tblock.get(
                        "description", base_tiers[tname].get("description", "")
                    ),
                    "prompts": _merge_prompt_lists(
                        base_tiers[tname].get("prompts", []), add
                    ),
                }
            else:
                base_tiers[tname] = tblock
        # also append scene_delta into A_high by default for visibility
        if delta and "A_high" in base_tiers:
            base_tiers["A_high"] = {
                **base_tiers["A_high"],
                "prompts": _merge_prompt_lists(
                    base_tiers["A_high"].get("prompts", []), delta
                ),
            }
        merged["tiers"] = base_tiers

    # Named sets: inherit base, extend with delta / child prompts
    for set_name in (
        "compact",
        "gap",
        "improved",
        "gap2",
        "gap2_core",
        "compact_plus",
    ):
        base_set = merged.get(set_name, {})
        base_prompts = list(base_set.get("prompts", [])) if isinstance(base_set, dict) else []
        child_set = cfg.get(set_name)
        if child_set is None:
            # default: base set + scene_delta for compact / improved
            if set_name in ("compact", "improved", "compact_plus") and delta:
                merged[set_name] = {
                    **(base_set if isinstance(base_set, dict) else {}),
                    "description": (
                        f"extends base {set_name} + scene_delta "
                        f"({cfg.get('id', 'child')})"
                    ),
                    "prompts": _merge_prompt_lists(base_prompts, delta),
                }
            continue
        if child_set.get("mode") == "replace":
            merged[set_name] = child_set
        else:
            # default mode: append (base + child + optional delta)
            extra = list(child_set.get("prompts", []))
            use_delta = child_set.get("include_scene_delta", True)
            parts = [base_prompts, extra]
            if use_delta:
                parts.append(delta)
            merged[set_name] = {
                **(base_set if isinstance(base_set, dict) else {}),
                **{k: v for k, v in child_set.items() if k != "prompts"},
                "description": child_set.get(
                    "description",
                    f"extends base {set_name} + additions",
                ),
                "prompts": _merge_prompt_lists(*parts),
            }

    # If child has no compact at all, still attach delta
    if "compact" not in cfg and delta:
        bp = list(base.get("compact", {}).get("prompts", []))
        merged["compact"] = {
            **base.get("compact", {}),
            "description": "base compact + scene_delta",
            "prompts": _merge_prompt_lists(bp, delta),
        }

    merged["inheritance_ko"] = (
        "1번 이미지 프롬프트 체계를 상속하고, 이 장면 전용 객체만 추가 확장."
    )
    return merged


def resolve_prompts(
    cfg: dict[str, Any],
    tier: str,
    extra_prompts: list[str] | None,
) -> list[str]:
    if extra_prompts:
        return [p.strip() for p in extra_prompts if p.strip()]

    # scene_delta only (this image's new objects vs base)
    if tier in ("delta", "scene_delta", "new"):
        return list(cfg.get("scene_delta", {}).get("prompts", []))

    # top-level named sets (compact / gap / improved / gap2...)
    if tier in cfg and isinstance(cfg[tier], dict) and "prompts" in cfg[tier]:
        return list(cfg[tier]["prompts"])

    if tier == "all":
        seen: set[str] = set()
        out: list[str] = []
        for key in (
            "A_high",
            "B_facility",
            "C_detail",
            "D_rail_domain",
            "E_missed",
        ):
            block = cfg.get("tiers", {}).get(key, {})
            for p in block.get("prompts", []):
                if p not in seen:
                    seen.add(p)
                    out.append(p)
        # include scene_delta
        for p in cfg.get("scene_delta", {}).get("prompts", []):
            if p not in seen:
                seen.add(p)
                out.append(p)
        return out

    if tier in cfg.get("tiers", {}):
        return list(cfg["tiers"][tier]["prompts"])

    # allow short aliases
    aliases = {
        "a": "A_high",
        "A": "A_high",
        "b": "B_facility",
        "B": "B_facility",
        "c": "C_detail",
        "C": "C_detail",
        "d": "D_rail_domain",
        "D": "D_rail_domain",
        "e": "E_missed",
        "E": "E_missed",
        "missed": "E_missed",
        "gap_fill": "gap",
    }
    if tier in aliases:
        key = aliases[tier]
        if key in cfg.get("tiers", {}):
            return list(cfg["tiers"][key]["prompts"])
        if key in cfg and "prompts" in cfg[key]:
            return list(cfg[key]["prompts"])

    raise ValueError(
        f"Unknown tier '{tier}'. Use compact | gap | gap2 | improved | all | "
        f"delta | A_high | B_facility | C_detail | D_rail_domain | E_missed | or --prompts"
    )


def resolve_image(path: Path | None, cfg: dict[str, Any]) -> Path:
    candidates: list[Path] = []
    if path is not None:
        candidates.append(path)
    candidates.append(ROOT / cfg.get("image", "data/DJI_20260306100802_0016.JPG"))
    candidates.append(DEFAULT_IMAGE)
    candidates.append(SAMPLE_FALLBACK)
    abs_hint = cfg.get("image_abs_hint")
    if abs_hint:
        candidates.append(Path(abs_hint))

    for c in candidates:
        if c and Path(c).is_file():
            return Path(c).resolve()
    raise FileNotFoundError(
        "Image not found. Place DJI JPG under data/ or pass --image. Tried:\n  "
        + "\n  ".join(str(c) for c in candidates)
    )


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Multi-prompt SamGeo3 segmentation (Grok+Gemini prompt sets)"
    )
    p.add_argument(
        "--prompts-json",
        type=Path,
        default=DEFAULT_PROMPTS_JSON,
        help="JSON with tiers/compact prompts",
    )
    p.add_argument(
        "--tier",
        default="compact",
        help=(
            "compact | gap | gap2 | improved | all | "
            "A_high | B_facility | C_detail | D_rail_domain | E_missed"
        ),
    )
    p.add_argument(
        "--prompts",
        default=None,
        help='Comma-separated overrides, e.g. "building,tree,car"',
    )
    p.add_argument("--image", type=Path, default=None)
    p.add_argument("--model-id", default="facebook/sam3.1")
    p.add_argument("--backend", default="meta", choices=["meta", "transformers"])
    p.add_argument("--confidence", type=float, default=0.3)
    p.add_argument("--mask-threshold", type=float, default=0.5)
    p.add_argument("--resolution", type=int, default=1008)
    p.add_argument("--min-size", type=int, default=0)
    p.add_argument("--max-size", type=int, default=None)
    p.add_argument("--device", default=None)
    p.add_argument(
        "--checkpoint",
        default=os.environ.get("SAM3_CHECKPOINT_PATH"),
    )
    p.add_argument(
        "--output-dir",
        type=Path,
        default=None,
        help="Default: output/<image_stem>_<tier>/",
    )
    p.add_argument("--no-viz", action="store_true")
    p.add_argument(
        "--skip-empty",
        action="store_true",
        default=True,
        help="Do not write mask files when 0 objects (default true)",
    )
    p.add_argument(
        "--keep-empty",
        action="store_true",
        help="Opposite of --skip-empty",
    )
    p.add_argument(
        "--list-only",
        action="store_true",
        help="Print resolved prompts and exit (no model load)",
    )
    return p.parse_args()


def main() -> int:
    args = parse_args()
    cfg_path = args.prompts_json.resolve()
    if not cfg_path.is_file():
        print(f"ERROR: prompts json not found: {cfg_path}", file=sys.stderr)
        return 2

    cfg = load_prompt_config(cfg_path)
    extra = None
    if args.prompts:
        extra = [x.strip() for x in args.prompts.split(",") if x.strip()]

    try:
        prompts = resolve_prompts(cfg, args.tier, extra)
    except ValueError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        return 2

    # de-dupe preserve order
    seen: set[str] = set()
    prompts = [p for p in prompts if not (p in seen or seen.add(p))]

    if args.list_only:
        print(f"prompts_json: {cfg_path}")
        print(f"tier: {args.tier if not extra else 'custom'}")
        print(f"count: {len(prompts)}")
        for i, pr in enumerate(prompts, 1):
            print(f"  {i:02d}. {pr}")
        return 0

    try:
        image_path = resolve_image(args.image, cfg)
    except FileNotFoundError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        return 2

    if args.backend == "transformers" and "sam3.1" in args.model_id:
        print(
            "ERROR: facebook/sam3.1 requires backend='meta'.",
            file=sys.stderr,
        )
        return 2

    tier_tag = "custom" if extra else args.tier
    out_dir = (
        args.output_dir.resolve()
        if args.output_dir
        else (ROOT / "output" / f"{image_path.stem}_{tier_tag}").resolve()
    )
    out_dir.mkdir(parents=True, exist_ok=True)
    skip_empty = not args.keep_empty

    print("=== SamGeo3 multi-prompt segmentation ===")
    print(f"image:       {image_path}")
    print(f"prompts_json:{cfg_path}")
    print(f"tier:        {tier_tag}")
    print(f"n_prompts:   {len(prompts)}")
    print(f"model_id:    {args.model_id}")
    print(f"backend:     {args.backend}")
    print(f"confidence:  {args.confidence}")
    print(f"output_dir:  {out_dir}")
    print("prompts:")
    for i, pr in enumerate(prompts, 1):
        print(f"  {i:02d}. {pr}")

    import numpy as np
    import torch
    from samgeo import SamGeo3

    if args.device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        device = args.device

    print(f"\ndevice: {device} | cuda={torch.cuda.is_available()}")
    if device == "cuda" and not torch.cuda.is_available():
        print("ERROR: --device cuda but CUDA is not available", file=sys.stderr)
        return 1

    init_kwargs: dict[str, Any] = dict(
        backend=args.backend,
        model_id=args.model_id,
        device=device,
        confidence_threshold=args.confidence,
        mask_threshold=args.mask_threshold,
        resolution=args.resolution,
        enable_segmentation=True,
        enable_inst_interactivity=False,
    )
    if args.checkpoint and os.path.isfile(args.checkpoint):
        init_kwargs["checkpoint_path"] = args.checkpoint
        init_kwargs["load_from_HF"] = False
        print(f"checkpoint: {args.checkpoint}")

    t0 = time.perf_counter()
    print("\nLoading model once...")
    sam = SamGeo3(**init_kwargs)
    print("set_image once...")
    sam.set_image(str(image_path))
    t_load = time.perf_counter() - t0
    print(f"load+set_image: {t_load:.1f}s\n")

    results: list[dict[str, Any]] = []
    stem = image_path.stem

    for i, prompt in enumerate(prompts, 1):
        t1 = time.perf_counter()
        print(f"[{i}/{len(prompts)}] generate_masks({prompt!r})...")
        gen_kwargs: dict[str, Any] = {"min_size": args.min_size, "quiet": True}
        if args.max_size is not None:
            gen_kwargs["max_size"] = args.max_size

        try:
            sam.generate_masks(prompt, **gen_kwargs)
        except Exception as e:
            print(f"  FAIL: {e}")
            results.append(
                {
                    "prompt": prompt,
                    "n_objects": 0,
                    "error": str(e),
                    "elapsed_s": round(time.perf_counter() - t1, 3),
                }
            )
            continue

        n = len(sam.masks) if getattr(sam, "masks", None) is not None else 0
        score_vals: list[float] = []
        scores = getattr(sam, "scores", None)
        if scores is not None and len(scores):
            try:
                score_vals = [
                    float(s.item() if hasattr(s, "item") else s) for s in scores
                ]
            except Exception:
                pass

        rec: dict[str, Any] = {
            "prompt": prompt,
            "n_objects": n,
            "scores_min": min(score_vals) if score_vals else None,
            "scores_max": max(score_vals) if score_vals else None,
            "scores_mean": float(np.mean(score_vals)) if score_vals else None,
            "elapsed_s": round(time.perf_counter() - t1, 3),
            "mask": None,
            "ann": None,
            "scores_npy": None,
        }

        if n == 0:
            print(f"  -> 0 objects ({rec['elapsed_s']}s)")
            results.append(rec)
            continue

        print(
            f"  -> {n} objects | score "
            f"{rec['scores_min']:.3f}~{rec['scores_max']:.3f} "
            f"({rec['elapsed_s']}s)"
        )

        if skip_empty and n == 0:
            results.append(rec)
            continue

        tag = safe_name(prompt)
        mask_path = out_dir / f"{stem}_{tag}_mask.png"
        ann_path = out_dir / f"{stem}_{tag}_ann.png"
        scores_path = out_dir / f"{stem}_{tag}_scores.npy"

        try:
            sam.save_masks(str(mask_path), unique=True)
            rec["mask"] = str(mask_path)
        except Exception as e:
            print(f"  [WARN] save_masks: {e}")

        if score_vals:
            np.save(str(scores_path), np.array(score_vals, dtype=np.float32))
            rec["scores_npy"] = str(scores_path)

        if not args.no_viz:
            try:
                sam.show_anns(output=str(ann_path))
                rec["ann"] = str(ann_path)
            except Exception as e:
                print(f"  [WARN] show_anns: {e}")

        results.append(rec)

    summary = {
        "image": str(image_path),
        "prompts_json": str(cfg_path),
        "tier": tier_tag,
        "model_id": args.model_id,
        "backend": args.backend,
        "confidence": args.confidence,
        "min_size": args.min_size,
        "device": device,
        "load_set_image_s": round(t_load, 3),
        "total_s": round(time.perf_counter() - t0, 3),
        "n_prompts": len(prompts),
        "n_with_objects": sum(1 for r in results if r.get("n_objects", 0) > 0),
        "results": results,
        "sources": cfg.get("sources"),
        "scene_notes_ko": cfg.get("scene_notes_ko"),
    }
    summary_path = out_dir / "summary.json"
    with open(summary_path, "w", encoding="utf-8") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)

    print("\n=== summary ===")
    print(f"with objects: {summary['n_with_objects']}/{summary['n_prompts']}")
    print(f"total time:   {summary['total_s']}s")
    print(f"summary:      {summary_path}")
    for r in results:
        n = r.get("n_objects", 0)
        mark = "OK" if n else "--"
        print(f"  [{mark}] {r['prompt']!r:30s} n={n}")

    return 0


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