"""
Docker example: run multiple models in isolated containers with budget enforcement.

Demonstrates:
- run_multi (parallel Docker runner)
- Custom eject policy
- Per-model progress tracking via on_event

Assumes you have:
1. A Docker image named MY_IMAGE that:
   - Reads its task from /workspace/TASK.md
   - Uses OPENROUTER_API_KEY env var
   - Optionally reads /workspace/output/budget.txt for cost awareness
2. A TASK.md file to mount

Usage:
    OR_ADMIN_KEY=sk-or-admin-... OPENROUTER_API_KEY=sk-or-... python examples/docker_run.py
"""

import os
from pathlib import Path

from agent_preflight import (
    BudgetFractionEject,
    CompositeEject,
    IdleTimeoutEject,
    run_multi,
)

ADMIN_KEY = os.environ["OR_ADMIN_KEY"]
OR_KEY = os.environ["OPENROUTER_API_KEY"]

MY_IMAGE = "my-agent:latest"
TASK_FILE = Path("TASK.md")
BUDGET = 0.25

MODELS = [
    ("meta-llama/llama-3.3-70b-instruct", "llama-70b"),
    ("google/gemini-flash-1.5", "gemini-flash"),
    ("anthropic/claude-3-haiku", "haiku"),
]


def on_event(event: dict) -> None:
    """Called for every JSON event emitted by a container."""
    if event.get("type") == "tool_execution_end":
        tool = event.get("toolName", "?")
        result_preview = str(event.get("result", ""))[:80]
        print(f"    tool:{tool} → {result_preview}")


def main():
    if not TASK_FILE.exists():
        TASK_FILE.write_text("# Task\nExplore the filesystem and report what you find.\n")

    # Custom eject: kill if no tool calls after 25% budget spent
    eject = CompositeEject([
        BudgetFractionEject(
            check_after_pct=0.25,
            progress_fn=lambda ctx: ctx.get("tool_calls", 0) > 0,
        ),
        IdleTimeoutEject(timeout_s=180.0),
    ])

    results = run_multi(
        models=MODELS,
        image_name=MY_IMAGE,
        admin_key=ADMIN_KEY,
        or_key=OR_KEY,
        budget_usd=BUDGET,
        mounts=[
            (str(TASK_FILE.resolve()), "/workspace/TASK.md", "ro"),
        ],
        eject_policy=eject,
        on_event=on_event,
        budget_mode="detailed",
        parallel=True,
    )

    # Results
    print("\nResults:")
    for r in results:
        print(f"  {r.alias}: cost=${r.cost_usd or 0:.3f} tools={r.tool_calls} "
              f"error={r.error or 'none'}")


if __name__ == "__main__":
    main()
