"""Strands Agents Integration with Beacon.

BeaconStrandsHandler is a native callback handler that traces all Strands
Agent events (agent lifecycle, model calls, tool invocations) directly as
Beacon spans. No OpenTelemetry setup required.

Key: Pass the handler as ``callback_handler`` when constructing the Agent.

Requirements:
    pip install 'lumenova-beacon[strands]' python-dotenv

Setup:
    AWS_ACCESS_KEY_ID=your-aws-access-key
    AWS_SECRET_ACCESS_KEY=your-aws-secret-key
    AWS_DEFAULT_REGION=us-east-1
    BEACON_ENDPOINT=http://localhost:8000 (optional)
    BEACON_API_KEY=your-api-key (optional)
"""

import dotenv
from lumenova_beacon import BeaconClient
from lumenova_beacon.tracing.integrations.strands import BeaconStrandsHandler
from strands import Agent
from strands.models.bedrock import BedrockModel
from strands_tools import calculator, current_time

dotenv.load_dotenv()


# === Step 1: Initialize Beacon Client ===
beacon_client = BeaconClient()

# === Step 2: Create Callback Handler ===
# The handler captures the full agent lifecycle as Beacon spans:
#   agent_span (AGENT)
#   └── model_span (GENERATION)  — one per event-loop cycle
#   └── tool_span  (TOOL)        — one per tool invocation
handler = BeaconStrandsHandler(
    session_id="user-session-123",
    agent_name="Research Assistant",
)

# === Step 3: Configure Strands Agent ===
model = BedrockModel(
    model_id="us.anthropic.claude-3-5-haiku-20241022-v1:0",
)

agent = Agent(
    model=model,
    system_prompt="You are a helpful assistant that provides concise answers to questions.",
    tools=[calculator, current_time],
    callback_handler=handler,
)

# === Step 4: Run Agent — Spans Automatically Traced to Beacon ===
try:
    print("Query: What is 2 + 2, and what is the current time?")
    result = agent("What is 2 + 2, and what is the current time?")
    print(f"Response: {result.message}\n")
    print(f"Beacon Trace ID: {handler.trace_id}")
except Exception as e:
    print(f"Error: {e}\n")

# === Example 2: Follow-up (same agent maintains conversation history) ===
try:
    print("Query: Double that result.")
    result = agent("Double the addition result.")
    print(f"Response: {result.message}\n")
    print(f"Beacon Trace ID: {handler.trace_id}")
except Exception as e:
    print(f"Error: {e}\n")
