#!/usr/bin/env python3
"""A provider-free Agent Loop demonstration with two deterministic cases."""

from dataclasses import dataclass
import json
from pathlib import Path
from typing import Dict, Optional


@dataclass
class Decision:
    kind: str
    name: Optional[str] = None
    arguments: Optional[Dict] = None
    answer: Optional[str] = None


class ScriptedModel:
    def __init__(self, decisions):
        self.decisions = iter(decisions)

    def decide(self, _messages):
        return next(self.decisions)


def weather(city):
    return {"city": city, "condition": "rain", "rain_probability": 0.8}


def run_agent(model, max_steps=3):
    messages = [{"role": "user", "content": "上海明天下雨吗？"}]
    trace = []

    for step in range(1, max_steps + 1):
        try:
            decision = model.decide(messages)
        except StopIteration:
            return {"terminal": "failed", "reason": "model_stopped_without_answer", "trace": trace}

        if decision.kind == "final":
            trace.append({"step": step, "decision": "final", "answer": decision.answer})
            return {"terminal": "success", "answer": decision.answer, "trace": trace}

        if decision.kind != "tool" or decision.name != "weather":
            return {"terminal": "failed", "reason": "invalid_action", "trace": trace}

        observation = weather(**(decision.arguments or {}))
        trace.append({"step": step, "decision": "tool", "tool": "weather", "observation": observation})
        messages.append({"role": "tool", "content": json.dumps(observation, ensure_ascii=False)})

    return {"terminal": "budget_exhausted", "reason": "max_steps", "trace": trace}


success_model = ScriptedModel([
    Decision("tool", "weather", {"city": "上海"}),
    Decision("final", answer="明天有雨，降雨概率约 80%。"),
])

looping_model = ScriptedModel([
    Decision("tool", "weather", {"city": "上海"}),
    Decision("tool", "weather", {"city": "上海"}),
    Decision("tool", "weather", {"city": "上海"}),
    Decision("tool", "weather", {"city": "上海"}),
])

result = {
    "success_case": run_agent(success_model),
    "stopping_case": run_agent(looping_model),
}

output = Path(__file__).with_name("experiment-result.json")
output.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
print(json.dumps(result, ensure_ascii=False, indent=2))