#!/usr/bin/env python3
"""Multimodal-Mind2Web hazırlığı: her adım için ilk-K aday listesi (görüntüsüz).

Çıktı (her split için bir jsonl.gz): adım başına görev, geçmiş eylemler, işlem türü,
adaylar (tag, kısa metin, kutu; K=0 ise hepsi, TF-IDF sırasıyla), doğru adayın sırası. Ekran görüntüleri BURADA
kaydedilmez; öznitelik çıkarırken akıştan tekrar okunur.

Buradaki sıralama yalnızca ön sıralama (TF-IDF). Gerçek ilk-50 süzgeci m2w_ranker.py'deki DeBERTa.
Gerçek bir satırda TF-IDF doğru öğeyi 554 adayın 349. sırasına koydu: tek başına yetmez.

Kullanım:
    python m2w_prep.py --splits train --max_rows 300 --out m2w_smoke
    python m2w_prep.py --splits train test_task test_website test_domain --out m2w_all
"""
import argparse
import gzip
import itertools
import json
import os
import re
import zlib
from collections import Counter

import numpy as np
from lxml import etree
from sklearn.feature_extraction.text import TfidfVectorizer

NAME = "osunlp/Multimodal-Mind2Web"
COLS = ["action_uid", "raw_html", "cleaned_html", "operation", "pos_candidates", "neg_candidates",
        "website", "domain", "subdomain", "annotation_id", "confirmed_task",
        "action_reprs", "target_action_index", "target_action_reprs"]
LABEL_ATTRS = ("aria_label", "placeholder", "title", "name", "alt", "value", "id")
# Yapısal alanlar (aday biçimi denemeleri için): metin öznitelikleri boş değilse, mantıksal öznitelikler varsa (boş değerle bile) saklanır
TEXT_AT = ("role", "type", "aria_label", "aria-label", "placeholder", "title", "name", "alt", "value", "id",
           "aria_checked", "aria-checked", "aria_expanded", "aria-expanded", "aria_disabled", "aria-disabled",
           "aria_selected", "aria-selected", "aria_haspopup", "aria-haspopup")
BOOL_AT = ("checked", "disabled", "readonly", "required", "selected")
ATTR_KEYS = Counter()   # doğru adaylarda hangi öznitelik anahtarları görülüyor (özet için)


def small_attrs(a):
    out = {}
    for k in TEXT_AT:
        v = a.get(k)
        if v is not None and str(v).strip():
            out[k] = str(v).strip()[:120]
    for k in BOOL_AT:
        if k in a:
            out[k] = str(a[k])[:10]
    return out


def _load(c):
    return json.loads(c) if isinstance(c, str) else c


def parse_cand(c):
    c = _load(c)
    a = _load(c["attributes"]) if c.get("attributes") else {}
    bbox = None
    r = a.get("bounding_box_rect")
    if r:
        try:
            v = [float(x) for x in (r.split(",") if isinstance(r, str) else r)]
            if len(v) == 4:
                bbox = v
        except ValueError:
            pass
    return {"tag": c.get("tag", ""), "bid": str(c.get("backend_node_id", a.get("backend_node_id", ""))),
            "attrs": a, "bbox": bbox, "orig": bool(c.get("is_original_target", False))}


def node_texts(cleaned_html, bids):
    """backend_node_id -> kısa iç metin (alt ağaçtan, ilk birkaç parça)."""
    out = {}
    if not cleaned_html or not bids:
        return out
    try:
        root = etree.fromstring(cleaned_html.encode("utf-8", "ignore"),
                                etree.HTMLParser(recover=True, huge_tree=True))
    except Exception:
        return out
    if root is None:
        return out
    for el in root.iter():
        b = el.get("backend_node_id")
        if b in bids:
            t = " ".join("".join(itertools.islice(el.itertext(), 40)).split())
            out[b] = t[:200]
    return out


def cand_text(cd, inner):
    a = cd["attrs"]
    parts = [cd["tag"]]
    for k in ("role", "type"):
        if a.get(k):
            parts.append(str(a[k]))
    for k in LABEL_ATTRS:
        v = a.get(k)
        if v and str(v).strip():
            parts.append(str(v).strip()[:60])
    if inner:
        parts.append(inner)
    return " | ".join(parts)[:160]


def page_size(raw_html):
    m = re.search(r'bounding_box_rect="([^"]+)"', raw_html[:4000])
    if m:
        v = m.group(1).split(",")
        if len(v) == 4:
            return float(v[2]), float(v[3])
    return None, None


def rank(query, texts, bboxes, drop_zero_area):
    """Benzerlik puanı ile sıralı indeksler. Alanı sıfır olan öğeler en sona."""
    n = len(texts)
    if n == 0:
        return []
    try:
        vec = TfidfVectorizer(token_pattern=r"[A-Za-z0-9]{2,}", lowercase=True, sublinear_tf=True)
        M = vec.fit_transform(texts + [query])
        sc = (M[:n] @ M[n].T).toarray().ravel()
    except ValueError:
        sc = np.zeros(n)
    sc = sc.astype(float)
    for i, b in enumerate(bboxes):
        if b is None or b[2] <= 0 or b[3] <= 0:
            sc[i] = -1.0 if drop_zero_area else sc[i] - 0.05
    return list(np.argsort(-sc, kind="stable"))


