"""
Practice: a tracer from scratch, and the two checks that find missing spans. Lab 14's
script 01 (Exercise 2): spans that nest through a context variable, a worker that does or
doesn't get the ticket, and the checks against the bill.

Fill in the TODOs, then run the tests:
    pytest test_trace_practice.py        (locally)
or press "Run tests" on the chapter page.
"""
import contextvars
import itertools
from contextlib import contextmanager
from dataclasses import dataclass, field

# The "ticket": the span we're currently inside. Each new span reads it to find its parent.
_current = contextvars.ContextVar("current_span", default=None)
_ids = itertools.count(1)


@dataclass
class Span:
    trace_id: str
    span_id: str
    parent_id: str | None
    name: str
    attributes: dict = field(default_factory=dict)


class Tracer:
    def __init__(self):
        self.spans = []

    @contextmanager
    def span(self, name, attributes=None):
        """Open a span named `name` and yield it.
        - The parent is the current span (_current.get()). With no parent, start a new trace:
          trace_id f"t{n}", parent_id None. Otherwise use the parent's trace_id, and its span_id
          as parent_id. Every span gets span_id f"s{n}" (n from next(_ids)).
        - While inside, this span is the current one (set _current, and reset it afterwards,
          even if the body raises).
        - When it ends, append it to self.spans."""
        parent = _current.get()
        if parent is None:
            s = Span(f"t{next(_ids)}", f"s{next(_ids)}", None, name, dict(attributes or {}))
        else:
            s = Span(parent.trace_id, f"s{next(_ids)}", parent.span_id, name, dict(attributes or {}))
        token = _current.set(s)
        try:
            yield s
        finally:
            _current.reset(token)
            self.spans.append(s)


def run_on_worker(fn, pass_ticket):
    """Run fn() the way a worker thread would, and return its result.
    A new thread starts with an EMPTY context: run fn inside contextvars.Context().
    Passing the ticket means running it in a copy of the caller's context:
    contextvars.copy_context().run(fn)."""
    if pass_ticket:
        return contextvars.copy_context().run(fn)
    return contextvars.Context().run(fn)


def is_model_call(span):
    """A model call's span has a "model" attribute."""
    return "model" in span.attributes


def rollup(spans, trace_id):
    """One trace's totals: {"cost": sum of the "cost" attributes, "model_calls": how many
    model-call spans}."""
    mine = [s for s in spans if s.trace_id == trace_id]
    return {"cost": sum(s.attributes.get("cost", 0) for s in mine),
            "model_calls": sum(1 for s in mine if is_model_call(s))}


def request_traces(spans):
    """The trace ids whose root span (parent_id None) is named "request"."""
    return {s.trace_id for s in spans if s.parent_id is None and s.name == "request"}


def share_of_bill(spans, bill, only_requests=False):
    """Traced model-call cost ÷ billed cost. `bill` is a list of {"model", "cost"}.
    With only_requests=True, count only spans inside request traces."""
    keep = request_traces(spans) if only_requests else None
    traced = sum(s.attributes["cost"] for s in spans
                 if is_model_call(s) and (keep is None or s.trace_id in keep))
    return traced / sum(b["cost"] for b in bill)


def check_calls_vs_bill(spans, bill):
    """Check 1: calls traced vs. calls billed, model by model. Return {model: (billed, traced)}
    for every model where the two counts differ ({} means the check passes)."""
    billed, traced = {}, {}
    for b in bill:
        billed[b["model"]] = billed.get(b["model"], 0) + 1
    for s in spans:
        if is_model_call(s):
            m = s.attributes["model"]
            traced[m] = traced.get(m, 0) + 1
    return {m: (billed.get(m, 0), traced.get(m, 0))
            for m in sorted(set(billed) | set(traced)) if billed.get(m, 0) != traced.get(m, 0)}


def complete_requests(spans, calls_per_request=3):
    """Check 2: how many request traces hold exactly `calls_per_request` model calls."""
    return sum(1 for t in request_traces(spans) if rollup(spans, t)["model_calls"] == calls_per_request)