#!/usr/bin/env python3
"""
Restore missing hitit_ocr/data/classification/all/<sign>/<tid>_<idx>.png crops.

Legacy prepare_data.py (deleted runs) produced crops from datasets/sources/hitit_local/<tid>/{image}+mark.txt
using enumerate(info["signs"]) with info["signs"] = deduplicate(parse_annotations_format(data)).
The manifest_v13_ultimate.jsonl records reference those exact filenames. We parse mark.txt the same way,
recompute bbox list, and write crops back to hitit_ocr/data/classification/all/<unified_label>/<tid>_<idx>.png.

Usage:
    python restore_hitit_crops.py --check TID      # dry-run one tablet
    python restore_hitit_crops.py --verify         # compare parsed bbox list vs manifest indices
    python restore_hitit_crops.py                  # full restore
"""
import argparse
import json
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path
from PIL import Image, ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True

MIN_CROP_SIZE = 8
CONTEXT_PAD_RATIO = 0.15
TINY_BBOX_THRESHOLD = 0.005

PROJECT = Path("/arf/scratch/stakan/hitit-proje")
SOURCES = PROJECT / "datasets/sources/hitit_local"
OUT_ALL = PROJECT / "hitit_ocr/data/classification/all"
MANIFEST = SOURCES / "manifest_v13_ultimate.jsonl"


def normalize_sign(s):
    """Mirror of legacy prepare_data.py normalize_sign (ham: sadece trim + sayı/rare filtresi)."""
    if s is None:
        return None
    s = s.strip()
    s = s.rstrip(".,;: ")
    s = s.strip()
    if not s or s in ["/", ".", ",", "-", "(X)", "(x)", "X", "x", "?", ""]:
        return None
    if s.isdigit():
        return None
    if len(s) == 1 and not s.isalpha():
        return None
    return s


def parse_annotations(data):
    results = []
    for ann in data.get("annotations", []):
        sign = normalize_sign(ann.get("comment", ""))
        if not sign:
            continue
        m = ann.get("mark", {})
        x, y, w, h = m.get("x", 0), m.get("y", 0), m.get("width", 0), m.get("height", 0)
        if w < 0:
            x += w
            w = abs(w)
        if h < 0:
            y += h
            h = abs(h)
        if w > MIN_CROP_SIZE and h > MIN_CROP_SIZE:
            results.append((x, y, w, h, sign))
    return dedup(results)


def parse_spots(data, img_w, img_h):
    results = []
    for sp in data.get("spots", []):
        title = sp.get("title", "")
        if "_" not in title:
            continue
        sign = normalize_sign(title.split("_", 1)[1])
        if not sign:
            continue
        x = sp.get("x", 0) / 100 * img_w
        y = sp.get("y", 0) / 100 * img_h
        w = sp.get("width", 0) / 100 * img_w
        h = sp.get("height", 0) / 100 * img_h
        if w > MIN_CROP_SIZE and h > MIN_CROP_SIZE:
            results.append((x, y, w, h, sign))
    return dedup(results)


def dedup(signs):
    seen = set()
    out = []
    for e in signs:
        k = (round(e[0], 2), round(e[1], 2), round(e[2], 2), round(e[3], 2), e[4])
        if k not in seen:
            seen.add(k)
            out.append(e)
    return out


def load_tablet_signs(tid):
    """Return (image_path, width, height, signs[]) for tablet dir <tid>."""
    tdir = SOURCES / tid
    if not tdir.is_dir():
        return None
    mark = tdir / "mark.txt"
    if not mark.exists() or mark.stat().st_size == 0:
        return None
    img_path = None
    for f in sorted(tdir.iterdir()):
        if f.suffix.lower() in (".jpg", ".jpeg", ".png") and not f.is_symlink():
            img_path = f
            break
        if f.suffix.lower() in (".jpg", ".jpeg", ".png"):
            img_path = f.resolve()
            break
    if not img_path or not img_path.exists():
        return None
    try:
        with Image.open(img_path) as im:
            w, h = im.size
    except Exception:
        return None
    try:
        data = json.loads(mark.read_text())
    except Exception:
        return None
    if "annotations" in data:
        signs = parse_annotations(data)
    elif "spots" in data:
        signs = parse_spots(data, w, h)
    else:
        return None
    return str(img_path), w, h, signs


def safe_name(s):
    for a, b in [("/", "_SLASH_"), ("\\", "_BSLASH_"), (":", "_COLON_"),
                 ("*", "_STAR_"), ("?", "_QMARK_"), ('"', "_DQUOTE_"),
                 ("<", "_LT_"), (">", "_GT_"), ("|", "_PIPE_"), (" ", "_")]:
        s = s.replace(a, b)
    return s if s else "_EMPTY_"


def manifest_records_by_tablet():
    """Group missing hitit manifest entries by tablet_id. Only include those under classification/all/."""
    by_tid = defaultdict(list)
    with open(MANIFEST) as f:
        for line in f:
            r = json.loads(line)
            if r.get("source") != "hitit":
                continue
            p = r.get("path", "")
            if "/classification/all/" not in p:
                continue
            # Derive tid from filename (more reliable than r["tablet_id"]
            # which has known mislabels, e.g. path=116.2_26.png with tablet_id="4").
            stem = Path(p).stem  # e.g. 116.2_26
            parts = stem.rsplit("_", 1)
            if len(parts) == 2 and parts[1].isdigit():
                tid = parts[0]
            else:
                tid = r.get("tablet_id") or stem
            by_tid[tid].append(r)
    return by_tid


