from __future__ import annotations

import pytest
from deg_cases import INVOICE, PRODUCT_LISTING

from deg.datagen.domains import TEST_DOMAINS, TRAIN_DOMAINS, sample_spec
from deg.datagen.labels import outcomes, parse_top, question_messages
from deg.datagen.run import item_hint
from deg.datagen.validate import Reject, check_fill
from deg.schema import compile_schema

CTX = ("Used Lego castle set (complete). Rated 4.5 stars by 120 buyers. Price: $1,045.00, "
       "was $1,200. Ships today.")


def product(**over):
    value = {"title": "Lego castle set", "category": "toys", "condition": "used",
             "price": 1045, "rating": 4.5, "in_stock": True,
             "description": "A complete Lego castle. Fun for builders of all ages.",
             "tags": ["lego", "castle"]}
    value.update(over)
    return value


def test_valid_fill_recovers_spans_and_decides():
    plan = compile_schema(PRODUCT_LISTING)
    checked = check_fill(plan, CTX, product())
    assert CTX[slice(*checked.spans["/title"])] == "Lego castle set"
    assert CTX[slice(*checked.spans["/price"])] == "$1,045.00"
    assert CTX[slice(*checked.spans["/rating"])] == "4.5"
    assert checked.decide == {"/category": "toys", "/condition": "used", "/in_stock": True}
    assert checked.generate_fields == ["/description", "/tags/0", "/tags/1"]
    assert checked.value == product()


def test_case_repair_copies_context_text():
    plan = compile_schema(PRODUCT_LISTING)
    checked = check_fill(plan, CTX, product(title="lego Castle set"))
    assert checked.value["title"] == "Lego castle set" and checked.repaired == 1


@pytest.mark.parametrize(
    ("over", "reason"),
    [
        ({"title": "Duplo castle"}, "span_not_found"),
        ({"price": 999}, "number_no_span"),
        ({"category": "games"}, "decide_invalid"),
        ({"in_stock": "yes"}, "decide_invalid"),
        ({"description": "too short"}, "generate_bounds"),
        ({"extra": 1}, "undeclared_key"),
        ({"price": 1200, "rating": 120}, "extract_constraint"),  # rating > maximum 5
        ({"tags": ["a", "b", "c", "d", "e", "f"]}, "schema_invalid_value"),  # maxItems 5
    ],
)
def test_rejects(over, reason):
    plan = compile_schema(PRODUCT_LISTING)
    with pytest.raises(Reject) as info:
        check_fill(plan, CTX, product(**over))
    assert info.value.reason == reason


def test_missing_decide_and_required_extract():
    plan = compile_schema(PRODUCT_LISTING)
    value = product()
    del value["in_stock"]
    with pytest.raises(Reject, match="decide_missing"):
        check_fill(plan, CTX, value)
    value = product()
    del value["title"]
    with pytest.raises(Reject, match="required_missing"):
        check_fill(plan, CTX, value)


def test_dates_and_arrays_in_invoice():
    ctx = ("Invoice INV-0042 from Acme Corp (billing@acme.example, https://acme.example), "
           "issued March 3, 2026, due 2026-04-02. Currency: USD.\n"
           "Consulting 1,200.00\nSupport plan 300.5\nTotal 1,500.50")
    plan = compile_schema(INVOICE)
    value = {"invoice_number": "INV-0042", "issued": "2026-03-03", "due": "2026-04-02",
             "vendor": {"name": "Acme Corp", "website": "https://acme.example"},
             "currency": "USD",
             "line_items": [{"description": "Consulting", "amount": 1200},
                            {"description": "Support plan", "amount": 300.5}],
             "total": 1500.5}
    checked = check_fill(plan, ctx, value)
    assert ctx[slice(*checked.spans["/issued"])] == "March 3, 2026"
    assert ctx[slice(*checked.spans["/line_items/1/amount"])] == "300.5"
    assert item_hint(checked.value, "/line_items/1/amount").startswith("/line_items/1 = ")
    assert item_hint(checked.value, "/total") is None
    bad = {**value, "issued": "2026-03-04"}
    with pytest.raises(Reject, match="date_no_span"):
        check_fill(plan, ctx, bad)


