"""Per-field results -> response assembly and validation (SPEC.md §6)."""

from __future__ import annotations

import pytest
from deg_cases import ARRAYS, INVOICE, PRODUCT_LISTING, SUPPORT_TICKET

from deg.response import (
    MODEL_ID,
    FieldResult,
    FillError,
    assemble,
    concrete_pointer,
    confidence,
    decide_choice,
    decide_noul,
    decide_score,
    extract,
    generate,
)
from deg.schema import compile_schema

TICKET_CONTEXT = (
    "From: dana@example.com\nOrder AB-1234 was charged twice. I need a refund today!"
)


def span_of(text, needle):
    start = text.index(needle)
    return (start, start + len(needle))


def ticket_results(plan):
    f = plan.field
    return {
        "/department": decide_choice(f("/department"), [0.8, 0.15, 0.05]),
        "/urgent": decide_noul(f("/urgent"), 0.91),
        "/anger": decide_score(f("/anger"), [0.0, 0.1, 0.6, 0.3, 0.0]),
        "/customer_email": extract(f("/customer_email"), TICKET_CONTEXT,
                                   span_of(TICKET_CONTEXT, "dana@example.com"), 0.97),
        "/order_id": extract(f("/order_id"), TICKET_CONTEXT,
                             span_of(TICKET_CONTEXT, "AB-1234"), 0.88),
        "/reply": generate(f("/reply"), "Sorry about the double charge; refund issued.", 11),
    }


def test_support_ticket_response():
    plan = compile_schema(SUPPORT_TICKET)
    body = assemble(plan, ticket_results(plan), {}, prompt_tokens=42)
    assert body["model"] == MODEL_ID
    assert body["value"] == {
        "department": "billing",
        "urgent": True,
        "anger": 2,
        "customer_email": "dana@example.com",
        "order_id": "AB-1234",
        "reply": "Sorry about the double charge; refund issued.",
    }
    assert list(body["value"]) == list(SUPPORT_TICKET["properties"])  # schema order
    assert body["usage"] == {"prompt_tokens": 42, "generated_tokens": 11}
    assert body["plan"] == plan.fingerprint
    meta = body["meta"]
    assert meta["/department"]["mode"] == "decide.choice"
    assert meta["/department"]["probability"] == pytest.approx(0.8)
    assert meta["/department"]["confidence"] == pytest.approx(confidence([0.8, 0.15, 0.05]))
    assert meta["/urgent"] == {"mode": "decide.noul", "probability": 0.91}
    assert meta["/anger"]["expected"] == pytest.approx(2.2)
    assert meta["/customer_email"]["span"] == list(span_of(TICKET_CONTEXT, "dana@example.com"))
    assert meta["/reply"] == {"mode": "generate", "generated_tokens": 11}


def test_optional_extract_not_found_is_omitted_nullable_is_null():
    plan = compile_schema(SUPPORT_TICKET)
    results = ticket_results(plan)
    results["/order_id"] = extract(plan.field("/order_id"), TICKET_CONTEXT, None, 0.7)
    results["/customer_email"] = extract(plan.field("/customer_email"), TICKET_CONTEXT, None, 0.6)
    body = assemble(plan, results, {}, prompt_tokens=1)
    assert "order_id" not in body["value"]
    assert body["value"]["customer_email"] is None
    assert body["meta"]["/order_id"]["null_reason"] == "not_found"
    # A span that breaks the pattern is treated as not found, with the reason recorded.
    bad = extract(plan.field("/order_id"), TICKET_CONTEXT, span_of(TICKET_CONTEXT, "refund"), 0.5)
    assert bad.omit and bad.meta["null_reason"] == "pattern"


def test_required_extract_with_nothing_raises():
    plan = compile_schema(INVOICE)
    with pytest.raises(FillError, match="required non-nullable"):
        extract(plan.field("/invoice_number"), "no invoice here", None, 0.2)


