"""Tests on Lab 14's script 01, Exercise 2, as rerun for the video (October 2026): 30 test questions,
each billed for a search embedding, an answer from gpt-5.4-nano and a judge from gpt-5.4-mini.
Run: pytest test_trace_practice.py"""
from trace_practice import (Tracer, check_calls_vs_bill, complete_requests, request_traces, rollup,
                            run_on_worker, share_of_bill)

EMBED, NANO, MINI = "text-embedding-3-small", "gpt-5.4-nano", "gpt-5.4-mini"

# What each question's three calls cost in the recorded run (US dollars, list price, from the API's replies).
COSTS = [
    (2.8e-07, 9.515e-05, 0.0004335),
    (1.8e-07, 8.2e-05, 0.00043425),
    (2.6e-07, 8.4e-05, 0.00041475),
    (3.2e-07, 0.000161, 0.00042225),
    (3.2e-07, 0.000121, 0.00044775),
    (1.8e-07, 0.00010955, 0.00042675),
    (2.2e-07, 0.00012815, 0.0003975),
    (2e-07, 0.0001083, 0.0004215),
    (2.4e-07, 0.0001123, 0.00038025),
    (2e-07, 9.305e-05, 0.00036825),
    (2.4e-07, 7.94e-05, 0.000375),
    (3.2e-07, 0.00012975, 0.000417),
    (1.8e-07, 0.0001083, 0.000399),
    (1.8e-07, 0.00011745, 0.00036975),
    (2.2e-07, 0.00014855, 0.00040125),
    (2.8e-07, 0.00013015, 0.0004905),
    (2e-07, 0.00015305, 0.00040875),
    (2.6e-07, 8.4e-05, 0.00041925),
    (2e-07, 7.78e-05, 0.0004095),
    (2.2e-07, 0.0001488, 0.0004305),
    (3.8e-07, 0.0001945, 0.00053775),
    (2.8e-07, 8.08e-05, 0.00041175),
    (3e-07, 0.0001233, 0.00040125),
    (2.6e-07, 9.815e-05, 0.000348),
    (3.4e-07, 0.0001497, 0.0004635),
    (2.2e-07, 0.00014595, 0.00038325),
    (3.4e-07, 0.00013905, 0.0004605),
    (3.4e-07, 0.000142, 0.000459),
    (2.8e-07, 0.0001489, 0.00047925),
    (4.2e-07, 0.0001352, 0.00042225),
]


def run_setup(setup):
    """The assistant on all 30 questions. setup 1: only the answer traced; 2: every call traced, the judge on a
    worker without the ticket; 3: the ticket passed. Every call is billed whatever is traced."""
    tracer, bill = Tracer(), []

    def call(model, cost, traced):
        bill.append({"model": model, "cost": cost})
        if traced:
            with tracer.span(f"chat {model}", {"model": model, "cost": cost}):
                pass

    for search, answer, judge in COSTS:
        with tracer.span("request"):
            with tracer.span("invoke_workflow support_assistant"):
                with tracer.span("retrieval docs"):
                    call(EMBED, search, traced=setup != 1)
                call(NANO, answer, traced=True)
            if setup == 1:
                call(MINI, judge, traced=False)
            else:
                run_on_worker(lambda: call(MINI, judge, traced=True), pass_ticket=setup == 3)
    return tracer.spans, bill


def test_spans_nest_into_one_trace():
    t = Tracer()
    with t.span("request") as req:
        with t.span("retrieval docs") as ret:
            with t.span("embeddings", {"model": EMBED, "cost": 3.2e-07}) as emb:
                pass
        with t.span("chat", {"model": NANO, "cost": 0.00012475}) as ans:
            pass
    assert req.parent_id is None
    assert ret.parent_id == req.span_id and emb.parent_id == ret.span_id and ans.parent_id == req.span_id
    assert len({s.trace_id for s in (req, ret, emb, ans)}) == 1
    assert [s.name for s in t.spans] == ["embeddings", "retrieval docs", "chat", "request"]  # recorded as each ends
    with t.span("next request") as nxt:
        pass
    assert nxt.trace_id != req.trace_id and nxt.parent_id is None


def test_the_ticket_is_reset_even_after_an_error():
    t = Tracer()
    try:
        with t.span("request"):
            raise ValueError("model call failed")
    except ValueError:
        pass
    with t.span("after") as after:
        pass
    assert after.parent_id is None
    assert [s.name for s in t.spans] == ["request", "after"]


def test_rollup_of_request_r0003():
    t = Tracer()
    with t.span("request") as req:
        with t.span("invoke_workflow support_assistant"):
            with t.span("retrieval docs"):
                with t.span("embeddings text-embedding-3-small", {"model": EMBED, "cost": 3.2e-07}):
                    pass
            with t.span("chat gpt-5.4-nano", {"model": NANO, "cost": 0.00012475}):
                pass
        with t.span("chat gpt-5.4-mini", {"model": MINI, "cost": 0.000468}):
            pass
    r = rollup(t.spans, req.trace_id)
    assert r["model_calls"] == 3
    assert abs(r["cost"] - 0.00059307) < 1e-12
    assert round(0.000468 / r["cost"] * 100) == 79  # the judge, on this one question


def test_setup_1_only_the_main_call_traced():
    spans, bill = run_setup(1)
    assert len(bill) == 90 and sum(1 for s in spans if "model" in s.attributes) == 30
    assert round(share_of_bill(spans, bill) * 100) == 22
    assert round(sum(b["cost"] for b in bill if b["model"] == MINI) / sum(b["cost"] for b in bill) * 100) == 78
    assert check_calls_vs_bill(spans, bill) == {EMBED: (30, 0), MINI: (30, 0)}  # check 1 catches it
    assert complete_requests(spans) == 0                                       # so does check 2


def test_setup_2_judge_on_a_worker_without_the_ticket():
    spans, bill = run_setup(2)
    assert round(share_of_bill(spans, bill) * 100) == 100       # every call recorded...
    assert round(share_of_bill(spans, bill, only_requests=True) * 100) == 22   # ...but not in its question
    assert len({s.trace_id for s in spans}) - len(request_traces(spans)) == 30  # 30 judge traces of their own
    assert check_calls_vs_bill(spans, bill) == {}               # check 1 passes, and misses it
    assert complete_requests(spans) == 0                        # check 2 catches it


def test_setup_3_the_ticket_passed():
    spans, bill = run_setup(3)
    assert len(request_traces(spans)) == 30 and len({s.trace_id for s in spans}) == 30
    assert round(share_of_bill(spans, bill, only_requests=True) * 100) == 100
    assert check_calls_vs_bill(spans, bill) == {}
    assert complete_requests(spans) == 30