def test_label_prompts_and_parse():
    plan = compile_schema({"type": "object", "properties": {
        "c": {"enum": ["a", "b", None]},
        "s": {"type": ["integer", "null"], "minimum": 1, "maximum": 3},
        "b": {"type": "boolean"}}})
    assert outcomes(plan.field("/c"))[-1].startswith("null")
    assert outcomes(plan.field("/s")) == ["1", "2", "3", "null (not stated / unknown)"]
    msg = question_messages(plan.field("/c"), "/c", "doc", None, [2, 1, 0])[1]["content"]
    assert "A. null" in msg and "C. a" in msg
    top = [{"token": "A", "logprob": -0.1}, {"token": " B", "logprob": -2.5},
           {"token": "x", "logprob": -3.0}]
    probs, cov = parse_top(top, [{"A"}, {"B"}], case_insensitive=False)
    assert sum(probs) == pytest.approx(1.0) and probs[0] > 0.9 and 0.9 < cov < 1.0
    yes = parse_top([{"token": "Yes", "logprob": -0.2}], [{"yes"}, {"no"}],
                    case_insensitive=True)
    assert yes[0] == [1.0, 0.0]


def test_spec_sampling_is_deterministic_and_covers_test_domains():
    assert sample_spec(0, 7) == sample_spec(0, 7)
    specs = [sample_spec(0, i) for i in range(20000)]
    assert all(50 <= s.words <= 1500 for s in specs)
    test = [s for s in specs if s.split_domain == "test"]
    assert {s.domain for s in test} == set(TEST_DOMAINS)
    assert 0.05 <= len(test) / len(specs) <= 0.08
    assert len({s.domain for s in specs}) > 250


def test_domains_held_out_and_brainstorm_filters():
    from deg.datagen.brainstorm import dedupe, leaks
    from deg.datagen.domains import BASE_TEST_DOMAINS

    assert set(BASE_TEST_DOMAINS) <= set(TEST_DOMAINS)
    assert len(TEST_DOMAINS) >= 13 and not set(TEST_DOMAINS) & set(TRAIN_DOMAINS)
    for train in TRAIN_DOMAINS[21:]:  # brainstormed train domains never touch a held-out topic
        assert not any(leaks(train, t) for t in TEST_DOMAINS), train
    assert leaks("warranty claim forms", "insurance claims")
    assert leaks("auto repair invoices", "vehicle service and repair records")
    assert not leaks("sales invoices", "insurance claims")
    assert dedupe(["Pharmacy Refill Requests", "pharmacy refill request", "x"],
                  ["invoices and receipts"]) == ["pharmacy refill requests"]


def test_v2_knobs_sample_modes_without_changing_v1_draws():
    base = [sample_spec(1, i) for i in range(3000)]
    mixed = [sample_spec(1, i, real_schema_frac=0.3, hard_frac=0.3) for i in range(3000)]
    for b, m in zip(base, mixed):
        assert (b.domain, b.archetype, b.words, b.style) == (m.domain, m.archetype, m.words,
                                                             m.style)
    assert not any(b.real_schema or b.hard for b in base)
    real = sum(m.real_schema for m in mixed) / len(mixed)
    hard = [m.hard for m in mixed if m.hard]
    assert 0.25 < real < 0.35 and 0.25 < len(hard) / len(mixed) < 0.35
    assert set(hard) == {"absent", "distractor", "correction"}


