#!/usr/bin/env python3

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

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-dir', required=True)
    return parser.parse_args()

def process_scan(scan_dir):
    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')
    frame_qr_file = os.path.join(scan_dir, 'frame-qr.jpg')
    std_qr_file = os.path.join(scan_dir, 'std-qr.jpg')
    sbs_no_margin_file = os.path.join(scan_dir, 'sbs-nomargin.jpg')
    try:
        if 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"
        if not os.path.exists(sbs_no_margin_file):
            frame_img = Image.open(frame_file)
            std_img = Image.open(std_file)
            if not os.path.exists(frame_qr_file) or not os.path.exists(std_qr_file):
                frame_qrcode, frame_qr_img = extract_qr(frame_img)
                std_qrcode, std_qr_img = extract_qr(std_img)
                frame_qr_img.save(frame_qr_file)
                std_qr_img.save(std_qr_file)
            else:
                frame_qr_img = Image.open(frame_qr_file)
                std_qr_img = Image.open(std_qr_file)
            sbs_no_margin_img = make_side_by_side_img(frame_qr_img, std_qr_img)
            sbs_no_margin_img.save(sbs_no_margin_file)
        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')
    pool = Pool(cpu_count())
    scan_ids = os.listdir(scans_dir)
    counts = defaultdict(int)
    for result in tqdm(pool.imap(process_scan, [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()