"""
CrewAI Integration Example — infralo-sdk

Demonstrates:
  1. Registering InfraloCrewAIListener globally on the CrewAI event bus.
  2. Running a Crew and having tool spans captured automatically.
  3. Tool error recording via ToolUsageErrorEvent.

Requirements:
    pip install infralo[crewai]

Run:
    cd examples/crewai
    uv run crewai_integration.py
    # or: python crewai_integration.py
"""

import asyncio
import logging
import os

from infralo import Infralo
from infralo.integrations.crewai import InfraloCrewAIListener


logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("infralo.examples.crewai")

INFRALO_API_KEY = os.environ.get(
    "INFRALO_API_KEY", "your_workspace_api_key"
)
INFRALO_ENDPOINT = os.environ.get("INFRALO_ENDPOINT", "http://localhost:8000")
MODEL_DEPLOYMENT = os.environ.get("INFRALO_MODEL_DEPLOYMENT", "your_model_name")

# ── 1. Infralo client ─────────────────────────────────────────────────────────

infralo = Infralo(
    api_key=INFRALO_API_KEY,
    endpoint=INFRALO_ENDPOINT,
    flush_interval=1.0,
)

# ── 2. Register the listener ONCE at module level ─────────────────────────────
#       InfraloCrewAIListener() self-registers on the CrewAI global event bus.

InfraloCrewAIListener()


async def run_crewai_example() -> None:
    try:
        from crewai import Agent, Crew, Task  # type: ignore[import-untyped]
        from crewai.tools import BaseTool  # type: ignore[import-untyped]
    except ImportError:
        logger.error("crewai is not installed. Run: pip install infralo[crewai]")
        return

    print("=== CrewAI + Infralo Integration Example (Async) ===")

    # ── 3. Define custom tools ────────────────────────────────────────────────

    class SearchWebTool(BaseTool):
        name: str = "search_web"
        description: str = "Search the web for up-to-date information on a topic."

        def _run(self, query: str) -> str:
            logger.info("  [tool] search_web(query=%r)", query)
            return f"Search results for '{query}': article 1, article 2."

        async def _arun(self, query: str) -> str:
            return self._run(query)

    class SummariseTool(BaseTool):
        name: str = "summarise"
        description: str = "Summarise a piece of text."

        def _run(self, text: str) -> str:
            logger.info("  [tool] summarise(text=%r...)", text[:40])
            return f"Summary: {text[:100]}..."

        async def _arun(self, text: str) -> str:
            return self._run(text)

    from crewai import LLM  # type: ignore[import-untyped]

    llm = LLM(
        model=f"openai/{MODEL_DEPLOYMENT}",
        base_url=f"{INFRALO_ENDPOINT.rstrip('/')}/v1",
        api_key=INFRALO_API_KEY,
    )

    # ── 4. Define agents & tasks ──────────────────────────────────────────────

    researcher = Agent(
        role="Research Analyst",
        goal="Gather comprehensive information on a given topic.",
        backstory="You are an expert researcher with access to the web.",
        tools=[SearchWebTool()],
        llm=llm,
        verbose=True,
    )

    writer = Agent(
        role="Content Writer",
        goal="Produce clear, concise summaries from research findings.",
        backstory="You are a skilled writer specialising in technology topics.",
        tools=[SummariseTool()],
        llm=llm,
        verbose=True,
    )

    research_task = Task(
        description="Research the current state of AI observability tools. You MUST use the search_web tool to look up current tools.",
        expected_output="A bullet-point list of key findings.",
        agent=researcher,
    )

    write_task = Task(
        description="Summarise the research findings. You MUST use the summarise tool to generate the summary.",
        expected_output="A 2-paragraph summary suitable for a blog post.",
        agent=writer,
    )

    crew = Crew(agents=[researcher, writer], tasks=[research_task, write_task], verbose=True)

    # ── 5. Run inside a trace ─────────────────────────────────────────────────

    with infralo.start_trace(session_id="crewai-demo") as trace:
        print(f"Trace ID: {trace.trace_id}")
        result = await crew.kickoff_async()
        print(f"\nCrew result:\n{result}")

    print("\nDone. Check Infralo dashboard for trace spans.")


if __name__ == "__main__":
    asyncio.run(run_crewai_example())
    infralo.flush()