#!/usr/bin/env python3
"""
Lesson 07 Demo — Tool Use with Claude
A self-contained calculator + unit_converter tool loop.
Run: python demo/tool_demo.py
"""

import os
import json
from dotenv import load_dotenv
import anthropic

load_dotenv()

# ── Tool definitions ──────────────────────────────────────────────────────────

TOOLS = [
    {
        "name": "calculator",
        "description": "Perform basic arithmetic: add, subtract, multiply, divide.",
        "input_schema": {
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": ["add", "subtract", "multiply", "divide"],
                    "description": "The arithmetic operation to perform."
                },
                "a": {"type": "number", "description": "First operand."},
                "b": {"type": "number", "description": "Second operand."}
            },
            "required": ["operation", "a", "b"]
        }
    },
    {
        "name": "unit_converter",
        "description": "Convert a value between common units (km/miles, kg/lb, celsius/fahrenheit).",
        "input_schema": {
            "type": "object",
            "properties": {
                "value": {"type": "number", "description": "The numeric value to convert."},
                "from_unit": {
                    "type": "string",
                    "enum": ["km", "miles", "kg", "lb", "celsius", "fahrenheit"],
                    "description": "Source unit."
                },
                "to_unit": {
                    "type": "string",
                    "enum": ["km", "miles", "kg", "lb", "celsius", "fahrenheit"],
                    "description": "Target unit."
                }
            },
            "required": ["value", "from_unit", "to_unit"]
        }
    }
]

# ── Tool execution ────────────────────────────────────────────────────────────

def execute_tool(tool_name: str, tool_input: dict) -> str:
    """Dispatch a tool call and return the result as a string."""
    if tool_name == "calculator":
        op = tool_input["operation"]
        a, b = tool_input["a"], tool_input["b"]
        if op == "add":
            result = a + b
        elif op == "subtract":
            result = a - b
        elif op == "multiply":
            result = a * b
        elif op == "divide":
            result = a / b if b != 0 else "Error: division by zero"
        else:
            result = f"Unknown operation: {op}"
        return str(result)

    elif tool_name == "unit_converter":
        value = tool_input["value"]
        from_u = tool_input["from_unit"]
        to_u = tool_input["to_unit"]
        conversions = {
            ("km", "miles"): lambda v: v * 0.621371,
            ("miles", "km"): lambda v: v * 1.60934,
            ("kg", "lb"): lambda v: v * 2.20462,
            ("lb", "kg"): lambda v: v * 0.453592,
            ("celsius", "fahrenheit"): lambda v: (v * 9 / 5) + 32,
            ("fahrenheit", "celsius"): lambda v: (v - 32) * 5 / 9,
        }
        key = (from_u, to_u)
        if key in conversions:
            converted = conversions[key](value)
            return f"{converted:.4f} {to_u}"
        else:
            return f"Unsupported conversion: {from_u} -> {to_u}"

    return f"Unknown tool: {tool_name}"

# ── Agentic tool loop ─────────────────────────────────────────────────────────

def run_tool_loop(user_message: str, max_iterations: int = 10) -> str:
    """Send a message to Claude with tools and handle the full tool-use loop."""
    client = anthropic.Anthropic()
    messages = [{"role": "user", "content": user_message}]

    print(f"\n{'='*60}")
    print(f"USER: {user_message}")
    print(f"{'='*60}")

    for iteration in range(max_iterations):
        response = client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=1024,
            tools=TOOLS,
            messages=messages,
        )

        print(f"\n── Round {iteration + 1} | stop_reason={response.stop_reason} ──")

        # Process each content block
        tool_results = []
        final_text = None

        for block in response.content:
            if block.type == "text":
                print(f"  CLAUDE: {block.text}")
                final_text = block.text
            elif block.type == "tool_use":
                print(f"  TOOL CALL: {block.name}({json.dumps(block.input)})")
                result = execute_tool(block.name, block.input)
                print(f"  TOOL RESULT: {result}")
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": result,
                })

        # If Claude is done talking, return
        if response.stop_reason == "end_turn":
            print(f"\n{'='*60}")
            return final_text or "(no text)"

        # Otherwise, feed tool results back
        messages.append({"role": "assistant", "content": response.content})
        messages.append({"role": "user", "content": tool_results})

    return "(max iterations reached)"

# ── Main ──────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("Lesson 07 Demo: Calculator + Unit Converter Tool Loop")
    print("=" * 60)

    # Demo 1: Simple calculation
    run_tool_loop("What is 1729 divided by 13?")

    # Demo 2: Unit conversion
    run_tool_loop("Convert 42 km to miles.")

    # Demo 3: Multi-step reasoning requiring both tools
    run_tool_loop(
        "I ran a marathon (42.195 km). Convert that to miles, "
        "then multiply the result by my pace of 8.5 minutes per mile "
        "to get my total time in minutes."
    )

    print("\n\nDemo complete.")