from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

from ultralytics import YOLO

DEFAULT_PROJECT_ROOT = Path(__file__).resolve().parents[4]


def project_root() -> Path:
    raw = os.getenv("SEG_DATA_SERVER_ROOT")
    if not raw:
        return DEFAULT_PROJECT_ROOT
    path = Path(raw).expanduser()
    if path.is_absolute():
        return path.resolve()
    return (DEFAULT_PROJECT_ROOT / path).resolve()


PROJECT_ROOT = project_root()


def resolve_project_path(value: str | Path) -> Path:
    path = Path(value).expanduser()
    if path.is_absolute():
        return path.resolve()
    return (PROJECT_ROOT / path).resolve()


def main() -> None:
    parser = argparse.ArgumentParser(description="Predict segmentation masks with a supplied YOLO checkpoint.")
    parser.add_argument("--weights", required=True, help="Path to best.pt/last.pt or another YOLO segmentation checkpoint.")
    parser.add_argument("--source", required=True, help="Image file or image directory to predict.")
    parser.add_argument("--project", default=str(PROJECT_ROOT / "var" / "custom_yolo_runs"))
    parser.add_argument("--name", default="custom_predict")
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--conf", type=float, default=0.25)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--save-txt", action="store_true")
    parser.add_argument("--save-conf", action="store_true")
    parser.add_argument("--exist-ok", action="store_true")
    args = parser.parse_args()

    weights = resolve_project_path(args.weights)
    source = resolve_project_path(args.source)
    project = resolve_project_path(args.project)
    if not weights.exists():
        raise FileNotFoundError(f"weights not found: {weights}")
    if not source.exists():
        raise FileNotFoundError(f"source not found: {source}")

    model = YOLO(str(weights))
    results = model.predict(
        source=str(source),
        imgsz=args.imgsz,
        conf=args.conf,
        device=args.device,
        project=str(project),
        name=args.name,
        save=True,
        save_txt=args.save_txt,
        save_conf=args.save_conf,
        exist_ok=args.exist_ok,
        verbose=True,
    )
    save_dir = Path(getattr(model.predictor, "save_dir", project / args.name))
    metadata = {
        "weights": str(weights),
        "source": str(source),
        "save_dir": str(save_dir),
        "count": len(results),
    }
    save_dir.mkdir(parents=True, exist_ok=True)
    (save_dir / "predict_manifest.json").write_text(json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(metadata, ensure_ascii=False))


if __name__ == "__main__":
    main()