def test_invoice_with_arrays_and_parsers():
    ctx = ("INVOICE #INV-0042  Issued: March 5, 2026  Due: 2026-04-04\n"
           "Acme Corp https://acme.example\n"
           "Widget x3  $1,200.00\nSupport plan  $300.50\nTOTAL: $1,500.50 USD")
    plan = compile_schema(INVOICE)
    f = plan.field
    results = {
        "/invoice_number": extract(f("/invoice_number"), ctx, span_of(ctx, "INV-0042"), 0.99),
        "/issued": extract(f("/issued"), ctx, span_of(ctx, "March 5, 2026"), 0.95),
        "/due": extract(f("/due"), ctx, span_of(ctx, "2026-04-04"), 0.95),
        "/vendor/name": extract(f("/vendor/name"), ctx, span_of(ctx, "Acme Corp"), 0.9),
        "/vendor/website": extract(f("/vendor/website"), ctx,
                                   span_of(ctx, "https://acme.example"), 0.9),
        "/currency": decide_choice(f("/currency"), [0.97, 0.02, 0.01]),
        "/line_items/0/description": extract(f("/line_items/*/description"), ctx,
                                             span_of(ctx, "Widget"), 0.9),
        "/line_items/0/quantity": extract(f("/line_items/*/quantity"), ctx, (ctx.index("x3") + 1,
                                          ctx.index("x3") + 2), 0.8),
        "/line_items/0/amount": extract(f("/line_items/*/amount"), ctx,
                                        span_of(ctx, "$1,200.00"), 0.9),
        "/line_items/1/description": extract(f("/line_items/*/description"), ctx,
                                             span_of(ctx, "Support plan"), 0.9),
        "/line_items/1/quantity": extract(f("/line_items/*/quantity"), ctx, None, 0.6),
        "/line_items/1/amount": extract(f("/line_items/*/amount"), ctx,
                                        span_of(ctx, "$300.50"), 0.9),
        "/total": extract(f("/total"), ctx, span_of(ctx, "$1,500.50"), 0.9),
    }
    body = assemble(plan, results, {"/line_items": 2}, prompt_tokens=80)
    assert body["value"] == {
        "invoice_number": "INV-0042",
        "issued": "2026-03-05",
        "due": "2026-04-04",
        "vendor": {"name": "Acme Corp", "website": "https://acme.example"},
        "currency": "USD",
        "line_items": [
            {"description": "Widget", "quantity": 3, "amount": 1200},
            {"description": "Support plan", "amount": 300.5},
        ],
        "total": 1500.5,
    }
    assert body["usage"]["generated_tokens"] == 0
    assert body["meta"]["/line_items"] == {"mode": "array", "items": 2}


def test_product_listing_generated_fields_and_token_totals():
    plan = compile_schema(PRODUCT_LISTING)
    ctx = "Used Lego castle set, complete, 4.5 stars, $45. Ships today."
    f = plan.field
    results = {
        "/title": extract(f("/title"), ctx, span_of(ctx, "Lego castle set"), 0.9),
        "/category": decide_choice(f("/category"), [0.0, 0.05, 0.9, 0.05]),
        "/condition": decide_choice(f("/condition"), [0.05, 0.9, 0.05]),
        "/price": extract(f("/price"), ctx, span_of(ctx, "$45"), 0.95),
        "/rating": extract(f("/rating"), ctx, span_of(ctx, "4.5"), 0.9),
        "/in_stock": decide_noul(f("/in_stock"), 0.8),
        "/description": generate(f("/description"),
                                 "A complete Lego castle. Fun for builders of all ages.", 14),
        "/tags/0": generate(f("/tags/*"), "lego", 2),
        "/tags/1": generate(f("/tags/*"), "castle", 2),
    }
    body = assemble(plan, results, {"/tags": 2}, prompt_tokens=30)
    assert body["value"]["rating"] == 4.5
    assert body["value"]["price"] == 45
    assert body["value"]["tags"] == ["lego", "castle"]
    assert body["usage"] == {"prompt_tokens": 30, "generated_tokens": 18}


