from __future__ import annotations

import asyncio
import json
import random

import pytest
from fake_teacher import FakeTeacher

from deg.datagen.hard import (
    SIDECAR,
    check_hard,
    leaked,
    make_nullable,
    pick_targets,
    split_planted,
)
from deg.datagen.run import Runner, build_parser
from deg.datagen.validate import Reject, check_fill
from deg.schema import compile_schema

ORDER = {"type": "object", "additionalProperties": False, "properties": {
    "order_id": {"type": "string", "description": "The order number"},
    "ship_date": {"type": "string", "format": "date", "description": "Date it shipped"},
    "total": {"type": "number", "description": "Order total"},
    "items": {"type": "array", "items": {"type": "string"}},
    "rush": {"type": "boolean"}}, "required": ["order_id", "rush"]}


def _checked(ctx, value):
    plan = compile_schema(ORDER)
    return plan, check_fill(plan, ctx, value)


def test_pick_targets_and_nullable():
    plan = compile_schema(ORDER)
    for kind in ("absent", "distractor", "correction"):
        got = pick_targets(plan, kind, random.Random(0))
        assert got and all("*" not in f.pointer and f.mode.startswith("extract.") for f in got)
    out = make_nullable(ORDER, ["/order_id", "/total"])
    assert out["properties"]["order_id"]["type"] == ["string", "null"]
    assert out["required"] == ["rush"] and ORDER["required"] == ["order_id", "rush"]
    assert ORDER["properties"]["order_id"]["type"] == "string"
    assert compile_schema(out).field("/order_id").nullable


def test_split_planted():
    doc, planted = split_planted(f"Hello there.\n{SIDECAR}\n{{\"/total\": \"$40.00\"}}")
    assert doc == "Hello there." and planted == {"/total": "$40.00"}
    assert split_planted("just a doc ") == ("just a doc", None)


def test_distractor_and_correction_checks():
    ctx = ("Order A-1001 shipped 2026-03-05. Total due $120.50. An earlier quote said $99.00. "
           "Correction: the order number is actually A-1002.")
    value = {"order_id": "A-1002", "ship_date": "2026-03-05", "total": 120.5, "rush": False}
    plan, checked = _checked(ctx, value)
    total, oid = plan.field("/total"), plan.field("/order_id")
    got = check_hard("distractor", [total], ctx, checked.value, {"/total": "$99.00"})
    assert got["planted"] == {"/total": "$99.00"}
    check_hard("correction", [oid], ctx, checked.value, {"/order_id": "A-1001"})
    with pytest.raises(Reject, match="distractor_chosen"):
        check_hard("distractor", [total], ctx, checked.value, {"/total": 120.5})
    with pytest.raises(Reject, match="planted_not_in_context"):
        check_hard("distractor", [total], ctx, checked.value, {"/total": "$77"})
    with pytest.raises(Reject, match="planted_missing"):
        check_hard("correction", [oid], ctx, checked.value, None)
    # corrected value stated BEFORE the superseded one -> wrong order
    ctx2 = "Order A-1002 placed. Earlier note: order A-1001. Shipped 2026-03-05, rush."
    _, checked2 = _checked(ctx2, {"order_id": "A-1002", "rush": True})
    with pytest.raises(Reject, match="correction_order"):
        check_hard("correction", [oid], ctx2, checked2.value, {"/order_id": "A-1001"})


def test_absent_checks():
    schema = make_nullable(ORDER, ["/total"])
    plan = compile_schema(schema)
    ctx = "Order A-1001 shipped 2026-03-05; no rush."
    checked = check_fill(plan, ctx, {"order_id": "A-1001", "total": None, "rush": False})
    check_hard("absent", [plan.field("/total")], ctx, checked.value, None)
    ctx3 = "Order A-1001 total 15. shipped 2026-03-05; no rush."
    filled = check_fill(plan, ctx3, {"order_id": "A-1001", "total": 15, "rush": False})
    with pytest.raises(Reject, match="absent_filled"):
        check_hard("absent", [plan.field("/total")], ctx3, filled.value, None)
    f = [plan.field("/total")]
    assert leaked(f, ctx, '{"/total": null}') is None
    assert leaked(f, ctx, '{"/total": "not stated"}') is None
    assert leaked(f, ctx3, '{"/total": "total 15"}') == "/total"
    with pytest.raises(Reject, match="absent_check_json"):
        leaked(f, ctx, "nope")


