"""
CrewAI Complex Nesting Example — infralo-sdk

Demonstrates:
  1. Multi-crew delegation: a Root Crew delegating tasks to a Research Crew.
  2. Nested LLM and tool trace parenting.
  3. Automatic context propagation across nested crew boundaries.

Run:
    cd examples/crewai
    uv run crewai_complex_nesting.py
"""

import asyncio
import logging
import os

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

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

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")

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

# Register the global CrewAI listener
InfraloCrewAIListener()


# ── Nested Tool (CrewAI Style) ────────────────────────────────────────────────

from crewai.tools import BaseTool  # type: ignore[import-untyped]


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("    [Nested Tool] search_web(query=%r)", query)
        return (
            f"Search results for '{query}': CrewAI sub-crews are fully traceable "
            f"using the Infralo event bus and monkeypatched BaseTool context propagation."
        )

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


# ── Root Crew Tool (Delegates to Nested Crew) ─────────────────────────────────


class DelegateResearchTool(BaseTool):
    name: str = "delegate_research"
    description: str = (
        "Delegate in-depth research about a specific topic to a specialized Research Crew."
    )

    async def _arun(self, topic: str) -> str:
        logger.info("  [Root Tool] delegating research for topic=%r...", topic)

        from crewai import LLM, Agent, Crew, Task  # type: ignore[import-untyped]

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

        # Create the specialized Research Agent
        researcher = Agent(
            role="Research Specialist",
            goal="Search the web and compile insights on the given topic.",
            backstory="You are an expert analyst specialized in deep research.",
            tools=[SearchWebTool()],
            llm=llm,
            verbose=True,
        )

        research_task = Task(
            description=f"Perform deep research on the topic: {topic}. You MUST use the search_web tool to find insights.",
            expected_output="A concise summary of the key findings.",
            agent=researcher,
        )

        research_crew = Crew(
            agents=[researcher],
            tasks=[research_task],
            verbose=True,
        )

        # Run the nested crew asynchronously. The patched BaseTool carries parent span context!
        result = await research_crew.kickoff_async()

        logger.info("  [Root Tool] delegation completed.")
        return str(result)

    def _run(self, topic: str) -> str:
        from infralo.trace import _active_span, _active_trace

        # Capture current thread's context variables (set by the SDK's
        # to_structured_tool patch) so they survive the cross-thread hop below.
        span = _active_span.get()
        trace = _active_trace.get()

        async def wrapped_arun() -> str:
            # Re-apply context vars on the main event loop task so that any
            # async LLM calls made by the child crew see the correct active span.
            span_token = _active_span.set(span)
            trace_token = _active_trace.set(trace)
            try:
                return await self._arun(topic)
            finally:
                _active_span.reset(span_token)
                _active_trace.reset(trace_token)

        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = None
        if loop and loop.is_running():
            future = asyncio.run_coroutine_threadsafe(wrapped_arun(), loop)
            return future.result()
        else:
            return asyncio.run(self._arun(topic))


# ── Main Execution ────────────────────────────────────────────────────────────


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

    print("=== Running CrewAI Complex Nesting Example ===")

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

    with infralo.start_trace(session_id="crewai-complex-nesting") as trace:
        print(f"Trace ID: {trace.trace_id}")

        # Create the Root Agent
        root_agent = Agent(
            role="Root Coordinator",
            goal="Coordinate and answer high-level research questions by delegating to specialized crews.",
            backstory="You are the lead manager responsible for organizing research teams.",
            tools=[DelegateResearchTool()],
            llm=llm,
            verbose=True,
        )

        coordinator_task = Task(
            description="Perform deep research on multi-agent collaboration frameworks. You MUST delegate the research by using the delegate_research tool.",
            expected_output="A concise answer explaining the findings.",
            agent=root_agent,
        )

        root_crew = Crew(
            agents=[root_agent],
            tasks=[coordinator_task],
            verbose=True,
        )

        result = await root_crew.kickoff_async()
        print(f"\nFinal Root Crew response:\n{result}")

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


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