#!/usr/bin/env python3
"""Build SSL continual-pretrain manifest from unlabeled cuneiform tablet images.

Sources:
  - cuneiml (83k Ur-III tablet photos + lineart)
  - heicubeda (11.8k Heidelberg tablet images)
  - tlhdig tablet images (if image files are extracted)
  - hitit_local raw tablets
  - transliterated_fragments tablet images

Output: JSONL with {path, source, width, height} per image — no labels, no
cross-source filtering. Used for DINOv3 continual pretraining (masked-image
modeling / DINOv2-style SSL).
"""
import json, argparse, os
from pathlib import Path
from collections import Counter

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


def iter_source(src_name, src_dir, min_size=96):
    mf = src_dir / 'manifest.jsonl'
    if not mf.exists():
        return
    for line in open(mf):
        try: r = json.loads(line)
        except: continue
        if r.get('storage') != 'fs': continue
        p = r.get('path')
        if not p or not os.path.exists(p): continue
        w = r.get('width') or 0
        h = r.get('height') or 0
        # Skip degenerate tiny crops that SSL can't learn from
        if w and h and (w < min_size or h < min_size): continue
        if r.get('integrity_ok') is False: continue
        yield {
            'path': p,
            'source': src_name,
            'width': w,
            'height': h,
            'period': r.get('period') or 'unknown',
        }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--output', required=True)
    ap.add_argument('--min-size', type=int, default=96)
    ap.add_argument('--include-hitit', action='store_true', default=True,
                    help='Also include hitit_local + transliterated_fragments')
    args = ap.parse_args()

    sources = [
        ('cuneiml', ROOT / 'datasets/sources/cuneiml'),
        ('heicubeda', ROOT / 'datasets/sources/heicubeda'),
    ]
    if args.include_hitit:
        sources += [
            ('hitit_local', ROOT / 'datasets/sources/hitit_local'),
            ('transliterated_fragments', ROOT / 'datasets/sources/transliterated_fragments'),
        ]

    stats = Counter()
    with open(args.output, 'w') as out:
        for name, dir_ in sources:
            n = 0
            for rec in iter_source(name, dir_, min_size=args.min_size):
                out.write(json.dumps(rec) + '\n')
                n += 1
            stats[name] = n
            print(f"  {name}: {n} images")

    total = sum(stats.values())
    print(f"\nSSL manifest total: {total}")
    print(f"Saved → {args.output}")


if __name__ == '__main__':
    main()