#!/usr/bin/env python3
"""
Immich YOLO Photo Classifier Pipeline
Classifies photos as: has-people, scenic, or neither.
Moves non-people/non-scenic photos to NO PEOPLE PHOTOS folder.

Usage:
  /mnt/storage/yolo_venv/bin/python3 this_script.py

Requirements:
  - /mnt/storage/yolo_venv/ with ultralytics installed
  - ~/yolov8n_openvino_model/ (exported from yolov8n.pt)
  - Immich PostgreSQL container running
  - Photos on /mnt/wd-passport/immich/photos/
"""
import multiprocessing, os, time, warnings, subprocess, shutil
warnings.filterwarnings("ignore")

PHOTO_BASE = "/mnt/wd-passport/immich/photos"
NO_PEOPLE_DIR = "/mnt/wd-passport/immich/NO PEOPLE PHOTOS"
DB_QUERY = '''docker exec immich_postgres psql -U postgres -d immich -t -A -c "SELECT \\\"originalPath\\\" FROM asset WHERE type='IMAGE' AND visibility='timeline' AND \\\"originalPath\\\" LIKE '/data/library/%' ORDER BY \\\"originalPath\\\";" 2>/dev/null'''

PERSON_CLASS = 0
NATURE_CLASSES = {16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 58, 77, 80}
MOVE_CLASSES = {56, 57, 59, 60, 61, 62, 63, 64, 65, 67, 70, 71, 72, 73, 74, 75, 76, 78, 79,
                1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}

def worker(args):
    model_path, paths = args
    from ultralytics import YOLO
    import warnings
    warnings.filterwarnings("ignore")
    model = YOLO(model_path, task="detect")
    results = []
    for ap in paths:
        try:
            r = model(ap, verbose=False)
            has_person = has_nature = has_move = False
            for r2 in r:
                if r2.boxes:
                    for box in r2.boxes:
                        cls = int(box.cls[0])
                        if cls == PERSON_CLASS:           has_person = True
                        elif cls in NATURE_CLASSES:        has_nature = True
                        elif cls in MOVE_CLASSES:          has_move = True
            if has_person:        results.append((ap, "keep_people"))
            elif has_nature:      results.append((ap, "keep_scenic"))
            elif has_move:        results.append((ap, "move"))
            else:                 results.append((ap, "check_size"))
        except Exception as e:
            results.append((ap, f"error:{e}"))
    return results

if __name__ == "__main__":
    multiprocessing.set_start_method("forkserver", force=True)
    num_workers = min(multiprocessing.cpu_count(), 8)
    os.makedirs(NO_PEOPLE_DIR, exist_ok=True)
    model_path = os.path.expanduser("~/yolov8n_openvino_model/")

    print(f"[1] Querying DB...", flush=True)
    result = subprocess.run(DB_QUERY, shell=True, capture_output=True, text=True, timeout=120)
    db_paths = [p.strip() for p in result.stdout.strip().split('\n') if p.strip()]

    print(f"[2] Mapping {len(db_paths)} paths...", flush=True)
    actual = []
    for p in db_paths:
        ap = PHOTO_BASE + "/" + p[6:]
        if os.path.exists(ap):
            actual.append(ap)
    print(f"  {len(actual)} accessible files", flush=True)

    print(f"[3] Classifying with {num_workers} workers...", flush=True)
    chunks = [actual[i::num_workers] for i in range(num_workers)]
    chunk_args = [(model_path, c) for c in chunks]
    t_start = time.time()
    with multiprocessing.Pool(num_workers) as pool:
        all_results = pool.map(worker, chunk_args)
    elapsed = time.time() - t_start

    keep_people = keep_scenic = errors = 0
    to_move, check_size = [], []
    for chunk in all_results:
        for ap, d in chunk:
            if d == "keep_people":    keep_people += 1
            elif d == "keep_scenic":  keep_scenic += 1
            elif d == "move":         to_move.append(ap)
            elif d == "check_size":   check_size.append(ap)
            else:                     errors += 1

    print(f"\n  YOLO: {elapsed/60:.1f}min ({len(actual)/elapsed:.1f} img/s)", flush=True)
    print(f"  People:{keep_people} Scenic:{keep_scenic} Move:{len(to_move)} Check:{len(check_size)} Err:{errors}", flush=True)

    print(f"[4] Size check on {len(check_size)} undetected...", flush=True)
    for ap in check_size:
        try:
            if os.path.getsize(ap) < 400000:
                to_move.append(ap)
        except: pass
    print(f"  Now {len(to_move)} to move", flush=True)

    print(f"[5] Moving {len(to_move)} photos...", flush=True)
    moved = file_errors = 0
    for ap in to_move:
        try:
            rel = os.path.relpath(ap, PHOTO_BASE)
            dest = os.path.join(NO_PEOPLE_DIR, rel)
            os.makedirs(os.path.dirname(dest), exist_ok=True)
            shutil.move(ap, dest)
            moved += 1
        except Exception as e:
            file_errors += 1

    print(f"\nDONE! Moved {moved} photos", flush=True)
    print(f"  Total:{len(actual)} Kept(people):{keep_people} Kept(scenic):{keep_scenic} Moved:{moved} Errors:{file_errors}", flush=True)
    print(f"  Time: {(time.time()-t_start)/60:.1f}min", flush=True)
    print(f"  Output: {NO_PEOPLE_DIR}", flush=True)