"""Demo agent — simple Anthropic SDK agent to test Litmus recording.

Usage:
    # Terminal 1: Start the Litmus proxy
    litmus proxy --mode record --output-dir ./traces

    # Terminal 2: Run this agent (it will route through the proxy)
    ANTHROPIC_BASE_URL=http://localhost:8787/anthropic python examples/demo_agent.py

    # Or use the SDK shim:
    python examples/demo_agent.py --auto
"""

from __future__ import annotations

import argparse
import json
import os
import sys

# Simple math "tools" the agent can call
TOOLS = [
    {
        "name": "calculate",
        "description": "Perform a math calculation. Input should be a mathematical expression.",
        "input_schema": {
            "type": "object",
            "properties": {
                "expression": {"type": "string", "description": "Math expression like '2 + 3 * 4'"},
            },
            "required": ["expression"],
        },
    },
    {
        "name": "get_weather",
        "description": "Get current weather for a city.",
        "input_schema": {
            "type": "object",
            "properties": {
                "city": {"type": "string", "description": "City name"},
            },
            "required": ["city"],
        },
    },
]


def handle_tool_call(name: str, input_data: dict) -> str:
    """Simulate tool execution."""
    if name == "calculate":
        try:
            # Safe eval for demo purposes only
            result = eval(input_data["expression"], {"__builtins__": {}})
            return json.dumps({"result": result})
        except Exception as e:
            return json.dumps({"error": str(e)})
    elif name == "get_weather":
        # Fake weather data
        city = input_data.get("city", "Unknown")
        return json.dumps({
            "city": city,
            "temperature": 72,
            "condition": "sunny",
            "humidity": 45,
        })
    return json.dumps({"error": f"Unknown tool: {name}"})


def run_agent(prompt: str = "What's 15 * 7? Also, what's the weather in San Francisco?"):
    """Run a simple agent loop using the Anthropic SDK."""
    try:
        import anthropic
    except ImportError:
        print("Error: anthropic SDK not installed. Run: pip install anthropic")
        sys.exit(1)

    client = anthropic.Anthropic()

    print(f"Agent prompt: {prompt}")
    print("-" * 60)

    messages = [{"role": "user", "content": prompt}]

    # Agent loop — up to 5 iterations
    for i in range(5):
        print(f"\n[Step {i + 1}] Calling Claude...")
        response = client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=1024,
            tools=TOOLS,
            messages=messages,
        )

        print(f"  Stop reason: {response.stop_reason}")

        # Check for tool use
        if response.stop_reason == "tool_use":
            # Process all tool calls in the response
            tool_results = []
            for block in response.content:
                if block.type == "tool_use":
                    print(f"  Tool call: {block.name}({json.dumps(block.input)})")
                    result = handle_tool_call(block.name, block.input)
                    print(f"  Tool result: {result}")
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result,
                    })

            # Add assistant response + tool results to messages
            messages.append({"role": "assistant", "content": response.content})
            messages.append({"role": "user", "content": tool_results})
        else:
            # Final response
            for block in response.content:
                if hasattr(block, "text"):
                    print(f"\n[Final answer]\n{block.text}")
            break

    print("\n" + "-" * 60)
    print("Agent finished.")


def main():
    parser = argparse.ArgumentParser(description="Litmus demo agent")
    parser.add_argument("--auto", action="store_true", help="Auto-start Litmus proxy via SDK shim")
    parser.add_argument("--prompt", default=None, help="Custom prompt")
    args = parser.parse_args()

    if args.auto:
        import litmus
        litmus.init(mode="record", output_dir="./traces")
        print("[Litmus] Proxy started, recording trace...\n")

    try:
        run_agent(args.prompt or "What's 15 * 7? Also, what's the weather in San Francisco?")
    finally:
        if args.auto:
            import litmus
            litmus.shutdown()
            print("\n[Litmus] Trace saved.")


if __name__ == "__main__":
    main()
