import argparse
import subprocess
import sys
from pathlib import Path


def find_latest_best_weight(root: Path) -> Path:
    train_dir = root
    candidates = sorted(
        train_dir.glob("run/train/weights/best.pt"),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    if not candidates:
        raise FileNotFoundError("No best.pt found under run/train/weights/")
    return candidates[0]


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Detect graph bounding boxes in testing folder using trained YOLOv7 model."
    )
    parser.add_argument(
        "--weights",
        type=str,
        default="",
        help="Optional path to model weights. If empty, latest run/train/exp5_full_finetune3/weights/best.pt is used.",
    )
    parser.add_argument(
        "--source",
        type=str,
        default="testing",
        help="Input image folder or file path.",
    )
    parser.add_argument("--img-size", type=int, default=640, help="Inference image size")
    parser.add_argument("--conf-thres", type=float, default=0.45, help="Confidence threshold")
    parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold")
    parser.add_argument("--device", type=str, default="0", help="Device, e.g. 0 or cpu")
    parser.add_argument("--name", type=str, default="testing_graph_bbox", help="Output run name")
    parser.add_argument("--view-img", action="store_true", help="Show images while running")
    parser.add_argument("--no-save-txt", action="store_true", help="Do not save YOLO txt labels")
    args = parser.parse_args()

    root = Path(__file__).resolve().parent
    source = Path(args.source)
    if not source.is_absolute():
        source = (root / source).resolve()

    if args.weights:
        weights = Path(args.weights)
        if not weights.is_absolute():
            weights = (root / weights).resolve()
    else:
        weights = find_latest_best_weight(root)

    if not weights.exists():
        print(f"[ERROR] Weights not found: {weights}")
        return 1
    if not source.exists():
        print(f"[ERROR] Source not found: {source}")
        return 1

    cmd = [
        sys.executable,
        str(root / "detect.py"),
        "--weights",
        str(weights),
        "--source",
        str(source),
        "--img-size",
        str(args.img_size),
        "--conf-thres",
        str(args.conf_thres),
        "--iou-thres",
        str(args.iou_thres),
        "--device",
        args.device,
        "--classes",
        "0",
        "--project",
        "run/detect",
        "--name",
        args.name,
        "--exist-ok",
    ]

    if not args.no_save_txt:
        cmd.extend(["--save-txt", "--save-conf"])
    if args.view_img:
        cmd.append("--view-img")

    print("Using weights:", weights)
    print("Source:", source)
    print("Running command:")
    print(" ".join(cmd))

    return subprocess.call(cmd, cwd=str(root))


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