"""Single-image text-prompt segmentation with SamGeo3 (meta backend).

Usage (from project root, with venv active):
    python scripts/text_segment.py
    python scripts/text_segment.py --image data/test_image.jpg --prompt person
    python scripts/text_segment.py --model-id facebook/sam3.1 --confidence 0.4
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path

# Project root = parent of scripts/
ROOT = Path(__file__).resolve().parents[1]
DEFAULT_DATA = ROOT / "data"
DEFAULT_OUTPUT = ROOT / "output"
SAMPLE_URL = (
    "https://raw.githubusercontent.com/facebookresearch/sam3/"
    "refs/heads/main/assets/images/test_image.jpg"
)


def ensure_sample_image(path: Path) -> Path:
    if path.is_file():
        return path
    path.parent.mkdir(parents=True, exist_ok=True)
    print(f"Downloading sample image -> {path}")
    try:
        from samgeo import download_file

        download_file(SAMPLE_URL, str(path))
    except Exception:
        import urllib.request

        urllib.request.urlretrieve(SAMPLE_URL, str(path))
    if not path.is_file():
        raise FileNotFoundError(f"failed to obtain image: {path}")
    return path


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="SamGeo3 text segmentation smoke test")
    p.add_argument("--image", type=Path, default=DEFAULT_DATA / "test_image.jpg")
    p.add_argument("--prompt", default="person", help="Text prompt for grounding")
    p.add_argument("--model-id", default="facebook/sam3.1")
    p.add_argument("--backend", default="meta", choices=["meta", "transformers"])
    # sam3.1 meta scores on this sample often peak ~0.3–0.4; 0.5 can yield zero masks
    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, help="Filter tiny masks (pixels)")
    p.add_argument("--max-size", type=int, default=None)
    p.add_argument("--device", default=None, help="cuda | cpu (default: auto)")
    p.add_argument(
        "--checkpoint",
        default=os.environ.get("SAM3_CHECKPOINT_PATH"),
        help="Local .pt path (or set SAM3_CHECKPOINT_PATH)",
    )
    p.add_argument(
        "--output-dir",
        type=Path,
        default=DEFAULT_OUTPUT,
        help="Directory for mask / annotation outputs",
    )
    p.add_argument(
        "--no-viz",
        action="store_true",
        help="Skip matplotlib annotation PNG (mask file still saved)",
    )
    return p.parse_args()


def main() -> int:
    args = parse_args()

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

    image_path = ensure_sample_image(args.image.resolve())
    out_dir = args.output_dir.resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    stem = image_path.stem
    safe_prompt = "".join(c if c.isalnum() or c in "-_" else "_" for c in args.prompt)
    mask_path = out_dir / f"{stem}_{safe_prompt}_mask.png"
    ann_path = out_dir / f"{stem}_{safe_prompt}_ann.png"
    scores_path = out_dir / f"{stem}_{safe_prompt}_scores.npy"

    print("=== SamGeo3 text segmentation ===")
    print(f"image:      {image_path}")
    print(f"prompt:     {args.prompt}")
    print(f"model_id:   {args.model_id}")
    print(f"backend:    {args.backend}")
    print(f"output_dir: {out_dir}")

    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"device:     {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(
        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}")

    print("\nLoading model...")
    sam = SamGeo3(**init_kwargs)

    print("set_image...")
    sam.set_image(str(image_path))

    print(f'generate_masks("{args.prompt}")...')
    gen_kwargs = {"min_size": args.min_size}
    if args.max_size is not None:
        gen_kwargs["max_size"] = args.max_size
    sam.generate_masks(args.prompt, **gen_kwargs)

    n = len(sam.masks) if getattr(sam, "masks", None) is not None else 0
    if n == 0:
        print("No masks found. Try another prompt or lower --confidence.")
        return 0

    scores = getattr(sam, "scores", None)
    if scores is not None and len(scores):
        try:
            vals = [float(s.item() if hasattr(s, "item") else s) for s in scores]
            print(f"scores (n={len(vals)}): min={min(vals):.3f} max={max(vals):.3f}")
        except Exception:
            pass

    print(f"Saving masks -> {mask_path}")
    # PNG cannot store float score maps; save mask first, scores as .npy
    sam.save_masks(str(mask_path), unique=True)

    scores = getattr(sam, "scores", None)
    if scores is not None and len(scores):
        import numpy as np

        score_vals = np.array(
            [float(s.item() if hasattr(s, "item") else s) for s in scores],
            dtype=np.float32,
        )
        np.save(str(scores_path), score_vals)
        print(f"Saved per-object scores -> {scores_path}")

    if not args.no_viz:
        try:
            print(f"Saving annotations -> {ann_path}")
            sam.show_anns(output=str(ann_path))
        except Exception as e:
            print(f"[WARN] show_anns failed: {e}")

    print(f"\nDone. Found {n} object(s).")
    print(f"  mask:   {mask_path}")
    if scores_path.is_file():
        print(f"  scores: {scores_path}")
    if ann_path.is_file():
        print(f"  ann:    {ann_path}")
    return 0


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