"""
Pre-label images using an existing YOLO model.
Generates .txt labels for unlabeled images so the annotator can load them for review.

Usage:
    python prelabel.py [image_dir] [--model boss-v1] [--conf 0.20]
"""
import argparse
import glob
import os

def run_prelabel(args):
    """Run pre-labeling. Called from main() or manage.py."""
    img_dir = os.path.abspath(args.img_dir)
    model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models", f"{args.model}.pt")

    if not os.path.exists(model_path):
        print(f"Model not found: {model_path}")
        return

    from ultralytics import YOLO
    model = YOLO(model_path)

    extensions = ("*.jpg", "*.jpeg", "*.png")
    files = []
    for ext in extensions:
        files.extend(glob.glob(os.path.join(img_dir, ext)))
    files.sort()

    # Only process unlabeled images
    unlabeled = []
    for f in files:
        label_path = os.path.splitext(f)[0] + ".txt"
        if not os.path.exists(label_path):
            unlabeled.append(f)

    print(f"Found {len(files)} images, {len(files) - len(unlabeled)} already labeled, {len(unlabeled)} to pre-label")

    if not unlabeled:
        print("All images already have labels!")
        return

    labeled = 0
    skipped = 0
    for filepath in unlabeled:
        results = model(filepath, conf=args.conf, verbose=False)
        boxes = results[0].boxes

        if len(boxes) == 0:
            skipped += 1
            continue

        label_path = os.path.splitext(filepath)[0] + ".txt"
        with open(label_path, "w") as f:
            for box in boxes:
                cls = int(box.cls[0])
                xywhn = box.xywhn[0]  # normalized center x, y, w, h
                cx, cy, w, h = xywhn.tolist()
                f.write(f"{cls} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n")

        labeled += 1
        fname = os.path.basename(filepath)
        conf = boxes.conf[0].item()
        print(f"  {fname}: {len(boxes)} box(es), best conf={conf:.2f}")

    print(f"\nPre-labeled {labeled} images, skipped {skipped} (no detections)")


def main():
    parser = argparse.ArgumentParser(description="Pre-label images with YOLO model")
    parser.add_argument("img_dir", nargs="?", default="../../training-data/kulemak/raw")
    parser.add_argument("--model", default="boss-kulemak", help="Model name in models/")
    parser.add_argument("--conf", type=float, default=0.20, help="Confidence threshold")
    args = parser.parse_args()
    run_prelabel(args)

if __name__ == "__main__":
    main()