"""
Agent Status integration for CrewAI.

Provides automatic monitoring for CrewAI crews, agents, and tasks.

Usage:

    # Option 1: Event Listener (recommended)
    from crewai import Crew, Agent, Task
    from agent_status.integrations.crewai import AgentStatusEventListener

    listener = AgentStatusEventListener(
        api_key="fb_live_xxx",
        crew_name="My Research Crew",
    )

    crew = Crew(agents=[...], tasks=[...])
    result = crew.kickoff()

    print(listener.metrics)

    # Option 2: Expose for outside-in monitoring
    from agent_status.integrations.crewai import expose

    expose(crew, api_key="fb_live_xxx", crew_name="My Crew")
"""

from __future__ import annotations

import time
import logging
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, TYPE_CHECKING

CREWAI_AVAILABLE = False
crewai_event_bus = None

try:
    from crewai.events import BaseEventListener, crewai_event_bus
    from crewai.events import (
        CrewKickoffStartedEvent,
        CrewKickoffCompletedEvent,
        CrewKickoffFailedEvent,
        TaskCompletedEvent,
        ToolUsageFinishedEvent,
        ToolUsageErrorEvent,
    )
    CREWAI_AVAILABLE = True
except ImportError:
    BaseEventListener = object

import agent_status
from agent_status.client import AgentStatusClient

logger = logging.getLogger(__name__)


class AgentStatusEventListener(BaseEventListener if CREWAI_AVAILABLE else object):
    """
    CrewAI event listener that reports metrics to Agent Status.

    Automatically tracks:
    - Crew executions (total time, success/failure)
    - Agent executions (per-agent timing and results)
    - Task completions (per-task timing and output)
    - Tool usage and errors

    Args:
        api_key: Your API key
        crew_name: Name for this crew in the dashboard
        base_url: API base URL

    Example:
        from crewai import Crew
        from agent_status.integrations.crewai import AgentStatusEventListener

        listener = AgentStatusEventListener(
            api_key="fb_live_xxx",
            crew_name="Research Crew",
        )

        crew = Crew(agents=[researcher, writer], tasks=[research_task, write_task])
        result = crew.kickoff()

        print(listener.metrics)
    """

    def __init__(
        self,
        api_key: str,
        crew_name: str = "CrewAI Crew",
        base_url: str = "https://api.fabric.carmel.so",
        **kwargs,
    ):
        if not CREWAI_AVAILABLE:
            raise ImportError(
                "CrewAI is not installed. Install with: pip install crewai"
            )

        super().__init__(**kwargs)

        self.api_key = api_key
        self.crew_name = crew_name
        self.base_url = base_url

        self._client = AgentStatusClient(api_key=api_key, base_url=base_url)

        self._metrics: Dict[str, Any] = {
            "total_executions": 0,
            "successful_executions": 0,
            "failed_executions": 0,
            "total_latency_ms": 0,
            "agent_executions": 0,
            "task_completions": 0,
            "tool_calls": 0,
            "errors": 0,
        }

        self._current_execution: Optional[Dict[str, Any]] = None
        self._execution_start: Optional[float] = None

        self._register_handlers()

    def _register_handlers(self) -> None:
        """Register event handlers with CrewAI's event bus."""
        if not CREWAI_AVAILABLE or crewai_event_bus is None:
            return

        try:
            crewai_event_bus.on(CrewKickoffStartedEvent)(self._on_crew_start)
            crewai_event_bus.on(CrewKickoffCompletedEvent)(self._on_crew_complete)
            crewai_event_bus.on(CrewKickoffFailedEvent)(self._on_crew_failed)
            crewai_event_bus.on(TaskCompletedEvent)(self._on_task_complete)
            crewai_event_bus.on(ToolUsageFinishedEvent)(self._on_tool_finished)
            crewai_event_bus.on(ToolUsageErrorEvent)(self._on_tool_error)
            logger.debug("CrewAI event handlers registered successfully")
        except Exception as e:
            logger.warning(f"Could not register CrewAI event handlers: {e}")

    def setup_listeners(self, bus) -> None:
        """Required by BaseEventListener - handlers are registered in __init__ instead."""
        pass

    @property
    def metrics(self) -> Dict[str, Any]:
        """Get current metrics summary."""
        total = self._metrics["total_executions"]
        return {
            **self._metrics,
            "avg_latency_ms": (
                self._metrics["total_latency_ms"] / total if total > 0 else 0
            ),
            "success_rate": (
                self._metrics["successful_executions"] / total if total > 0 else 0
            ),
        }

    def _on_crew_start(self, source, event):
        """Called when crew execution starts."""
        self._execution_start = time.time()
        self._current_execution = {
            "crew_name": self.crew_name,
            "started_at": datetime.now(timezone.utc).isoformat(),
        }
        logger.debug(f"Crew started: {self.crew_name}")

    def _on_crew_complete(self, source, event):
        """Called when crew execution completes successfully."""
        if self._execution_start:
            latency_ms = (time.time() - self._execution_start) * 1000
            self._metrics["total_latency_ms"] += latency_ms
            logger.debug(f"Crew completed: {self.crew_name}, latency={latency_ms:.0f}ms")

        self._metrics["total_executions"] += 1
        self._metrics["successful_executions"] += 1
        self._current_execution = None
        self._execution_start = None

    def _on_crew_failed(self, source, event):
        """Called when crew execution fails."""
        if self._execution_start:
            latency_ms = (time.time() - self._execution_start) * 1000
            self._metrics["total_latency_ms"] += latency_ms

        self._metrics["total_executions"] += 1
        self._metrics["failed_executions"] += 1
        self._metrics["errors"] += 1
        self._current_execution = None
        self._execution_start = None
        logger.error(f"Crew failed: {self.crew_name}")

    def _on_agent_complete(self, source, event):
        """Called when an agent completes execution."""
        self._metrics["agent_executions"] += 1
        logger.debug("Agent execution completed")

    def _on_task_complete(self, source, event):
        """Called when a task completes."""
        self._metrics["task_completions"] += 1
        logger.debug("Task completed")

    def _on_tool_finished(self, source, event):
        """Called when a tool finishes."""
        self._metrics["tool_calls"] += 1
        logger.debug("Tool call completed")

    def _on_tool_error(self, source, event):
        """Called when a tool errors."""
        self._metrics["tool_calls"] += 1
        self._metrics["errors"] += 1
        logger.error("Tool call failed")

    def reset_metrics(self) -> None:
        """Reset all tracked metrics."""
        self._metrics = {
            "total_executions": 0,
            "successful_executions": 0,
            "failed_executions": 0,
            "total_latency_ms": 0,
            "agent_executions": 0,
            "task_completions": 0,
            "tool_calls": 0,
            "errors": 0,
        }

    def report(self) -> Dict[str, Any]:
        """Generate a summary report of metrics."""
        return {
            "crew_name": self.crew_name,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "metrics": self.metrics,
        }