def test_score_null_outcome_and_expected_level():
    plan = compile_schema({"type": "object", "properties": {
        "n": {"type": ["integer", "null"], "minimum": 1, "maximum": 3},
    }})
    n = decide_score(plan.field("/n"), [0.1, 0.1, 0.1, 0.7])
    assert n.value is None and n.meta["expected"] == pytest.approx(2.0)
    n2 = decide_score(plan.field("/n"), [0.1, 0.5, 0.1, 0.3])
    assert n2.value == 2


def test_dev_fit_decode_rules():
    plan = compile_schema({"type": "object", "properties": {
        "b": {"type": "boolean"},
        "s": {"type": "integer", "minimum": 1, "maximum": 5},
        "n": {"type": ["integer", "null"], "minimum": 1, "maximum": 3},
    }})
    assert decide_noul(plan.field("/b"), 0.4).value is False
    assert decide_noul(plan.field("/b"), 0.4, threshold=0.36).value is True
    probs = [0.0, 0.45, 0.1, 0.0, 0.45]  # argmax level 2, expected 3.35
    assert decide_score(plan.field("/s"), probs).value == 2
    assert decide_score(plan.field("/s"), probs, expected_bias=0.2).value == 4
    assert decide_score(plan.field("/s"), probs, expected_bias=0.0).value == 3
    # a winning null outcome stays null under the expected-level rule
    assert decide_score(plan.field("/n"), [0.1, 0.1, 0.1, 0.7], expected_bias=0.2).value is None


def test_bounded_number_is_extracted_and_bounds_validate():
    plan = compile_schema({"type": "object", "properties": {
        "q": {"type": "number", "minimum": 0, "maximum": 1, "multipleOf": 0.25},
    }, "required": ["q"]})
    fld = plan.field("/q")
    assert fld.mode == "extract.number" and fld.levels is None
    ctx = "ratio 0.75, later 0.8"
    assert extract(fld, ctx, span_of(ctx, "0.75"), 0.9).value == 0.75
    with pytest.raises(FillError):  # 0.8 is not a multiple of 0.25; required non-nullable
        extract(fld, ctx, span_of(ctx, "0.8"), 0.9)


def test_tie_breaks_and_noul_threshold():
    plan = compile_schema({"type": "object", "properties": {
        "c": {"enum": ["a", "b"]}, "b": {"type": "boolean"}}})
    assert decide_choice(plan.field("/c"), [0.5, 0.5]).value == "a"
    assert decide_noul(plan.field("/b"), 0.5).value is True
    assert decide_noul(plan.field("/b"), 0.4999).value is False


def test_arrays_nested_assembly():
    plan = compile_schema(ARRAYS)
    ctx = "Bob (aka Rob, Bobby) and Ann. Grid: 1 2 / 3"
    f = plan.field
    results = {
        "/labels/0": decide_choice(f("/labels/*"), [0.9, 0.05, 0.05]),
        "/people/0/name": extract(f("/people/*/name"), ctx, span_of(ctx, "Bob"), 0.9),
        "/people/0/aliases/0": extract(f("/people/*/aliases/*"), ctx, span_of(ctx, "Rob"), 0.8),
        "/people/0/aliases/1": extract(f("/people/*/aliases/*"), ctx, span_of(ctx, "Bobby"), 0.8),
        "/people/1/name": extract(f("/people/*/name"), ctx, span_of(ctx, "Ann"), 0.9),
        "/matrix/0/0": extract(f("/matrix/*/*"), ctx, span_of(ctx, "1"), 0.9),
        "/matrix/0/1": extract(f("/matrix/*/*"), ctx, span_of(ctx, "2"), 0.9),
        "/matrix/1/0": extract(f("/matrix/*/*"), ctx, span_of(ctx, "3"), 0.9),
    }
    lengths = {"/labels": 1, "/people": 2, "/people/0/aliases": 2, "/people/1/aliases": 0,
               "/matrix": 2, "/matrix/0": 2, "/matrix/1": 1}
    body = assemble(plan, results, lengths, prompt_tokens=5)
    assert body["value"] == {
        "labels": ["bug"],
        "people": [{"name": "Bob", "aliases": ["Rob", "Bobby"]},
                   {"name": "Ann", "aliases": []}],
        "matrix": [[1, 2], [3]],
    }