def restore_tablet(tid, records, verify_only=False):
    """For a tablet, parse bbox list, then write each record's crop.
    Returns dict of stats."""
    t = load_tablet_signs(tid)
    if t is None:
        return {"tid": tid, "status": "no_tablet", "n_records": len(records)}
    img_path, w, h, signs = t
    # Build filename → idx map from records
    # File name pattern: <tid>_<idx>.png
    stats = {"tid": tid, "n_signs": len(signs), "n_records": len(records),
             "written": 0, "skipped_label_mismatch": 0, "idx_out_of_range": 0,
             "already_exists": 0, "save_failed": 0}
    if verify_only:
        # Check idx vs manifest label/size
        by_idx = {}
        for r in records:
            fn = Path(r["path"]).stem  # e.g. 100.2_116
            parts = fn.rsplit("_", 1)
            if len(parts) != 2 or not parts[1].isdigit():
                continue
            idx = int(parts[1])
            by_idx[idx] = r
        matched = 0
        sz_mismatch = 0
        label_mismatch_examples = []
        for idx, r in by_idx.items():
            if idx >= len(signs):
                stats["idx_out_of_range"] += 1
                continue
            x, y, bw, bh, sign = signs[idx]
            # Expected pad
            px = bw * CONTEXT_PAD_RATIO
            py = bh * CONTEXT_PAD_RATIO
            x1 = max(0, int(x - px)); y1 = max(0, int(y - py))
            x2 = min(w, int(x + bw + px)); y2 = min(h, int(y + bh + py))
            expect_w = x2 - x1; expect_h = y2 - y1
            man_w = r.get("width"); man_h = r.get("height")
            if man_w != expect_w or man_h != expect_h:
                sz_mismatch += 1
            matched += 1
        stats["verify_matched"] = matched
        stats["verify_size_mismatch"] = sz_mismatch
        return stats

    # Open image once
    try:
        img = Image.open(img_path).convert("RGB")
    except Exception:
        return {**stats, "status": "img_open_failed"}
    for r in records:
        fn = Path(r["path"]).stem
        parts = fn.rsplit("_", 1)
        if len(parts) != 2 or not parts[1].isdigit():
            continue
        idx = int(parts[1])
        if idx >= len(signs):
            stats["idx_out_of_range"] += 1
            continue
        x, y, bw, bh, sign = signs[idx]
        px = bw * CONTEXT_PAD_RATIO
        py = bh * CONTEXT_PAD_RATIO
        x1 = max(0, int(x - px))
        y1 = max(0, int(y - py))
        x2 = min(w, int(x + bw + px))
        y2 = min(h, int(y + bh + py))
        if x2 - x1 < MIN_CROP_SIZE or y2 - y1 < MIN_CROP_SIZE:
            continue
        out = Path(r["path"])
        if out.exists():
            stats["already_exists"] += 1
            continue
        out.parent.mkdir(parents=True, exist_ok=True)
        try:
            img.crop((x1, y1, x2, y2)).save(str(out))
            stats["written"] += 1
        except Exception as e:
            stats["save_failed"] += 1
    img.close()
    return stats


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--check", help="dry-run one tablet")
    ap.add_argument("--verify", action="store_true", help="parse all tablets, compare bbox sizes vs manifest")
    ap.add_argument("--limit", type=int, default=0, help="process at most N tablets")
    ap.add_argument("--workers", type=int, default=8)
    args = ap.parse_args()

    by_tid = manifest_records_by_tablet()
    tids = sorted(by_tid.keys())
    print(f"[plan] {len(tids)} tablets, {sum(len(v) for v in by_tid.values())} crops to restore")

    if args.check:
        tid = args.check
        t = load_tablet_signs(tid)
        if t is None:
            print(f"tablet {tid} not loadable"); return
        img_path, w, h, signs = t
        print(f"tablet {tid}: img={img_path} size={w}x{h} bbox_count={len(signs)}")
        print("first 3 bboxes:", signs[:3])
        print("records for this tablet:", len(by_tid.get(tid, [])))
        if by_tid.get(tid):
            print("first record:", json.dumps(by_tid[tid][0], ensure_ascii=False)[:200])
        return

    if args.verify:
        agg = Counter()
        for i, tid in enumerate(tids):
            st = restore_tablet(tid, by_tid[tid], verify_only=True)
            agg["tablets"] += 1
            for k in ("n_records", "verify_matched", "verify_size_mismatch", "idx_out_of_range"):
                agg[k] += st.get(k, 0)
            if args.limit and i + 1 >= args.limit:
                break
        print("VERIFY:", dict(agg))
        return

    # Full restore
    import multiprocessing as mp
    work = [(tid, by_tid[tid]) for tid in tids]
    if args.limit:
        work = work[:args.limit]
    agg = Counter()
    if args.workers <= 1:
        for i, (tid, recs) in enumerate(work):
            st = restore_tablet(tid, recs)
            for k, v in st.items():
                if isinstance(v, int):
                    agg[k] += v
            if (i + 1) % 20 == 0:
                print(f"  [{i+1}/{len(work)}] written={agg['written']} skipped_idx={agg['idx_out_of_range']}")
    else:
        with mp.Pool(args.workers) as pool:
            for i, st in enumerate(pool.starmap(restore_tablet, work)):
                for k, v in st.items():
                    if isinstance(v, int):
                        agg[k] += v
                if (i + 1) % 20 == 0:
                    print(f"  [{i+1}/{len(work)}] written={agg['written']} skipped_idx={agg['idx_out_of_range']}")
    print("DONE:", dict(agg))


if __name__ == "__main__":
    main()