#!/usr/bin/env python3

import os
import sys
import json
import random
import argparse
from functools import partial
from common import *

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-dir', required=True)
    parser.add_argument('--scan-id', help='Process only this specific scan ID')
    parser.add_argument('--scan-id-range', help='Process scan IDs in range (e.g., "358626-358630" or "358626-358630,359000-359010")')
    parser.add_argument('--force', action='store_true', help='Overwrite existing sbs.jpg files')
    return parser.parse_args()

def process_scan(scan_dir, force=False):
    if not os.path.isdir(scan_dir):
        return "scan_dir not found"
    frame_file = os.path.join(scan_dir, 'frame.jpg')
    std_file = os.path.join(scan_dir, 'std.jpg')
    if not os.path.exists(frame_file) or not os.path.exists(std_file):
        return "frame.jpg or std.jpg not found"
    sbs_file = os.path.join(scan_dir, 'sbs.jpg')
    try:
        if force or not os.path.exists(sbs_file):
            frame_img = Image.open(frame_file)
            std_img = Image.open(std_file)
            sbs_img = make_side_by_side_img_with_margins(frame_img, std_img)
            if sbs_img:
                sbs_img.save(sbs_file)
            else:
                return "make_side_by_side_img_with_margins failed"
        return "ok"
    except Exception as e:
        return f"error: {e}"

def main():
    args = parse_args()
    data_dir = args.data_dir
    scans_dir = os.path.join(data_dir, 'scans')
    scan_ids = os.listdir(scans_dir)
    if args.scan_id:
        if args.scan_id not in scan_ids:
            print(f"Error: scan-id '{args.scan_id}' not found")
            sys.exit(1)
        scan_ids = [args.scan_id]
    elif args.scan_id_range:
        ranges = parse_ranges(args.scan_id_range)
        filtered_scan_ids = []
        for scan_id in scan_ids:
            try:
                scan_id_int = int(scan_id)
                for val_range in ranges:
                    if in_range(scan_id_int, val_range):
                        filtered_scan_ids.append(scan_id)
                        break
            except ValueError:
                # Skip non-numeric scan IDs
                continue
        scan_ids = filtered_scan_ids
        if not scan_ids:
            print(f"Error: No scan IDs found in range '{args.scan_id_range}'")
            sys.exit(1)
    pool = Pool(cpu_count())
    counts = defaultdict(int)
    process_scan_with_force = partial(process_scan, force=args.force)
    for result in tqdm(pool.imap(process_scan_with_force, [os.path.join(scans_dir, scan_id) for scan_id in scan_ids]), total=len(scan_ids)):
        counts[result] += 1
    for k, v in counts.items():
        print(f"{k}: {v}")

if __name__ == '__main__':
    main()