"""
CrewAI integration for scandar-guard.

Hooks into CrewAI's step_callback and task_callback to monitor
agent behavior and send ASG telemetry.

Usage:
    from crewai import Crew, Agent, Task
    from scandar_guard.integrations.crewai import ScandarCrewAIMonitor

    monitor = ScandarCrewAIMonitor(api_key="sk_xxx")

    crew = Crew(
        agents=[...],
        tasks=[...],
        step_callback=monitor.step_callback,
        task_callback=monitor.task_callback,
    )
"""
from __future__ import annotations

from typing import Any, Optional

from ..asg_telemetry import AsgTelemetrySender
from ..config import GuardConfig
from ..inspector import MessageInspector
from ..tool_inspector import ToolCallInspector


class ScandarCrewAIMonitor:
    """Monitor CrewAI crew execution with scandar-guard."""

    def __init__(
        self,
        api_key: str,
        agent_prefix: str = "crewai",
        endpoint: Optional[str] = None,
    ) -> None:
        self._api_key = api_key
        self._agent_prefix = agent_prefix
        self._endpoint = endpoint
        self._msg_inspector = MessageInspector()
        self._tool_inspector = ToolCallInspector()
        self._senders: dict[str, AsgTelemetrySender] = {}
        self._session_keys: dict[str, str] = {}
        self._tool_counts: dict[str, int] = {}
        self._finding_counts: dict[str, dict[str, int]] = {}
        self._max_severity: dict[str, str] = {}

    def _get_sender(self, agent_role: str) -> AsgTelemetrySender:
        agent_id = f"{self._agent_prefix}-{agent_role}".replace(" ", "-").lower()
        if agent_id not in self._senders:
            self._senders[agent_id] = AsgTelemetrySender(
                api_key=self._api_key,
                agent_id=agent_id,
                endpoint=self._endpoint or "https://scandar.ai/api/v1/asg/events",
            )
        return self._senders[agent_id]

    def _get_session_key(self, agent_role: str) -> str:
        if agent_role not in self._session_keys:
            import uuid
            self._session_keys[agent_role] = f"crewai_{uuid.uuid4().hex[:12]}"
        return self._session_keys[agent_role]

    def step_callback(self, step_output: Any) -> None:
        """Called after each agent step in CrewAI."""
        try:
            agent_role = getattr(step_output, "agent", "unknown")
            if hasattr(agent_role, "role"):
                agent_role = agent_role.role

            sender = self._get_sender(str(agent_role))
            session_key = self._get_session_key(str(agent_role))

            # Check if this step involved a tool call
            tool_name = None
            if hasattr(step_output, "tool"):
                tool_name = str(step_output.tool)
            elif hasattr(step_output, "tool_name"):
                tool_name = str(step_output.tool_name)

            if tool_name:
                # Inspect the tool call
                result = self._tool_inspector.inspect_tool_call(
                    tool_name, {}, session_key
                )
                top_finding = result.findings[0] if result.findings else None
                sender.send_tool_call(
                    session_key, tool_name,
                    finding_category=top_finding.category if top_finding else None,
                    finding_severity=top_finding.severity if top_finding else None,
                )
                # Track for session_end aggregation
                role_key = str(agent_role)
                self._tool_counts[role_key] = self._tool_counts.get(role_key, 0) + 1
                for f in result.findings:
                    if role_key not in self._finding_counts:
                        self._finding_counts[role_key] = {}
                    self._finding_counts[role_key][f.category] = self._finding_counts[role_key].get(f.category, 0) + 1
                    sev_rank = {"critical": 4, "high": 3, "medium": 2, "low": 1, "info": 0}
                    curr = self._max_severity.get(role_key, "info")
                    if sev_rank.get(f.severity, 0) > sev_rank.get(curr, 0):
                        self._max_severity[role_key] = f.severity

            # Inspect the output text
            output_text = ""
            if hasattr(step_output, "output"):
                output_text = str(step_output.output)
            elif hasattr(step_output, "text"):
                output_text = str(step_output.text)

            if output_text:
                self._msg_inspector.inspect_response(output_text, session_key)
        except Exception:
            pass  # Never affect the agent

    def task_callback(self, task_output: Any) -> None:
        """Called after each task completes in CrewAI."""
        try:
            agent_role = "unknown"
            if hasattr(task_output, "agent"):
                agent_role = str(getattr(task_output.agent, "role", "unknown"))

            sender = self._get_sender(agent_role)
            session_key = self._get_session_key(agent_role)

            # Send session start (if first callback for this agent)
            if agent_role not in self._session_keys or self._session_keys.get(agent_role) == session_key:
                sender.send_session_start(session_key)

            # Aggregate findings for this agent
            role_key = agent_role
            tool_count = self._tool_counts.get(role_key, 0)
            findings = self._finding_counts.get(role_key, {})
            max_sev = self._max_severity.get(role_key, "safe")
            total_findings = sum(findings.values())
            threat_score = min(100, total_findings * 15 + (25 if max_sev == "critical" else 10 if max_sev == "high" else 0))
            classification = "critical" if max_sev == "critical" else "suspicious" if max_sev == "high" else "elevated" if total_findings > 0 else "safe"

            sender.send_session_end(
                session_key,
                message_count=0,
                tool_call_count=tool_count,
                threat_score=threat_score,
                classification=classification,
                findings_summary=findings,
            )

            # Reset accumulators for next task
            self._session_keys.pop(agent_role, None)
            self._tool_counts.pop(role_key, None)
            self._finding_counts.pop(role_key, None)
            self._max_severity.pop(role_key, None)
        except Exception:
            pass