def process_row(r, K, drop_zero_area):
    pos = [parse_cand(c) for c in r["pos_candidates"]]
    neg = [parse_cand(c) for c in r["neg_candidates"]]
    allc = pos + neg
    bids = {c["bid"] for c in allc}
    texts_in = node_texts(r["cleaned_html"], bids)
    texts = [cand_text(c, texts_in.get(c["bid"], "")[:100]) for c in allc]   # eski metin: iç metin 100 karakter
    idx = int(r["target_action_index"])
    hist = list(r["action_reprs"][:idx])
    query = r["confirmed_task"] + " " + " ".join(re.sub(r"->.*", "", h) for h in hist[-3:])
    order = rank(query, texts, [c["bbox"] for c in allc], drop_zero_area)
    if K:
        order = order[:K]
    pos_b = {c["bid"] for c in pos}
    for c in pos:
        ATTR_KEYS.update(c["attrs"].keys())
    cands, pos_idx = [], []
    for j, i in enumerate(order):
        c = allc[i]
        cands.append({"bid": c["bid"], "tag": c["tag"], "text": texts[i], "bbox": c["bbox"],
                      "inner": texts_in.get(c["bid"], ""), "at": small_attrs(c["attrs"])})
        if c["bid"] in pos_b:
            pos_idx.append(j)
    pw, ph = page_size(r["raw_html"])
    op = _load(r["operation"])   # gerçek veride JSON metni
    return {
        "action_uid": r["action_uid"], "ann": r["annotation_id"], "step": idx,
        "website": r["website"], "domain": r["domain"], "task": r["confirmed_task"],
        "history": hist, "target": r["target_action_reprs"],
        "op": op["op"], "value": op.get("value", ""),
        "page_w": pw, "page_h": ph, "n_total": len(allc), "n_pos": len(pos),
        "cands": cands, "pos_idx": pos_idx,
    }


def open_stream(split):
    from datasets import load_dataset
    try:
        return load_dataset(NAME, split=split, streaming=True, columns=COLS)
    except Exception as e:  # columns desteklenmiyorsa
        print("columns= çalışmadı, tüm kolonlarla akış:", type(e).__name__)
        ds = load_dataset(NAME, split=split, streaming=True)
        return ds.remove_columns([c for c in ("screenshot",) if c])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--splits", nargs="+", default=["train"])
    ap.add_argument("--out", default="m2w")
    ap.add_argument("--K", type=int, default=0, help="0 = tüm adaylar")
    ap.add_argument("--max_rows", type=int, default=0)
    ap.add_argument("--val_mod", type=int, default=8, help="eğitim sitelerinin 1/val_mod kadarı doğrulamaya")
    ap.add_argument("--keep_zero_area", action="store_true")
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)

    for split in args.splits:
        ds = open_stream(split)
        if args.max_rows:
            ds = ds.take(args.max_rows)
        outs = {}
        st = Counter()
        n_cand = []
        for k, r in enumerate(ds):
            ex = process_row(r, args.K, not args.keep_zero_area)
            name = split
            if split == "train":
                name = "val" if zlib.crc32(ex["website"].encode()) % args.val_mod == 0 else "train"
            if name not in outs:
                outs[name] = gzip.open(os.path.join(args.out, f"{name}.jsonl.gz"), "wt", encoding="utf-8")
            ex["split"] = name
            outs[name].write(json.dumps(ex, ensure_ascii=False) + "\n")
            st[f"{name}/n"] += 1
            st[f"{name}/op_{ex['op']}"] += 1
            st[f"{name}/pos_yok"] += ex["n_pos"] == 0
            st[f"{name}/pos_ilkK_icinde"] += len(ex["pos_idx"]) > 0
            st[f"{name}/pos_ilk1"] += 0 in ex["pos_idx"]
            n_cand.append(ex["n_total"])
            if (k + 1) % 500 == 0:
                print(f"  {split}: {k + 1} adım", flush=True)
        for f in outs.values():
            f.close()
        print(f"\n== {split} ==  adaylar/adım: ortanca {int(np.median(n_cand))}, en çok {max(n_cand)}")
        for name in sorted({k.split('/')[0] for k in st}):
            n = st[f"{name}/n"]
            print(f"  [{name}] adım={n} | doğru aday listede yok={st[name+'/pos_yok']} "
                  f"| recall@{args.K or 'hepsi'}={st[name+'/pos_ilkK_icinde']/n:.3f} | top1={st[name+'/pos_ilk1']/n:.3f} "
                  f"| CLICK={st[name+'/op_CLICK']} TYPE={st[name+'/op_TYPE']} SELECT={st[name+'/op_SELECT']}")
    print("\ndoğru adaylarda görülen öznitelik anahtarları (ilk 40):", ATTR_KEYS.most_common(40))
    json.dump(dict(ATTR_KEYS.most_common(200)), open(os.path.join(args.out, "oznitelik_anahtarlari.json"), "w"), indent=1)
    print("\nyazıldı:", args.out)


if __name__ == "__main__":
    main()