"""
Anthropic (Claude) + Agent-Gantry integration demo.

Demonstrates how to use Agent-Gantry with Anthropic's Claude API, including:
- Dynamic tool retrieval with Anthropic schema conversion
- Using the @with_semantic_tools decorator with dialect="anthropic"
- Tool execution handling specific to Anthropic's response format
"""

import asyncio
import os
from typing import Any

from dotenv import load_dotenv

from agent_gantry import AgentGantry, set_default_gantry, with_semantic_tools
from agent_gantry.schema.execution import ToolCall
from agent_gantry.schema.query import ConversationContext, ToolQuery

# Load environment variables
load_dotenv()


async def main() -> None:
    print("=== Agent-Gantry + Anthropic (Claude) Integration Demo ===\n")

    # 1. Check for API Key
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print("❌ Error: ANTHROPIC_API_KEY not found in environment.")
        print("   Please set it in your .env file.")
        return

    # 2. Initialize Gantry
    try:
        from agent_gantry.adapters.embedders.nomic import NomicEmbedder

        gantry = AgentGantry(embedder=NomicEmbedder())
        print("✅ Initialized with Nomic Embeddings")
    except ImportError:
        gantry = AgentGantry()
        print("⚠️  Initialized with Simple Embeddings")

    # 3. Register Tools
    @gantry.register(tags=["system"])
    def get_server_status(env: str) -> str:
        """Check the status of a server environment (prod/staging)."""
        return f"Environment '{env}' is HEALTHY. CPU: 45%"

    await gantry.sync()
    print(f"✅ Registered {gantry.tool_count} tools\n")

    # 4. Initialize Anthropic Client
    from anthropic import AsyncAnthropic

    client = AsyncAnthropic(api_key=api_key)

    # --- Scenario A: Dynamic Retrieval with Manual Schema Conversion ---
    print("--- Scenario A: Dynamic Retrieval with Anthropic Schema ---")
    query = "Check the production server status"
    print(f"User Query: '{query}'")

    # A. Retrieve Tools (Returns RetrievalResult)
    # We use the lower-level `retrieve` method to get the internal ToolDefinition objects
    # so we can convert them to Anthropic's format.
    retrieval_result = await gantry.retrieve(
        ToolQuery(context=ConversationContext(query=query), limit=1, score_threshold=0.1)
    )

    # B. Convert to Anthropic Schema
    # Agent-Gantry provides a helper method `to_anthropic_schema()` on ToolDefinition
    anthropic_tools = [t.tool.to_anthropic_schema() for t in retrieval_result.tools]

    print(f"Gantry retrieved {len(anthropic_tools)} tool(s)")
    # print(json.dumps(anthropic_tools, indent=2)) # Uncomment to see schema

    # C. Call Claude
    response = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=1024,
        messages=[{"role": "user", "content": query}],
        tools=anthropic_tools,
    )

    # D. Handle Tool Use
    # Anthropic returns a list of content blocks. We look for 'tool_use'.
    for block in response.content:
        if block.type == "tool_use":
            print(f"Claude decided to call: {block.name}({block.input})")

            # Execute securely via Gantry
            result = await gantry.execute(ToolCall(tool_name=block.name, arguments=block.input))
            print(f"Execution Result: {result.result}")

    # --- Scenario B: Using the Decorator (RECOMMENDED) ---
    print("\n--- Scenario B: Using @with_semantic_tools Decorator (RECOMMENDED) ---")

    # Set default gantry for decorator usage
    set_default_gantry(gantry)

    # The decorator handles retrieval AND schema conversion (dialect="anthropic")
    # Note: score_threshold=0.1 is used because we may be using SimpleEmbedder.
    # With Nomic or OpenAI embeddings, use the default (0.5).
    @with_semantic_tools(limit=1, dialect="anthropic", score_threshold=0.1)
    async def chat_with_claude(
        messages: list[dict[str, str]],
        tools: list[dict[str, Any]] | None = None,
    ):
        print(f"   [Decorator] Injected {len(tools) if tools else 0} tools (Anthropic format)")
        return await client.messages.create(
            model="claude-sonnet-4-5", max_tokens=1024, messages=messages, tools=tools
        )

    query_dec = "Check staging status"
    print(f"User Query: '{query_dec}'")

    response_dec = await chat_with_claude(messages=[{"role": "user", "content": query_dec}])

    for block in response_dec.content:
        if block.type == "tool_use":
            print(f"Claude decided to call: {block.name}")


if __name__ == "__main__":
    asyncio.run(main())
