"""
CrewAI Integration for Agendex Governance.

Provides a simplified integration surface for CrewAI agents with
automatic reasoning capture from step callbacks.

Usage:
    from agendex.crewai import governed_crew

    # Create governed context
    client, step_callback = governed_crew(task="research_task")

    # Get governed clients
    db = client.db()
    http = client.http()

    # Attach callback to crew
    crew = Crew(
        agents=[...],
        tasks=[...],
        step_callback=step_callback,
    )
"""
from __future__ import annotations

from typing import Any, Callable, Dict, Optional, Tuple

from .client import AgendexClient
from .context import clear_reasoning, set_reasoning
from .crm import GovernedCRMClient
from .db import GovernedDBClient
from .http import GovernedHTTPSession
from .langchain import GovernedContext
from .s3 import GovernedS3Client


class CrewAIStepCallback:
    """
    CrewAI step callback that captures reasoning into contextvars.

    When a crew step executes, this callback extracts reasoning from
    the step output and stores it in a contextvar. Governed clients
    automatically read from this contextvar.

    Usage:
        client, step_callback = governed_crew(task="research")
        crew = Crew(..., step_callback=step_callback)
    """

    def __init__(self, on_event: Optional[Callable[[Dict], None]] = None):
        """
        Initialize the step callback.

        Args:
            on_event: Optional callback for step events
        """
        self.on_event = on_event

    def __call__(self, step_output: Any) -> None:
        """
        Called for each crew step.

        Extracts reasoning from step output and stores in contextvar.
        """
        reasoning = None

        # Try to extract reasoning from different step output formats
        if hasattr(step_output, "thought"):
            reasoning = step_output.thought
        elif hasattr(step_output, "reasoning"):
            reasoning = step_output.reasoning
        elif hasattr(step_output, "log"):
            reasoning = step_output.log
        elif isinstance(step_output, dict):
            reasoning = (
                step_output.get("thought")
                or step_output.get("reasoning")
                or step_output.get("log")
            )

        if reasoning:
            set_reasoning(str(reasoning))
            if self.on_event:
                self.on_event({"type": "reasoning", "content": str(reasoning)})

    def on_task_complete(self) -> None:
        """
        Called when a task completes.

        Clears reasoning to prevent leaking to next task.
        """
        clear_reasoning()


def governed_crew(
    task: str,
    base_url: Optional[str] = None,
    agent_id: Optional[str] = None,
    token: Optional[str] = None,
    on_event: Optional[Callable[[Dict], None]] = None,
) -> Tuple[GovernedContext, CrewAIStepCallback]:
    """
    Create a governed crew context with auto-reasoning capture.

    This is the main entry point for CrewAI integration. It returns
    a GovernedContext for creating governed clients and a step callback
    for attaching to the Crew.

    Args:
        task: Default task context for all operations
        base_url: Override Agendex URL (default: from env)
        agent_id: Override agent ID (default: from env)
        token: Override auth token (default: from env)
        on_event: Optional callback for governance events

    Returns:
        Tuple of (GovernedContext, CrewAIStepCallback)

    Example:
        client, step_callback = governed_crew(task="research")
        db = client.db()
        crew = Crew(
            agents=[...],
            tasks=[...],
            step_callback=step_callback,
        )
    """
    agendex = AgendexClient(
        base_url=base_url,
        agent_id=agent_id,
        token=token,
    )
    context = GovernedContext(agendex, task, on_event=on_event)
    callback = CrewAIStepCallback(on_event=on_event)
    return context, callback


__all__ = [
    "CrewAIStepCallback",
    "governed_crew",
]