def test_clef_questions_and_parse():
    from deg.datagen.clef import agree, build_request, parse_response

    plan = compile_schema({"type": "object", "properties": {
        "c": {"enum": ["a", 1, None]},
        "b": {"type": "boolean"},
        "s": {"type": "integer", "minimum": 1, "maximum": 5},
        "w": {"type": "integer", "minimum": 0, "maximum": 10},
        "n": {"type": ["integer", "null"], "minimum": 1, "maximum": 3}}})
    items = [(f.pointer, f, None) for f in plan.fields]
    req, keys = build_request("ctx", items)
    q = req["questions"]
    assert req["state"] == "ctx"
    assert q["/c"]["type"] == "choice" and list(q["/c"]["criteria"]) == ["a", "1", "null"]
    assert q["/b"] == {"type": "noul", "instructions": "JSON field /b"}
    assert q["/s"]["type"] == "score" and q["/s"]["criteria"] == ["1", "2", "3", "4", "5"]
    assert q["/w"]["type"] == "choice" and len(q["/w"]["criteria"]) == 11  # > 10 levels
    assert list(q["/n"]["criteria"]) == ["1", "2", "3", "null"]
    resp = {"answers": {
        "/c": {"type": "choice", "probabilities": {"a": 0.2, "1": 0.7, "null": 0.1}},
        "/b": {"type": "noul", "noul": 0.9},
        "/s": {"type": "score", "probabilities": {str(i): 0.2 for i in range(5)}},
        "/w": {"type": "choice", "probabilities": {str(i): 1 / 11 for i in range(11)}},
        "/n": {"type": "choice", "probabilities": {"1": 0.1, "2": 0.1, "3": 0.1, "null": 0.7}},
    }}
    got = parse_response(resp, keys)
    assert got["/c"] == pytest.approx([0.2, 0.7, 0.1]) and got["/b"] == pytest.approx([0.9, 0.1])
    assert got["/n"][-1] == pytest.approx(0.7)
    assert agree([0.6, 0.4], [0.8, 0.2]) is True
    assert agree([0.5, 0.5], [0.9, 0.1]) is False  # |dp| = 0.4
    assert agree([0.4, 0.6], [0.6, 0.4]) is False and agree(None, [1.0]) is None


def test_ambiguous_spans_listed_and_marked():
    from deg.datagen.labels import mark_spans

    ctx = "1. Widget A, Qty 10, $5.00 ea\n2. Widget B, Qty 5, $2.00 ea"
    plan = compile_schema({"type": "object", "properties": {"items": {
        "type": "array", "items": {"type": "object", "properties": {
            "name": {"type": "string"}, "qty": {"type": "integer"}}}}}})
    checked = check_fill(plan, ctx, {"items": [{"name": "Widget A", "qty": 10},
                                               {"name": "Widget B", "qty": 5}]})
    alts = checked.alternatives["/items/1/qty"]
    assert [ctx[a:b] for a, b in alts] == ["$5.00", "5"]
    assert checked.alternatives["/items/0/qty"] == [checked.spans["/items/0/qty"]]
    marked = mark_spans(ctx, alts)
    assert "[[A]]$5.00" in marked and "Qty [[B]]5," in marked


def test_teacher_spec_parse(monkeypatch):
    from deg.datagen.client import parse_teacher_spec

    got = parse_teacher_spec("url=127.0.0.1:8780,model=Q27,concurrency=24")
    assert got == {"base_url": "http://127.0.0.1:8780/v1", "api_key": "EMPTY",
                   "model": "Q27", "concurrency": 24}
    monkeypatch.setenv("K", "secret")
    assert parse_teacher_spec("url=http://h/v1,model=m,key_env=K")["api_key"] == "secret"
    with pytest.raises(SystemExit):
        parse_teacher_spec("url=http://h,concurrency=3")
    with pytest.raises(SystemExit):
        parse_teacher_spec("url=http://h,model=m,bogus=1")


def test_find_string_folds_typographic_quotes_and_dashes():
    from deg.datagen.validate import _find_string

    spans, text = _find_string("Contact: Kevin O’Connor – Sales", "kevin o'connor - sales")
    assert text == "Kevin O’Connor – Sales" and len(spans) == 1