import json
import subprocess
import sys
from pathlib import Path
import torch
from glint.joint_schema_model import ClefModel, JointSchemaHead, collate_records, encode_record
from glint.schema import GlintEncodedRecord, encode_glint_record, image_token_coords, BOX_OUTPUT_TEXT
from glint.glint_model import GlintModel
from glint.box import BoxHead
from glint.data import derive_presence_gold, normalize_box_gold, normalize_teacher, teacher_questions, load_jsonl
from glint.losses import decision_loss, record_box_loss

ROOT = Path(__file__).resolve().parents[1]


def test_no_box_is_bit_identical_to_clef(tok, records):
    a = encode_record(tok, records["choice"]); b = encode_glint_record(tok, records["choice"])
    assert type(b) is type(a) and b == a
    c = encode_glint_record(tok, records["choice"], force_glint_path=True)   # the mirrored encoder
    assert c.input_ids == a.input_ids and c.questions == a.questions and c.box_questions == ()


def test_mixed_record_encoding(tok, records):
    enc = encode_glint_record(tok, records["mixed"])
    assert isinstance(enc, GlintEncodedRecord)
    assert [q.question_id for q in enc.questions] == ["screen", "has_error", "btn_visible"]
    assert [q.question_id for q in enc.box_questions] == ["btn", "badge"]
    s, e = enc.box_questions[0].question_span
    assert enc.input_ids[s:e] == tuple(tok("the primary button in `screen`").input_ids)
    # choice fields keep the exact same token text as Clef would produce for them
    choice_only = {**records["mixed"], "questions": {k: v for k, v in records["mixed"]["questions"].items() if v["type"] != "box"}}
    ref = encode_record(tok, choice_only)
    for q, r in zip(enc.questions, ref.questions):
        assert [enc.input_ids[x:y] for x, y in q.option_spans] == [ref.input_ids[x:y] for x, y in r.option_spans]
    ids = list(enc.input_ids)
    for needle in ("TYPE: box", BOX_OUTPUT_TEXT):
        n = tok(needle).input_ids
        assert any(ids[i:i + len(n)] == n for i in range(len(ids))), needle


def test_image_token_coords():
    pad = 99
    ids = [1, 2] + [pad] * 6 + [3] + [pad] * 2 + [4]
    grid = [[1, 4, 6], [1, 2, 4]]                                 # merge 2 -> 2x3 and 1x2 tokens
    pos, yx = image_token_coords(ids, pad, grid, 2)
    assert pos == (2, 3, 4, 5, 6, 7, 9, 10)
    assert yx[0] == (0.25, 1 / 6) and yx[5] == (0.75, 5 / 6) and yx[6] == (0.5, 0.25)
    assert image_token_coords(ids[:-2], pad, grid, 2) is None    # count mismatch -> no position encoding


def _model(tok, box=True, hidden=32):
    from tests.conftest import FakeBackbone
    torch.manual_seed(0)
    head = JointSchemaHead(hidden_size=hidden, width=16, routing_layers=1, layers=1, heads=2, feedforward=32)
    bh = BoxHead(hidden_size=hidden, width=16, layers=1, heads=2, feedforward=32, bins=16) if box else None
    return GlintModel(FakeBackbone(hidden=hidden), head, bh)


def test_glint_model_forward_matches_clef_without_boxes(tok, records):
    m = _model(tok).eval()
    enc = encode_glint_record(tok, records["choice"])
    batch = collate_records([enc], tok.pad_token_id, torch.device("cpu"))
    ref = ClefModel(m.language_model, m.head).eval()(batch)
    out = m(batch)
    assert all(torch.equal(a, b) for a, b in zip(ref[0], out[0]))
    logits, boxes = m.decide(batch)
    assert boxes == [None]


def test_glint_model_mixed_batch_and_train_step(tok, records):
    m = _model(tok).train()
    encs = [encode_glint_record(tok, records[k]) for k in ("mixed", "box_only", "choice")]
    batch = collate_records(encs, tok.pad_token_id, torch.device("cpu"))
    logits, boxes = m.decide(batch)
    assert [len(l) for l in logits] == [3, 0, 2]
    assert boxes[0].logits.shape == (2, 4, 16) and boxes[1].logits.shape == (1, 4, 16) and boxes[2] is None
    ex = {"id": "m", "questions": records["mixed"]["questions"],
          "gold": {"screen": "error", "btn": {"top": 0.6, "bottom": 0.4, "left": 0.1, "right": 0.5}, "badge": {"present": False}},
          "teacher": {"screen": {"home": 0.2, "error": 0.8}, "btn": {"type": "box", "box": {"top": 0.1}}}}
    gold = derive_presence_gold(ex)
    assert gold["btn_visible"] == "true" and "badge_visible" not in gold
    assert normalize_box_gold(ex) == {"btn": [0.4, 0.6, 0.1, 0.5], "badge": None}
    assert "btn" not in normalize_teacher(ex)                     # teacher entries for box questions are ignored
    dl, _ = decision_loss(logits[0], encs[0], normalize_teacher(ex), gold)
    bl, parts, n = record_box_loss(boxes[0], normalize_box_gold(ex))
    assert n == 1 and set(parts) >= {"dfl", "l1", "giou", "conf"}
    (dl + bl).backward()
    # everything but the image-position projection (no images in this record) receives gradient
    assert all(p.grad is not None for n, p in m.box_head.named_parameters() if not n.startswith("position_projection"))
    assert record_box_loss(boxes[2], {}) == (None, {}, 0)


def test_box_without_box_head_raises(tok, records):
    m = _model(tok, box=False)
    batch = collate_records([encode_glint_record(tok, records["box_only"])], 0, torch.device("cpu"))
    try:
        m.decide(batch)
    except ValueError as e:
        assert "box head" in str(e)
    else:
        raise AssertionError


def test_presence_gold_does_not_override_explicit():
    ex = {"questions": {"b": {"type": "box"}, "b_visible": {"type": "noul"}}, "gold": {"b": {"present": False}, "b_visible": "true"}}
    assert derive_presence_gold(ex)["b_visible"] == "true"
    ex = {"questions": {"b": {"type": "box"}, "b_visible": {"type": "noul"}}, "gold": {"b": {"present": False}}}
    assert derive_presence_gold(ex)["b_visible"] == "false"
    ex = {"questions": {"s": {"type": "noul"}}, "gold": {"s": "true"}}
    assert derive_presence_gold(ex) is ex["gold"]


def test_teacher_questions_drops_box(records):
    assert set(teacher_questions(records["mixed"]["questions"])) == {"screen", "has_error", "btn_visible"}


def test_make_synthetic_box(tmp_path, tok):
    out = tmp_path / "syn"
    subprocess.run([sys.executable, str(ROOT / "make_synthetic.py"), "--out", str(out), "--n-train", "5", "--n-eval", "0", "--box"], check=True, capture_output=True)
    rows = load_jsonl(out / "train.jsonl")
    seen_present = set()
    for ex in rows:
        g = ex["gold"]["primary_button_box"]
        assert ex["questions"]["primary_button_box"]["type"] == "box"
        assert "primary_button_box" not in ex["teacher"]
        assert ex["gold"]["primary_button_box_visible"] == ("true" if g["present"] else "false")
        b = normalize_box_gold(ex)["primary_button_box"]
        seen_present.add(b is not None)
        if b is not None:
            assert 0 <= b[0] < b[1] <= 1 and 0 <= b[2] < b[3] <= 1
        enc = encode_glint_record(tok, {**ex, "images": []})
        assert [q.question_id for q in enc.box_questions] == ["primary_button_box"]
    assert seen_present == {True, False}