"""
Basic example: track spend for a single model run.

Demonstrates:
- preflight validation
- sub-key provisioning
- BudgetTracker polling
- EventLogger
"""

import os
from pathlib import Path

from agent_preflight import (
    BudgetTracker,
    EventLogger,
    preflight,
    provision_key,
)

ADMIN_KEY = os.environ["OR_ADMIN_KEY"]
OR_KEY = os.environ["OPENROUTER_API_KEY"]
MODEL = "meta-llama/llama-3.3-70b-instruct"
BUDGET = 0.10


def main():
    # 1. Preflight
    preflight(
        admin_key=ADMIN_KEY,
        or_key=OR_KEY,
        models=[MODEL],
        budget_per_model=BUDGET,
        checks=["balance", "math", "models", "key_smoke"],
    )

    # 2. Provision a capped sub-key
    key = provision_key(ADMIN_KEY, label="basic-example", limit_usd=BUDGET)
    print(f"Provisioned key: {key[:16]}…")

    # 3. Set up budget tracker
    tracker = BudgetTracker(
        api_key=key,
        limit_usd=BUDGET,
        poll_interval_s=10.0,
        mode="detailed",
        on_exceeded=lambda: print("!! Budget exceeded — would kill agent here"),
    )

    # 4. Event logger
    run_dir = Path("runs/basic-example")
    run_dir.mkdir(parents=True, exist_ok=True)

    with EventLogger(run_dir / "events.jsonl") as ev:
        ev.log("run_start", model=MODEL, budget=BUDGET)

        # Simulate an agent loop — in real use you'd poll inside your agent's event loop
        status = tracker.poll(force=True)
        print(f"Initial budget: {status}")
        ev.log("budget_check", used=status.used, remaining=status.remaining)

        if status.exceeded:
            print("Budget already exceeded — aborting.")
            return

        ev.log("run_end", final_status=str(status))

    print(f"Done. Events saved to {run_dir / 'events.jsonl'}")


if __name__ == "__main__":
    main()
