"""
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."""
        # TODO: find the parent, nest under it, make this span current while inside.
        s = Span(f"t{next(_ids)}", f"s{next(_ids)}", None, name, dict(attributes or {}))
        yield s
        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)."""
    # TODO
    return 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}."""
    # TODO
    return {"cost": 0.0, "model_calls": 0}


def request_traces(spans):
    """The trace ids whose root span (parent_id None) is named "request"."""
    # TODO
    return set()


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."""
    # TODO
    return 0.0


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)."""
    # TODO
    return {}


def complete_requests(spans, calls_per_request=3):
    """Check 2: how many request traces hold exactly `calls_per_request` model calls."""
    # TODO
    return 0