def test_assembly_rejects_missing_extra_and_invalid():
    plan = compile_schema(SUPPORT_TICKET)
    results = ticket_results(plan)
    missing = dict(results)
    del missing["/urgent"]
    with pytest.raises(FillError, match="missing field result"):
        assemble(plan, missing, {}, prompt_tokens=0)
    with pytest.raises(FillError, match="does not contain"):
        assemble(plan, {**results, "/ghost": FieldResult(1, {"mode": "extract.span"})}, {},
                 prompt_tokens=0)
    with pytest.raises(FillError, match="meta.mode"):
        assemble(plan, {**results, "/urgent": FieldResult(True, {"mode": "generate"})}, {},
                 prompt_tokens=0)
    with pytest.raises(FillError, match="is not one of"):  # schema validation catches it
        assemble(plan, {**results, "/department": FieldResult("hr", {"mode": "decide.choice"})},
                 {}, prompt_tokens=0)


def test_assembly_enforces_array_bounds():
    plan = compile_schema(INVOICE)
    ctx = "INV-1 2026-01-01 Acme USD"
    f = plan.field
    results = {
        "/invoice_number": extract(f("/invoice_number"), ctx, (0, 5), 1.0),
        "/issued": extract(f("/issued"), ctx, (6, 16), 1.0),
        "/due": extract(f("/due"), ctx, None, 0.5),
        "/vendor/name": extract(f("/vendor/name"), ctx, (17, 21), 1.0),
        "/vendor/website": extract(f("/vendor/website"), ctx, None, 0.5),
        "/currency": decide_choice(f("/currency"), [1, 0, 0]),
        "/total": FieldResult(0, {"mode": "extract.number"}),
    }
    with pytest.raises(FillError, match="should be non-empty|too short"):
        assemble(plan, results, {"/line_items": 0}, prompt_tokens=0)  # minItems: 1
    with pytest.raises(FillError, match="array length"):
        assemble(plan, results, {}, prompt_tokens=0)


def test_generate_must_respect_bounds():
    plan = compile_schema(PRODUCT_LISTING)
    with pytest.raises(FillError, match="minLength"):
        generate(plan.field("/description"), "too short", 3)
    with pytest.raises(FillError, match="maxLength"):
        generate(plan.field("/tags/*"), "x" * 25, 3)


def test_bad_inputs_rejected():
    plan = compile_schema(SUPPORT_TICKET)
    with pytest.raises(FillError):
        decide_choice(plan.field("/department"), [0.5, 0.5])  # wrong arity
    with pytest.raises(FillError):
        decide_choice(plan.field("/department"), [0, 0, 0])
    with pytest.raises(FillError):
        decide_noul(plan.field("/urgent"), 1.5)
    with pytest.raises(FillError, match="out of range"):
        extract(plan.field("/order_id"), "abc", (2, 9), 0.5)
    with pytest.raises(FillError, match="mode"):
        decide_noul(plan.field("/department"), 0.5)


def test_date_time_assumed_offset_and_concrete_pointer():
    plan = compile_schema({"type": "object", "properties": {
        "at": {"type": "string", "format": "date-time"}}, "required": ["at"]})
    ctx = "Seen 2026-10-08 14:30 local"
    r = extract(plan.field("/at"), ctx, span_of(ctx, "2026-10-08 14:30"), 0.9,
                default_offset="-05:00")
    assert r.value == "2026-10-08T14:30:00-05:00" and r.meta["assumed_offset"] is True
    assert concrete_pointer(["a/b", 0, "c~"]) == "/a~1b/0/c~0"