#!/usr/bin/env python3
"""
Build tablet-level structured manifest from mark.txt annotations.

Each tablet → list of lines → list of signs with:
  - bbox (x, y, w, h) normalized
  - sign label (ABZ/comment)
  - language (Hittite/Sumerian/Akkadian/...)
  - damage (inferred: 'intact' | 'partial' | 'broken')
  - reading order (row_no, col_no, x-sort within line)

Output: datasets/sources/hitit_local/tablet_structure.jsonl
  {tablet_id, image_path, width, height, lines: [{row, col, signs:[...]}]}
"""
import argparse
import json
from collections import defaultdict
from pathlib import Path

from PIL import Image

SOURCES = Path("/arf/scratch/stakan/hitit-proje/datasets/sources/hitit_local")


LANG_MAP = {
    "Hititçe": "hit", "Hittite": "hit", "Hit": "hit",
    "Sümerce": "sum", "Sumerian": "sum", "Sum": "sum",
    "Akadca": "akk", "Akkadian": "akk", "Akk": "akk",
    "Luvice": "luw", "Luwian": "luw",
    "Hurrice": "hur", "Hurrian": "hur",
    "Hattice": "hat",
    "Palaca": "pal", "Palaic": "pal",
    "": "unk",
}


# Damage inference from comment
DAMAGE_MARKS = {
    'x': 'broken',   # 'x' standard Assyriology marker for lost sign
    'X': 'broken',
    '?': 'uncertain',
    '▒': 'broken',
    '…': 'broken',
    '[': 'partial',  # half-bracket usually
    ']': 'partial',
    '⸢': 'partial',
    '⸣': 'partial',
}


def classify_damage(comment, quality=None):
    """Return 'intact' | 'partial' | 'broken' | 'uncertain'."""
    if comment is None:
        return 'broken'
    c = comment.strip().rstrip(",.")
    if not c or c in {"x", "X", "_SLASH__SLASH_", "//", "▒", "…"}:
        return 'broken'
    if any(m in c for m in ('[', ']', '⸢', '⸣', '〈', '〉')):
        return 'partial'
    if '?' in c:
        return 'uncertain'
    return 'intact'


def normalize_bbox(mark, W, H):
    x = mark.get("x", 0); y = mark.get("y", 0)
    w = mark.get("width", 0); h = mark.get("height", 0)
    if w < 0: x += w; w = abs(w)
    if h < 0: y += h; h = abs(h)
    return x / W, y / H, w / W, h / H


def process_tablet(tid):
    tdir = SOURCES / tid
    mark = tdir / "mark.txt"
    if not mark.exists() or mark.stat().st_size == 0:
        return None
    # Find image
    img_path = None
    for ext in (".JPG", ".jpg", ".jpeg", ".png", ".PNG"):
        for f in tdir.glob(f"*{ext}"):
            img_path = f
            break
        if img_path: break
    if not img_path: 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
    anns = data.get("annotations", [])
    if not anns:
        return None

    # Build sign list with reading-order metadata; dedup by (pk_id, bbox)
    signs = []
    seen_pk = set()
    for idx, a in enumerate(anns):
        pk = a.get("pk_id")
        if pk is not None and pk in seen_pk:
            continue
        m = a.get("mark", {})
        bx, by, bw, bh = normalize_bbox(m, W, H)
        comment = (a.get("comment") or "").strip().rstrip(",.")
        lang_raw = a.get("lang") or ""
        signs.append({
            "idx": idx,
            "sign": comment if comment else "x",
            "lang": LANG_MAP.get(lang_raw, "unk"),
            "lang_raw": lang_raw,
            "damage": classify_damage(comment),
            "bbox": [round(bx, 4), round(by, 4), round(bw, 4), round(bh, 4)],
            "row": a.get("row_no", 0),
            "col": a.get("col_no", 0),
            "pk": pk,
        })
        if pk is not None:
            seen_pk.add(pk)

    # Group by (row, col), sort within row by x (reading order)
    by_line = defaultdict(list)
    for s in signs:
        by_line[(s["row"], s["col"])].append(s)
    lines = []
    for (row, col), items in sorted(by_line.items()):
        items.sort(key=lambda s: s["bbox"][0])  # left → right
        lines.append({"row": row, "col": col, "signs": items})

    return {
        "tablet_id": tid,
        "image_path": str(img_path),
        "width": W,
        "height": H,
        "n_signs": len(signs),
        "n_lines": len(lines),
        "lang_dist": dict((k, sum(1 for s in signs if s["lang"] == k))
                          for k in set(s["lang"] for s in signs)),
        "damage_dist": dict((k, sum(1 for s in signs if s["damage"] == k))
                             for k in set(s["damage"] for s in signs)),
        "lines": lines,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--output", default=str(SOURCES / "tablet_structure.jsonl"))
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args()

    tids = sorted(d.name for d in SOURCES.iterdir()
                   if d.is_dir() and (d / "mark.txt").exists())
    print(f"Scanning {len(tids)} tablets")
    n = 0
    with open(args.output, "w") as out:
        for tid in tids:
            t = process_tablet(tid)
            if t is None:
                continue
            out.write(json.dumps(t, ensure_ascii=False) + "\n")
            n += 1
            if args.limit and n >= args.limit:
                break
    print(f"DONE: {n} tablets → {args.output}")


if __name__ == "__main__":
    main()