from __future__ import annotations

import argparse
from pathlib import Path

import cv2

from utils import draw_obb_res
from yolo_obb import Yolo26Obb


def parse_args():
    parser = argparse.ArgumentParser(description="Run YOLO26l-OBB QNN model with AidLite.")
    parser.add_argument(
        "--target_model",
        type=str,
        default="./models/QCS8550/FP16/yolo26l-obb_qcs8550_fp16.qnn248.ctx.bin",
        help=(
            "ctx.bin 路径，可换成 ./models/QCS8550/W8A16/ 或 ./models/QCS8550/W8A8/ 下的量化模型"
        ),
    )
    parser.add_argument("--imgs", type=str, default="./code/python/boats.jpg")
    parser.add_argument("--height", type=int, default=1024)
    parser.add_argument("--width", type=int, default=1024)
    parser.add_argument("--invoke_nums", type=int, default=10)
    parser.add_argument("--model_type", type=str, default="QNN")
    parser.add_argument("--conf_thres", type=float, default=0.25)
    parser.add_argument("--iou_thres", type=float, default=0.45)
    parser.add_argument("--save_path", type=str, default="./code/python/result.jpg")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    frame = cv2.imread(args.imgs)
    if frame is None:
        raise FileNotFoundError(args.imgs)

    model = Yolo26Obb(args.target_model, args.width, args.height, args.model_type)
    detections = model(frame, args.invoke_nums, args.conf_thres, args.iou_thres)
    print("===================")
    print(f"Detect {len(detections)} targets.")
    result = draw_obb_res(frame, detections)
    Path(args.save_path).parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(args.save_path, result)
    return 0


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