#!/usr/bin/env python3
"""Beginner lab: trace an AI request across service boundaries."""
from __future__ import annotations
import json
import sys

REQUEST_ID = "REQ-18427"

SCENARIOS = {
    "healthy": [
        ("client", 18, "ok"),
        ("application", 12, "ok"),
        ("gateway", 20, "ok"),
        ("queue", 120, "ok"),
        ("inference", 900, "ok"),
        ("stream", 420, "ok"),
    ],
    "queue-spike": [
        ("client", 18, "ok"),
        ("application", 12, "ok"),
        ("gateway", 20, "ok"),
        ("queue", 4200, "slow"),
        ("inference", 900, "ok"),
        ("stream", 420, "ok"),
    ],
    "auth-denied": [
        ("client", 18, "ok"),
        ("application", 12, "ok"),
        ("gateway", 11, "auth_denied"),
    ],
}


def run(name: str) -> int:
    if name not in SCENARIOS:
        print("usage: 02_trace_ai_request.py healthy|queue-spike|auth-denied", file=sys.stderr)
        return 2
    stages = SCENARIOS[name]
    print(f"scenario={name} request_id={REQUEST_ID}")
    total = 0
    for stage, duration_ms, status in stages:
        total += duration_ms
        print(json.dumps({"request_id": REQUEST_ID, "stage": stage, "duration_ms": duration_ms, "status": status}, separators=(",", ":")))
    print(f"total_ms={total}")
    if name == "queue-spike":
        print("diagnosis=capacity_or_queue_pressure")
        print("reason=queue_delay_dominates_before_inference")
    elif name == "auth-denied":
        print("diagnosis=policy_failure")
        print("reason=request_stopped_before_inference;do_not_retry_model")
    else:
        print("diagnosis=healthy_trace")
        print("reason=all_stages_within_teaching_threshold")
    return 0


if __name__ == "__main__":
    scenario = sys.argv[1] if len(sys.argv) > 1 else "healthy"
    raise SystemExit(run(scenario))