# -*- coding: utf-8 -*-
"""Extract soup-only masks from frame images using the reference RF-DETR model.

This script intentionally keeps the trained segmentation model in the reference
project path. It writes only derived crop/mask images into this workspace.
"""

from __future__ import annotations

import argparse
import csv
import sys
from pathlib import Path
from typing import Any


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_REFERENCE_SRC = (
    Path.home()
    / "Documents"
    / "開発しているもの"
    / "制作中"
    / "AIラーメン開発"
    / "ramen-soup-ai-suite"
    / "03_soup-density-ai"
    / "src"
)
DEFAULT_CHECKPOINT = (
    Path.home()
    / "Documents"
    / "開発しているもの"
    / "制作中"
    / "AIラーメン開発"
    / "ramen-soup-ai-suite"
    / "03_soup-density-ai"
    / "models"
    / "crop"
    / "soup_valid_area_rfdetr_seg_v2"
    / "checkpoint_best_ema.pth"
)
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}


def main() -> int:
    parser = argparse.ArgumentParser(description="Extract soup masks from saved frame images.")
    parser.add_argument("--input-dir", required=True, help="Directory containing extracted frame images.")
    parser.add_argument("--output-dir", default="local_batch_api/density_features/soup_masks")
    parser.add_argument("--report", default="local_batch_api/density_features/soup_masks_report.csv")
    parser.add_argument("--reference-src", default=str(DEFAULT_REFERENCE_SRC))
    parser.add_argument("--checkpoint-path", default=str(DEFAULT_CHECKPOINT))
    parser.add_argument("--model-size", default="small", choices=["small", "medium"])
    parser.add_argument("--threshold", type=float, default=0.5)
    parser.add_argument("--low-confidence-threshold", type=float, default=0.35)
    parser.add_argument("--output-kind", default="crop", choices=["crop", "masked", "alpha", "all"])
    parser.add_argument("--limit", type=int, default=0)
    args = parser.parse_args()

    reference_src = Path(args.reference_src).expanduser().resolve()
    checkpoint_path = Path(args.checkpoint_path).expanduser().resolve()
    input_dir = resolve_path(args.input_dir)
    output_dir = resolve_path(args.output_dir)
    report_path = resolve_path(args.report)

    if not reference_src.exists():
        raise FileNotFoundError(f"reference src not found: {reference_src}")
    if not checkpoint_path.exists():
        raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}")
    if not input_dir.exists():
        raise FileNotFoundError(f"input dir not found: {input_dir}")

    sys.path.insert(0, str(reference_src))
    try:
        from crop.extract_soup_area import build_predictor, extract_one_image
    except ImportError as exc:
        raise RuntimeError(
            "Failed to import the reference soup extractor. "
            "Install its dependencies first, especially `rfdetr`, `torch`, `pillow`, and `opencv-python`."
        ) from exc

    images = list_images(input_dir)
    if args.limit > 0:
        images = images[: args.limit]
    output_dir.mkdir(parents=True, exist_ok=True)

    predictor = build_predictor(args.model_size, str(checkpoint_path), args.threshold)
    rows: list[dict[str, Any]] = []
    for image_path in images:
        print(f"MASK {image_path.name}")
        row = extract_one_image(
            image_path=image_path,
            output_dir=output_dir,
            checkpoint_path=str(checkpoint_path),
            model_size=args.model_size,
            threshold=args.threshold,
            low_confidence_threshold=args.low_confidence_threshold,
            predictor=predictor,
            save_alpha=args.output_kind in {"alpha", "all"},
            preserve_filename=True,
            output_kind=args.output_kind,
            skip_fallback_output=False,
        )
        rows.append(row)

    write_report(report_path, rows)
    ok = sum(1 for row in rows if int(row.get("extract_status", 0)) > 0)
    print(f"done images={len(rows)} extracted={ok} report={report_path}")
    return 0


def resolve_path(value: str) -> Path:
    path = Path(value)
    return path if path.is_absolute() else ROOT / path


def list_images(path: Path) -> list[Path]:
    return sorted(p for p in path.rglob("*") if p.suffix.lower() in IMAGE_EXTENSIONS)


def write_report(path: Path, rows: list[dict[str, Any]]) -> None:
    fields: list[str] = []
    for row in rows:
        for key in row:
            if key not in fields:
                fields.append(key)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8-sig", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)


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