#!/usr/bin/env python3
"""Robust LLM function calling with schema validation, repair, and capped retries.

Demonstrates the full validate→repair→retry→execute loop.
"""

from __future__ import annotations

import json
from typing import Any, Callable

import jsonschema


# ── Tool definitions (JSON Schema) ──────────────────────────────────────

TOOL_SCHEMAS = {
    "send_email": {
        "type": "object",
        "properties": {
            "to": {"type": "string", "format": "email", "description": "Recipient email address"},
            "subject": {"type": "string", "minLength": 1, "description": "Email subject line"},
            "body": {"type": "string", "description": "Plain text email body"},
        },
        "required": ["to", "subject", "body"],
        "additionalProperties": False,
    },
    "get_weather": {
        "type": "object",
        "properties": {
            "city": {"type": "string", "minLength": 1, "description": "City name"},
            "units": {
                "type": "string",
                "enum": ["celsius", "fahrenheit"],
                "description": "Temperature unit",
            },
        },
        "required": ["city"],
        "additionalProperties": False,
    },
}


# ── Tool implementations ────────────────────────────────────────────────

def _send_email(to: str, subject: str, body: str) -> str:
    return f"Email sent to {to}: '{subject}'"

def _get_weather(city: str, units: str = "celsius") -> str:
    return f"Weather in {city}: 22°{'C' if units == 'celsius' else 'F'}"

TOOL_IMPLS: dict[str, Callable] = {
    "send_email": _send_email,
    "get_weather": _get_weather,
}


# ── Validation & execution ──────────────────────────────────────────────

def validate_and_execute(
    tool_name: str, arguments: dict, max_retries: int = 3
) -> dict:
    """Validate arguments against schema, retry with repair, execute."""
    schema = TOOL_SCHEMAS.get(tool_name)
    if schema is None:
        return {"tool": tool_name, "error": f"Unknown tool: {tool_name}"}

    retries = 0
    while retries <= max_retries:
        try:
            jsonschema.validate(instance=arguments, schema=schema)
            result = TOOL_IMPLS[tool_name](**arguments)
            return {"tool": tool_name, "result": result, "retries": retries}
        except jsonschema.ValidationError as e:
            if retries < max_retries:
                arguments = _repair_arguments(arguments, schema)
                retries += 1
            else:
                return {
                    "tool": tool_name,
                    "error": f"Validation failed after {max_retries} retries: {e.message}",
                    "retries": retries,
                }
        except Exception as e:
            return {"tool": tool_name, "error": str(e), "retries": retries}

    return {"tool": tool_name, "error": "Max retries exceeded"}


def _repair_arguments(args: dict, schema: dict) -> dict:
    """Simple repair: drop unknown keys, ensure required keys exist."""
    known = set(schema.get("properties", {}).keys())
    repaired = {k: v for k, v in args.items() if k in known}
    for req in schema.get("required", []):
        if req not in repaired:
            repaired[req] = None
    return repaired


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

if __name__ == "__main__":
    test_cases = [
        ("send_email", {"to": "a@b.com", "subject": "Hi", "body": "Hello"}),
        ("send_email", {"to": "a@b.com", "body": "Hello"}),  # missing subject
        ("send_email", {"to": "a@b.com", "subject": "Hi", "body": "Hello", "cc": "x@y.com"}),
        ("get_weather", {"city": "London", "units": "kelvin"}),  # invalid enum
        ("delete_everything", {}),  # unknown tool
    ]

    for tool, args in test_cases:
        result = validate_and_execute(tool, args)
        status = "✓" if "error" not in result else "✗"
        detail = result.get("result") or result.get("error")
        print(f"{status} {tool}({json.dumps(args)}) → {detail} [{result.get('retries', 0)} retries]")