"""CrewAI Integration with Beacon.

BeaconCrewAIListener is a native event listener that traces the full CrewAI
execution lifecycle (crew, tasks, agents, LLM calls, tool usage) directly as
Beacon spans. No third-party instrumentation libraries required.

Key: Instantiate the listener BEFORE calling crew.kickoff().

Span hierarchy produced:
    crew_span   (AGENT)                — one per kickoff
    └── task_span   (TASK)            — one per task
        └── agent_span  (AGENT)       — one per agent execution
            ├── llm_span  (GENERATION) — each LLM call
            └── tool_span (TOOL)       — each tool invocation

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

Setup:
    OPENAI_API_KEY=your-openai-api-key   (or configure any LLM CrewAI supports)
    BEACON_ENDPOINT=http://localhost:8000 (optional)
    BEACON_API_KEY=your-api-key          (optional)
"""

import dotenv
from crewai import Agent, Crew, Process, Task
from crewai.tools import tool

from lumenova_beacon import BeaconClient
from lumenova_beacon.tracing.integrations.crewai import BeaconCrewAIListener

dotenv.load_dotenv()


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

# === Step 2: Create Listener ===
# Instantiate before kickoff — it registers itself with the CrewAI event bus.
listener = BeaconCrewAIListener(
    session_id="research-session-1",
    crew_name="AI Research Crew",
    metadata={"project": "beacon-demo"},
)


# === Step 3: Define Custom Tools ===
@tool("Search Simulator")
def search_simulator(query: str) -> str:
    """Simulates a web search and returns relevant results.
    Input should be a search query string."""
    mock_data = {
        "ai": "Recent AI highlights: reasoning models, agentic frameworks, multimodal capabilities.",
        "llm": "LLM trends: extended context windows, improved tool use, reduced hallucinations.",
    }
    for key, value in mock_data.items():
        if key in query.lower():
            return f"Results for '{query}': {value}"
    return f"Results for '{query}': No specific data found — try 'ai' or 'llm'."


@tool("Word Counter")
def word_counter(text: str) -> str:
    """Counts words and sentences in a text.
    Input should be any text string."""
    words = len(text.split())
    sentences = text.count(".") + text.count("!") + text.count("?")
    return f"{words} words, {sentences} sentences"


# === Step 4: Define Agents ===
researcher = Agent(
    role="Senior Research Analyst",
    goal="Uncover the latest developments in AI",
    backstory=(
        "You work at a leading tech think tank. Your expertise lies in identifying "
        "emerging trends and translating complex data into actionable insights."
    ),
    verbose=True,
    allow_delegation=False,
    tools=[search_simulator],
)

writer = Agent(
    role="Tech Content Strategist",
    goal="Craft an engaging summary of AI advancements",
    backstory=(
        "You are a renowned content strategist known for making complex topics "
        "accessible. You always verify your content quality before submitting."
    ),
    verbose=True,
    allow_delegation=False,
    tools=[word_counter],
)

# === Step 5: Define Tasks ===
research_task = Task(
    description=(
        "Research the latest AI advancements. Use the search simulator to gather "
        "information about recent AI and LLM developments. Summarise the top 3 trends."
    ),
    expected_output="A bullet-point summary of the top 3 AI trends.",
    agent=researcher,
)

writing_task = Task(
    description=(
        "Using the research findings, write a short 2-paragraph blog post about "
        "AI advancements. Then use the word counter tool to report the length."
    ),
    expected_output="A 2-paragraph blog post followed by its word count.",
    agent=writer,
)

# === Step 6: Create and Run Crew ===
crew = Crew(
    agents=[researcher, writer],
    tasks=[research_task, writing_task],
    verbose=True,
    process=Process.sequential,
)

try:
    result = crew.kickoff()
    print("\n=== CREW EXECUTION COMPLETED ===")
    print(result)
    print(f"\nBeacon Trace ID: {listener.trace_id}")
except Exception as e:
    print(f"\nError during crew execution: {e}")
    print(f"Beacon Trace ID: {listener.trace_id}")