def _runner(tmp_path):
    args = build_parser().parse_args(["--out", str(tmp_path / "o"), "--no-clef",
                                      "--no-flash-labels", "--hard-frac", "1.0"])
    return Runner(args, [], None)


def _index_for(kind):
    from deg.datagen.domains import sample_spec
    return next(i for i in range(500) if sample_spec(0, i, 0.0, 1.0).hard == kind)


def test_runner_distractor_end_to_end(tmp_path):
    runner = _runner(tmp_path)
    state = {}
    doc = ("Hi team, order A-1002 is confirmed and shipped 2026-03-05. Total charged: $120.50 "
           "(the earlier quote of $99.00 expired). Not a rush order. Please file the "
           "paperwork with the usual vendor and let me know if anything looks off. Thanks!")

    def respond(prompt):
        if prompt.startswith("Domain:"):
            return json.dumps(ORDER)
        if prompt.startswith("Write one realistic document"):
            assert "HARD CASE (distractors)" in prompt
            targets = [p for p in ("/order_id", "/ship_date", "/total") if f"  * {p}:" in prompt]
            planted = {"/order_id": "A-1001", "/ship_date": "2026-02-01", "/total": "$99.00"}
            state["targets"] = {p: planted[p] for p in targets}
            return doc + " Old ref A-1001 from 2026-02-01."  # forgets the planted line
        if prompt.startswith("Now output ONLY a JSON object"):
            return json.dumps(state["targets"])
        if prompt.startswith("Fill the JSON Schema"):
            return json.dumps({"order_id": "A-1002", "ship_date": "2026-03-05",
                               "total": 120.5, "items": [], "rush": False})
        raise AssertionError(prompt[:60])

    row, kept = asyncio.run(runner.example(_index_for("distractor"), FakeTeacher(respond)))
    assert kept, row
    assert row["hard"] == "distractor" and row["hard_case"]["planted"]
    assert SIDECAR not in row["context"]


def test_runner_absent_rejects_leak(tmp_path):
    runner = _runner(tmp_path)

    def respond(prompt):
        if prompt.startswith("Domain:"):
            return json.dumps(ORDER)
        if prompt.startswith("Write one realistic document"):
            assert "HARD CASE (absent)" in prompt
            return ("Order A-1002 shipped 2026-03-05, total 120.50, not urgent at all. "
                    "Please confirm receipt with the warehouse and update the ledger today.")
        if prompt.startswith("Fill the JSON Schema"):
            return json.dumps({"order_id": None, "ship_date": None, "total": None,
                               "rush": False})
        if prompt.startswith("For each field below"):
            return json.dumps({"/order_id": "A-1002", "/ship_date": "2026-03-05",
                               "/total": "120.50"})
        raise AssertionError(prompt[:60])

    row, kept = asyncio.run(runner.example(_index_for("absent"), FakeTeacher(respond)))
    assert not kept and row["reason"] in ("absent_leaked", "required_missing")
    assert row["hard"] == "absent"


def test_planted_key_variants():
    from deg.datagen.hard import planted_for

    assert planted_for({"dimensions.length": 4}, "/dimensions/length", 2) == 4
    assert planted_for({"dimensions/length": 4}, "/dimensions/length", 2) == 4
    assert planted_for({"Length": 4}, "/dimensions/length", 2) == 4
    assert planted_for({"whatever": 4}, "/dimensions/length", 1) == 4
    assert planted_for({"whatever": 4}, "/dimensions/length", 2) is None