# Backward compatibility alias
RoraEventListener = AgentStatusEventListener


def expose(
    crew: Any,
    api_key: str,
    crew_name: str = "CrewAI Crew",
    *,
    input_key: str = "input",
    tunnel_url: str = "wss://rora-tunnel.carmel.so/ws/connect",
    blocking: bool = True,
) -> Optional[str]:
    """
    Expose a CrewAI crew for outside-in monitoring.

    This connects your local CrewAI crew to the distributed monitoring network,
    allowing nodes around the world to probe your crew.

    Args:
        crew: A CrewAI Crew instance
        api_key: Your API key
        crew_name: Name for this crew in the dashboard
        input_key: The input key to use when invoking the crew (default: "input")
        tunnel_url: WebSocket URL for the tunnel relay
        blocking: If True, blocks until interrupted. If False, runs in background.

    Returns:
        The public URL where the crew can be reached (if non-blocking)

    Example:
        from crewai import Crew, Agent, Task
        from agent_status.integrations.crewai import expose

        crew = Crew(agents=[...], tasks=[...])
        expose(crew, api_key="fb_live_xxx", crew_name="My Research Crew")
    """
    if not CREWAI_AVAILABLE:
        raise ImportError(
            "CrewAI is not installed. Install with: pip install crewai"
        )

    from agent_status.integrations.langchain import _TunnelHandler

    class CrewWrapper:
        def __init__(self, crew):
            self.crew = crew

        def invoke(self, inputs):
            input_value = inputs.get(input_key, "")
            result = self.crew.kickoff(inputs={input_key: input_value})
            return {"output": str(result)}

    handler = _TunnelHandler(
        chain=CrewWrapper(crew),
        api_key=api_key,
        agent_name=crew_name,
        input_key=input_key,
        tunnel_url=tunnel_url,
        auto_register=True,
    )

    if blocking:
        handler.run()
        return None
    else:
        import threading
        thread = threading.Thread(target=handler.run, daemon=True)
        thread.start()
        return handler.public_url
