"""Minimal end-to-end use of the agentblame API (no API keys, runs in < 1 s).

A 3-step pipeline: a mocked LLM picks a tool, the tool runs, a formatter
writes the answer. Each step has an alternative ("fix"); we ask which step
caused the failures in a small batch.

    python examples/quickstart.py
"""

from agentblame import Case, Pipeline, analyze_case, blame_batch, format_batch, format_case, step


@step("choose_tool")
def choose_tool(question, noise):
    # A sampled LLM decision: Gumbel-max over option logits. The noise is keyed by
    # (world, step, option), so the fixed variant below sees the *same* Gumbels.
    return noise.categorical({"calculator": 0.2, "web_search": 0.6, "none": 0.1})


@choose_tool.variant("better_prompt")
def choose_tool_fixed(question, noise):
    return noise.categorical({"calculator": 2.5, "web_search": 0.6, "none": 0.1})


@step("run_tool")
def run_tool(choose_tool, question):
    a, b = question
    return {"calculator": a * b, "web_search": a * b + 1, "none": None}[choose_tool]


@run_tool.variant("validated_tool")
def run_tool_fixed(choose_tool, question):
    a, b = question
    return a * b if choose_tool != "none" else None    # search results re-checked


@step("format")
def format_(run_tool):
    return "I don't know" if run_tool is None else f"The answer is {run_tool}"


@format_.variant("fallback_compute")
def format_fixed(run_tool):
    # Deliberately identical to as-run: a "fix" that changes nothing is a null player
    # and gets zero Shapley blame -- even though naive last-step-to-touch blames it.
    return "I don't know" if run_tool is None else f"The answer is {run_tool}"


def correct(out) -> float:
    a, b = out["question"]
    return float(out["format"] == f"The answer is {a * b}")


pipe = Pipeline([choose_tool, run_tool, format_], correct, name="quickstart")
cases = [Case(f"q{i}", {"question": (i, i + 3)}, seed=i) for i in range(40)]

if __name__ == "__main__":
    batch = blame_batch(pipe, cases)            # replays each failure under its own noise
    print(format_batch(batch), "\n")
    print(format_case(batch.cases[0]))
    # The same failing case with Monte Carlo permutation Shapley + CIs instead of exact:
    first = next(c for c in cases if c.id == batch.cases[0].case_id)
    cb = analyze_case(pipe, first, method="mc", mc={"max_permutations": 200})
    print("\nMC estimate for", cb.case_id,
          {k: (round(v, 3), tuple(round(x, 2) for x in cb.shapley.ci[k]))
           for k, v in cb.shapley.values.items()})