#!/usr/bin/env python3
"""Extract OB sign_crops.zip to fs + update manifest to storage='fs'."""
import json, argparse, zipfile, os, sys
from pathlib import Path

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--manifest', default=str(ROOT / 'datasets/sources/old_babylonian_signs/manifest.jsonl'))
    ap.add_argument('--archive', default=str(ROOT / 'datasets/sources/old_babylonian_signs/sign_crops.zip'))
    ap.add_argument('--extract-dir', default=str(ROOT / 'datasets/sources/old_babylonian_signs/sign_crops'))
    ap.add_argument('--output-manifest', default=str(ROOT / 'datasets/sources/old_babylonian_signs/manifest_fs.jsonl'))
    args = ap.parse_args()

    # Extract zip (idempotent)
    ed = Path(args.extract_dir)
    if not ed.exists() or len(list(ed.rglob('*.jpg'))) == 0:
        print(f"Extracting {args.archive} → {ed}")
        ed.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(args.archive) as z:
            z.extractall(ed.parent)  # extracts as sign_crops/
    n_extracted = sum(1 for _ in ed.rglob('*.jpg')) + sum(1 for _ in ed.rglob('*.png'))
    print(f"Images on disk: {n_extracted}")

    # Rewrite manifest
    n_fs = 0; n_skip = 0
    with open(args.manifest) as fin, open(args.output_manifest, 'w') as fout:
        for line in fin:
            r = json.loads(line)
            if r.get('storage') == 'zip' and r.get('member'):
                fs_path = str(ed.parent / r['member'])
                if os.path.exists(fs_path):
                    r['storage'] = 'fs'
                    r['path'] = fs_path
                    r['integrity_ok'] = True
                    fout.write(json.dumps(r) + '\n'); n_fs += 1
                else:
                    n_skip += 1
            else:
                fout.write(line); n_fs += 1
    print(f"Converted to fs: {n_fs}, missing files: {n_skip}")
    print(f"Output: {args.output_manifest}")

if __name__ == '__main__':
    main()