"""LangChain built-in agent with middleware and BeaconLangGraphHandler tracing.

Uses langchain.agents.create_agent (LangChain v1.0+) with @wrap_model_call and
@wrap_tool_call middleware for dynamic model selection and tool error handling.

The agent is a compiled LangGraph graph under the hood, so BeaconLangGraphHandler
traces it with the same span hierarchy as manually-built LangGraph agents.

Requirements:
    pip install lumenova-beacon[langchain]
    pip install langchain langchain-anthropic
    Set ANTHROPIC_API_KEY environment variable
"""

import dotenv
from langchain.agents import create_agent
from langchain.agents.middleware import wrap_model_call, wrap_tool_call
from langchain_anthropic import ChatAnthropic
from langchain_core.tools import tool

from lumenova_beacon import BeaconLangGraphHandler

dotenv.load_dotenv()

# ---------------------------------------------------------------------------
# Models -- used by the model-selection middleware
# ---------------------------------------------------------------------------
basic_model = ChatAnthropic(model="claude-sonnet-4-20250514", temperature=0)
advanced_model = ChatAnthropic(model="claude-sonnet-4-20250514", temperature=0.3)

# ---------------------------------------------------------------------------
# Tools (mock -- no external API keys needed beyond the LLM)
# ---------------------------------------------------------------------------
# NOTE: We do NOT add @trace here. BeaconLangGraphHandler already captures tool
# spans via on_tool_start / on_tool_end. Adding @trace would create duplicates.


@tool
def get_weather(location: str) -> str:
    """Get the current weather for a location."""
    weather_data = {
        "new york": "Sunny, 72\u00b0F",
        "london": "Cloudy, 61\u00b0F",
        "tokyo": "Rainy, 68\u00b0F",
        "paris": "Partly cloudy, 65\u00b0F",
        "san francisco": "Foggy, 58\u00b0F",
        "miami": "Hot and humid, 85\u00b0F",
    }
    return weather_data.get(
        location.lower(), f"Weather data not available for {location}"
    )


@tool
def calculator(expression: str) -> str:
    """Evaluate a mathematical expression. Example: '2 + 2' or '10 * 5'."""
    try:
        result = eval(expression, {"__builtins__": {}}, {})  # noqa: S307
        return f"The result is: {result}"
    except Exception as e:
        return f"Error calculating: {str(e)}"


@tool
def search_web(query: str) -> str:
    """Search the web for information about a topic."""
    responses = {
        "python": "Python is a high-level programming language known for its simplicity and readability.",
        "ai": "Artificial Intelligence refers to computer systems that can perform tasks requiring human intelligence.",
        "langgraph": "LangGraph is a library for building stateful, multi-actor applications with LLMs.",
        "anthropic": "Anthropic is an AI safety company that created the Claude family of models.",
    }
    for key, value in responses.items():
        if key in query.lower():
            return value
    return f"Search results for '{query}': General information available online."


tools = [get_weather, calculator, search_web]

# ---------------------------------------------------------------------------
# Middleware
# ---------------------------------------------------------------------------


@wrap_model_call
def log_and_select_model(request, handler):
    """Log model invocations and switch models based on conversation length.

    Short conversations (< 5 messages) use the basic model. Longer ones switch
    to a higher-temperature configuration for more creative responses.
    """
    message_count = len(request.state["messages"])
    print(f"  [middleware] Model call with {message_count} messages in state")

    if message_count > 5:
        print("  [middleware] Switching to advanced model for complex conversation")
        return handler(request.override(model=advanced_model))

    return handler(request)


@wrap_tool_call
def handle_tool_errors(request, handler):
    """Wrap tool calls with logging and graceful error handling."""
    tool_call = request.tool_call
    print(
        f"  [middleware] Executing tool: {tool_call['name']} "
        f"with args: {tool_call['args']}"
    )
    try:
        result = handler(request)
        print(f"  [middleware] Tool {tool_call['name']} succeeded")
        return result
    except Exception as e:
        error_msg = (
            f"Tool '{tool_call['name']}' failed: {e}. "
            "Please try a different approach."
        )
        print(f"  [middleware] Tool error caught: {error_msg}")
        return error_msg


# ---------------------------------------------------------------------------
# BeaconLangGraphHandler -- traces all LangGraph operations automatically
# ---------------------------------------------------------------------------
handler = BeaconLangGraphHandler(
    environment="development",
    agent_name="builtin-agent-demo",
    metadata={
        "app_name": "langchain-builtin-agent-demo",
        "version": "1.0.0",
        "agent_type": "create_agent",
    },
)

# ---------------------------------------------------------------------------
# Agent creation
# ---------------------------------------------------------------------------
agent = create_agent(
    model=basic_model,
    tools=tools,
    system_prompt=(
        "You are a helpful assistant with access to weather, calculator, and "
        "web search tools. Use the appropriate tool to answer the user's "
        "questions. When asked about weather, use the get_weather tool. "
        "When asked to calculate, use the calculator tool."
    ),
    middleware=[log_and_select_model, handle_tool_errors],
)


# ---------------------------------------------------------------------------
# Run queries
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    # === Query 1: Single tool call (weather) ===
    print("=" * 60)
    print("Query 1: Weather lookup (single tool call)")
    print("=" * 60)

    result = agent.invoke(
        {"messages": [{"role": "user", "content": "What's the weather in London?"}]},
        config={"callbacks": [handler]},
    )
    print(f"Response: {result['messages'][-1].content}\n")

    # === Query 2: Multi-tool use (weather + calculator) ===
    print("=" * 60)
    print("Query 2: Multi-tool (weather + calculator)")
    print("=" * 60)

    result = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "What's the weather in Tokyo and calculate 25 * 4 + 10?",
                }
            ]
        },
        config={"callbacks": [handler]},
    )
    print(f"Response: {result['messages'][-1].content}\n")

    # === Query 3: Search with reasoning ===
    print("=" * 60)
    print("Query 3: Search + reasoning")
    print("=" * 60)

    result = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "Search for information about Anthropic and tell me what you found.",
                }
            ]
        },
        config={"callbacks": [handler]},
    )
    print(f"Response: {result['messages'][-1].content}")


# Check your Beacon dashboard to see:
#
# Trace hierarchy for each query:
# - Root span with span.type = "agent" (LangGraph auto-detected)
#   - "agent" node span (chain) -- model reasoning
#     - ChatAnthropic LLM span (generation) -- actual Claude API call
#   - "tools" node span (chain) -- tool execution
#     - Individual tool spans (get_weather, calculator, search_web)
#   - "agent" node span (chain) -- model processes tool results
#     - ChatAnthropic LLM span (generation) -- final response
#
# Middleware does NOT create separate spans -- it wraps existing operations.
# The @wrap_model_call middleware fires inside the LLM call flow.
# The @wrap_tool_call middleware fires inside the tool execution flow.
#
# Custom metadata on all spans (using OTEL semantic conventions):
# - deployment.environment.name = "development"
# - gen_ai.agent.name = "builtin-agent-demo"
# - gen_ai.agent.id = "builtin-agent-001"
# - gen_ai.agent.description = "A demo agent using langchain create_agent..."
# - beacon.metadata.app_name = "langchain-builtin-agent-demo"
# - beacon.metadata.version = "1.0.0"
# - beacon.metadata.agent_type = "create_agent"
#
# LLM spans include auto-extracted attributes:
# - gen_ai.provider.name = "anthropic"
# - gen_ai.request.model = "claude-sonnet-4-20250514"
# - gen_ai.usage.input_tokens / output_tokens
# - gen_ai.request.temperature
