#!/usr/bin/env python3
"""Multi-agent orchestration patterns: demonstrating supervisor-worker.

This is a SKELETON that shows the structure. Plug in real agent/model calls
for production use. The value is the explicit state machine and guardrails.
"""

from __future__ import annotations

import time
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable


# ── Agent primitives ──────────────────────────────────────────────────────


@dataclass
class Agent:
    name: str
    role: str  # "supervisor" | "worker" | "reviewer"
    tools: list[str] = field(default_factory=list)
    handler: Callable[[str], str] = lambda self, x: f"[{self.name}] processed {x!r}"


@dataclass
class Task:
    id: str
    description: str
    assigned_to: str | None = None
    result: str | None = None


# ── Patterns ──────────────────────────────────────────────────────────────


class SupervisorWorker:
    """Supervisor decomposes task, assigns to workers, collects results."""

    def __init__(self, supervisor: Agent, workers: list[Agent], max_steps: int = 10):
        self.supervisor = supervisor
        self.workers = {w.name: w for w in workers}
        self.max_steps = max_steps
        self.trace: list[str] = []

    def run(self, task_text: str) -> dict[str, str]:
        """Execute a task through the supervisor-worker pattern."""
        run_id = uuid.uuid4().hex[:8]
        self.trace.append(f"[{run_id}] START: {task_text[:60]}")

        step = 0
        subtasks = [
            Task(f"{run_id}-1", "Analyze input and identify sub-problems"),
            Task(f"{run_id}-2", "Solve each sub-problem"),
            Task(f"{run_id}-3", "Merge results and verify"),
        ]

        results: dict[str, str] = {}

        for task in subtasks:
            if step >= self.max_steps:
                self.trace.append(f"[{run_id}] HALT: step limit reached")
                break

            # Assign round-robin to workers
            worker = list(self.workers.values())[step % len(self.workers)]
            task.assigned_to = worker.name
            task.result = worker.handler(task.description)
            results[task.id] = task.result
            self.trace.append(f"[{run_id}] {worker.name}: {task.result[:60]}")
            step += 1

        self.trace.append(f"[{run_id}] END: {len(results)}/{len(subtasks)} done")
        return results


class FanOut:
    """Parallel fan-out: one task to all workers simultaneously."""

    def __init__(self, workers: list[Agent]):
        self.workers = workers

    def run(self, task_text: str) -> dict[str, str]:
        """In production: asyncio.gather over worker LLM calls."""
        results = {}
        for w in self.workers:
            results[w.name] = w.handler(task_text)
        return results


# ── Guardrails ────────────────────────────────────────────────────────────


class Guardrails:
    """Hard caps and validation that stop runaway agent loops."""

    def __init__(
        self,
        max_steps: int = 10,
        max_tokens: int = 100_000,
        max_spend: float = 5.0,
    ):
        self.max_steps = max_steps
        self.max_tokens = max_tokens
        self.max_spend = max_spend
        self.spent = 0.0
        self.tokens_used = 0
        self.steps = 0

    def check(self, step_tokens: int = 0, step_cost: float = 0.0) -> str | None:
        """Return error message if limit breached, else None."""
        self.steps += 1
        self.tokens_used += step_tokens
        self.spent += step_cost
        if self.steps > self.max_steps:
            return f"Step limit ({self.max_steps}) exceeded"
        if self.tokens_used > self.max_tokens:
            return f"Token budget ({self.max_tokens}) exceeded"
        if self.spent > self.max_spend:
            return f"Spend cap (${self.max_spend:.2f}) exceeded"
        return None


# ── Demo ───────────────────────────────────────────────────────────────────


def main() -> None:
    sup = Agent("supervisor", "supervisor", handler=lambda x: f"PLAN: {x[:40]}")
    workers = [
        Agent("coder", "worker", tools=["read", "write"], handler=lambda x: f"CODE: {x[:40]}"),
        Agent("reviewer", "worker", tools=["diff", "lint"], handler=lambda x: f"REVIEW: {x[:40]}"),
    ]

    print("=== Supervisor-Worker Pattern ===")
    orch = SupervisorWorker(sup, workers, max_steps=5)
    results = orch.run("Implement a todo-list REST API")
    for tid, result in results.items():
        print(f"  {tid}: {result}")

    print(f"\n  Trace: {len(orch.trace)} steps")
    for t in orch.trace:
        print(f"    {t}")

    print("\n=== Fan-Out Pattern ===")
    fan = FanOut(workers)
    for name, result in fan.run("What is the time complexity of quicksort?").items():
        print(f"  {name}: {result}")

    print("\n=== Guardrails ===")
    g = Guardrails(max_steps=3, max_tokens=1000, max_spend=0.50)
    for i in range(5):
        limit = g.check(step_tokens=800, step_cost=0.30)
        if limit:
            print(f"  BREACHED at step {i+1}: {limit}")
            break
    else:
        print("  All within limits")


if __name__ == "__main__":
    main()