#!/usr/bin/env python3
"""Letterbox 224 preprocessed crop'ları WebDataset shard'larına yaz.

Config-driven preprocessing.yaml'dan okur.
Pipeline: read → CLAHE (conditional) → gamma → letterbox 224 → JPEG
MSII proxy optional 2. shard olarak.
"""
import json, os, io, time, argparse
from pathlib import Path
from concurrent.futures import ProcessPoolExecutor
import sys

ROOT = Path("/arf/scratch/stakan/hitit-proje")
SOURCES = ROOT / "datasets" / "sources"
sys.path.insert(0, str(ROOT / "hitit_ocr" / "src"))

def process_one(item):
    """Tek image: preprocess et, JPEG bytes döndür."""
    rid, path, label, source, fold = item
    try:
        import numpy as np
        from PIL import Image
        import cv2
        
        from preprocessing.pipeline import load_config, apply_clahe, apply_gamma, letterbox, is_low_quality, msii_proxy
        
        cfg = load_config()
        
        img = np.array(Image.open(path).convert('RGB'))
        enh = cfg.get('enhancement', {})
        if enh.get('clahe', {}).get('enabled') and (not enh['clahe'].get('conditional') or is_low_quality(img)):
            img = apply_clahe(img, clip_limit=enh['clahe'].get('clip_limit', 2.5))
        if enh.get('gamma_correction', {}).get('enabled') and (not enh['gamma_correction'].get('conditional') or is_low_quality(img)):
            img = apply_gamma(img, gamma=enh['gamma_correction'].get('gamma', 1.2))
        
        lb = cfg.get('geometric', {}).get('letterbox', {})
        img = letterbox(img, target=lb.get('target_size', 224), margin_ratio=lb.get('margin_ratio', 0.1))
        
        # Encode JPEG
        buf = io.BytesIO()
        Image.fromarray(img).save(buf, format='JPEG', quality=90)
        jpeg_bytes = buf.getvalue()
        
        return (rid, jpeg_bytes, label, source, fold)
    except Exception as e:
        return (rid, None, None, None, None)

def main():
    import webdataset as wds
    
    ap = argparse.ArgumentParser()
    ap.add_argument('--subset', default='curated', help='curated|all|hitit')
    ap.add_argument('--workers', type=int, default=64)
    ap.add_argument('--shard-size', type=int, default=2000)
    ap.add_argument('--out', default=str(ROOT / "datasets" / "streaming" / "preprocessed_224"))
    args = ap.parse_args()
    
    # Topla
    items = []
    for d in sorted(SOURCES.iterdir()):
        if not d.is_dir(): continue
        mp = d / "manifest_classification.jsonl"
        if not mp.exists(): continue
        with open(mp) as f:
            for line in f:
                r = json.loads(line)
                if r.get('storage') != 'fs' or not r.get('path') or not os.path.exists(r['path']): continue
                if r.get('quality_gate_pass') is False: continue
                if args.subset == 'curated' and not r.get('in_curated_pretrain'): continue
                if args.subset == 'hitit' and r.get('source') != 'hitit': continue
                items.append((r['id'], r['path'], r.get('unified_label',''), r.get('source',''), r.get('fold', 0)))
    print(f"İşlenecek: {len(items):,}", flush=True)
    
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    pattern = str(out_dir / "preprocessed-%06d.tar")
    
    sink = wds.ShardWriter(pattern, maxcount=args.shard_size, encoder=False)
    t0 = time.time()
    done = 0
    
    try:
        with ProcessPoolExecutor(max_workers=args.workers) as ex:
            for res in ex.map(process_one, items, chunksize=50):
                rid, jpeg_bytes, label, source, fold = res
                if jpeg_bytes:
                    sink.write({
                        '__key__': rid,
                        'jpeg': jpeg_bytes,
                        'json': json.dumps({
                            'id': rid, 'unified_label': label,
                            'source': source, 'fold': fold,
                        }, ensure_ascii=False).encode('utf-8'),
                    })
                    done += 1
                if done % 5000 == 0 and done > 0:
                    rate = done / max(time.time() - t0, 1)
                    print(f"  {done:,}/{len(items):,} ({rate:.0f} img/s)", flush=True)
    finally:
        sink.close()
    
    # Summary
    shards = sorted(out_dir.glob("preprocessed-*.tar"))
    total_size = sum(s.stat().st_size for s in shards)
    print(f"\n{len(shards)} shard, {done:,} kayıt, {total_size/1024/1024:.1f} MB")
    
    meta = {
        "format": "WebDataset (preprocessed letterbox 224, CLAHE+gamma conditional)",
        "n_shards": len(shards),
        "n_records": done,
        "total_size_mb": round(total_size/1024/1024, 1),
        "pattern": f"preprocessed-{{000000..{len(shards)-1:06d}}}.tar",
        "config_version": "1.0",
        "preprocessing_steps": ["CLAHE (conditional)", "Gamma (conditional)", "Letterbox 224 median-fill"],
    }
    with open(out_dir / "shards.json", 'w') as f:
        json.dump(meta, f, indent=2, ensure_ascii=False)

if __name__ == '__main__':
    main()