"""
AutoGen and CrewAI multi-agent orchestrator adapters for llm-guard-kit.

These adapters bridge the gap between multi-agent frameworks and AgentGuard's
confidence scoring.  A2ATrustObject handles handoff semantics; these adapters
capture each agent's trace and call score_chain() transparently.

AutoGenAdapter
--------------
Converts AutoGen GroupChat message history to normalized AgentGuard steps.
Supports ``autogen`` and ``autogen_agentchat`` message formats.
Attach as a ``register_hook()`` callback on GroupChat, or score offline from
a messages list with ``score_from_messages()``.

CrewAIAdapter
-------------
Converts CrewAI task execution output to AgentGuard steps.
Supports CrewAI TaskOutput objects and raw dict logs.
Use ``score_from_task_output()`` for post-hoc scoring, or wrap a task with
``wrap_task()`` to automatically score on completion.

Usage
-----
    from llm_guard import AgentGuard, AutoGenAdapter, CrewAIAdapter

    guard = AgentGuard(api_key="sk-ant-...")

    # AutoGen
    adapter = AutoGenAdapter(guard)
    result = adapter.score_from_messages(
        messages=groupchat.messages,
        question="What is the capital of France?",
        final_answer=last_message["content"],
    )
    print(result.risk_score, result.confidence_tier)

    # CrewAI
    crew_adapter = CrewAIAdapter(guard)
    result = crew_adapter.score_from_task_output(task_output, question="Summarise the report")

A2A trust handoff
-----------------
Both adapters optionally generate an ``A2ATrustObject`` for downstream agents:

    trust = adapter.to_trust_object(result, answer="Paris")
    # pass trust.to_dict() to the next agent's context
Purpose:    AutoGenAdapter converts GroupChat message history to AgentGuard steps; CrewAIAdapter converts TaskOutput to steps. Both enable transparent confidence scoring within multi-agent frameworks.
Exports:    AutoGenAdapter, CrewAIAdapter
Added:      v0.1.0
Status:     Stable
"""

from __future__ import annotations

import warnings
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from llm_guard.scoring.agent_guard import AgentGuard, ChainTrustResult
    from llm_guard.trust.trust_object import A2ATrustObject


# ── Step normalisation helpers ────────────────────────────────────────────────

def _autogen_messages_to_steps(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """
    Convert AutoGen GroupChat message history to normalized step dicts.

    Each AutoGen message is either:
    - A tool call / function call (mapped to action_type="Search" or "ToolCall")
    - A plain text agent reply (mapped to action_type="Think" with thought=content)
    - A final user-visible reply (last non-system message, mapped to "Finish")

    Returns a list of dicts compatible with ``normalize_steps()`` internal format.
    """
    steps: List[Dict[str, Any]] = []
    for i, msg in enumerate(messages):
        role    = msg.get("role", msg.get("sender", "assistant"))
        content = msg.get("content", "") or ""
        name    = msg.get("name", role)

        # Tool / function call
        tool_calls = msg.get("tool_calls") or msg.get("function_call")
        if tool_calls:
            if isinstance(tool_calls, dict):
                tool_calls = [tool_calls]
            for tc in (tool_calls if isinstance(tool_calls, list) else [tool_calls]):
                fn_name = (
                    tc.get("function", {}).get("name", "")
                    or tc.get("name", "tool_call")
                )
                fn_args = (
                    tc.get("function", {}).get("arguments", "")
                    or tc.get("arguments", str(tc))
                )
                steps.append({
                    "thought": f"[{name}] calling {fn_name}",
                    "action_type": "ToolCall",
                    "action_arg": f"{fn_name}({fn_args})",
                    "observation": "",
                })
        elif role in ("tool", "function"):
            # Tool result message — attach as observation to last step
            if steps:
                steps[-1]["observation"] = content[:1000]
        elif content:
            is_last = (i == len(messages) - 1)
            steps.append({
                "thought": f"[{name}] {content[:300]}",
                "action_type": "Finish" if is_last else "Think",
                "action_arg": content if is_last else "",
                "observation": "",
            })

    # Reclassify: if no ToolCall steps, treat Think steps as Search (for SC_OLD signal)
    has_tool = any(s["action_type"] == "ToolCall" for s in steps)
    if not has_tool:
        for s in steps:
            if s["action_type"] == "Think":
                s["action_type"] = "Search"
                s["action_arg"]  = s["thought"]

    return steps


def _crewai_task_to_steps(task_output: Any) -> List[Dict[str, Any]]:
    """
    Convert a CrewAI TaskOutput (or raw dict/str) to normalized step dicts.

    CrewAI TaskOutput attributes used:
    - ``raw`` — the raw LLM response (string)
    - ``pydantic`` — structured output (optional)
    - ``json_dict`` — JSON output dict (optional)
    - ``agent`` — agent name string
    - ``tasks_output`` — list of per-task outputs (for Crew.kickoff results)
    """
    steps: List[Dict[str, Any]] = []

    # Crew-level result (has .tasks_output list)
    if hasattr(task_output, "tasks_output"):
        for to in task_output.tasks_output:
            steps.extend(_crewai_task_to_steps(to))
        return steps

    # Single TaskOutput
    if hasattr(task_output, "raw"):
        raw    = str(task_output.raw or "")
        agent  = getattr(task_output, "agent", "agent") or "agent"
        output = (
            getattr(task_output, "pydantic", None)
            or getattr(task_output, "json_dict", None)
        )
        # Each task is one Search step + Finish
        steps.append({
            "thought": f"[{agent}] working on task",
            "action_type": "Search",
            "action_arg": raw[:500],
            "observation": str(output)[:500] if output else raw[:500],
        })
        steps.append({
            "thought": f"[{agent}] completed",
            "action_type": "Finish",
            "action_arg": raw,
            "observation": "",
        })
    elif isinstance(task_output, dict):
        content = str(task_output.get("output", task_output.get("result", str(task_output))))
        steps.append({
            "thought": "crew task result",
            "action_type": "Finish",
            "action_arg": content,
            "observation": "",
        })
    elif isinstance(task_output, str):
        steps.append({
            "thought": "crew task result",
            "action_type": "Finish",
            "action_arg": task_output,
            "observation": "",
        })

    return steps


# ── AutoGen Adapter ────────────────────────────────────────────────────────────

@dataclass
class OrchestratorScoringResult:
    """Scoring result from an orchestrator adapter."""

    chain_result: "ChainTrustResult"
    """Full ChainTrustResult from AgentGuard.score_chain()."""

    framework: str = ""
    """'autogen' or 'crewai'."""

    n_messages: int = 0
    """Number of messages / task outputs processed."""

    n_steps: int = 0
    """Number of normalized steps passed to score_chain()."""

    question: str = ""
    """The question scored."""

    final_answer: str = ""
    """The final answer scored."""

    @property
    def risk_score(self) -> float:
        return self.chain_result.risk_score

    @property
    def confidence_tier(self) -> str:
        return self.chain_result.confidence_tier

    @property
    def needs_alert(self) -> bool:
        return self.chain_result.needs_alert


class AutoGenAdapter:
    """
    Converts AutoGen GroupChat conversation history to AgentGuard confidence scores.

    Supports:
    - ``autogen`` (pyautogen) GroupChat message format
    - ``autogen_agentchat`` (new AutoGen 0.4+) message format
    - Any list of dicts with ``role`` / ``content`` / ``tool_calls`` keys

    Parameters
    ----------
    guard : AgentGuard
        Pre-configured AgentGuard instance.
    extract_answer_fn : callable, optional
        Function to extract the final answer from a message list.
        Default: last non-system message content.
    """

    def __init__(
        self,
        guard: "AgentGuard",
        extract_answer_fn: Optional[Callable[[List[Dict]], str]] = None,
    ) -> None:
        self._guard = guard
        self._extract_answer = extract_answer_fn or _default_final_answer

    def score_from_messages(
        self,
        messages: List[Dict[str, Any]],
        question: str,
        final_answer: Optional[str] = None,
    ) -> OrchestratorScoringResult:
        """
        Score a GroupChat conversation.

        Parameters
        ----------
        messages : list[dict]
            AutoGen GroupChat ``.messages`` list (list of message dicts).
        question : str
            The original user question / task.
        final_answer : str, optional
            If omitted, extracted from the last non-system message.

        Returns
        -------
        OrchestratorScoringResult
        """
        if not messages:
            warnings.warn(
                "AutoGenAdapter.score_from_messages(): empty message list.",
                UserWarning, stacklevel=2,
            )

        answer = final_answer or self._extract_answer(messages)
        steps  = _autogen_messages_to_steps(messages)

        if len(steps) < 2:
            warnings.warn(
                "AutoGenAdapter: fewer than 2 steps extracted from messages. "
                "SC_OLD accuracy is reduced for very short chains. "
                "Consider passing use_judge=True to your AgentGuard for reliability.",
                UserWarning, stacklevel=2,
            )

        result = self._guard.score_chain(question, steps, answer)
        return OrchestratorScoringResult(
            chain_result=result,
            framework="autogen",
            n_messages=len(messages),
            n_steps=len(steps),
            question=question,
            final_answer=answer,
        )

    def register_hook(self, groupchat: Any, question: str) -> None:
        """
        Register a reply callback on an AutoGen GroupChat.

        On each new message the guard scores the current chain and emits a
        UserWarning if ``needs_alert`` is True.

        Parameters
        ----------
        groupchat : autogen.GroupChat or agentchat.GroupChat
            The GroupChat instance to hook into.
        question : str
            The driving question for this conversation.
        """
        adapter = self

        def _on_message(message: Dict[str, Any], **_: Any) -> None:
            msgs = list(groupchat.messages)
            if not msgs:
                return
            result = adapter.score_from_messages(msgs, question=question)
            if result.needs_alert:
                warnings.warn(
                    f"AutoGenAdapter [alert]: risk={result.risk_score:.3f} "
                    f"tier={result.confidence_tier} after {result.n_messages} messages. "
                    "One or more agents may be producing unreliable output.",
                    UserWarning, stacklevel=1,
                )

        # Try AutoGen 0.2 API (register_reply)
        try:
            # GroupChat doesn't have register_reply directly; attach to manager
            if hasattr(groupchat, "register_reply"):
                groupchat.register_reply([object], _on_message)
                return
        except Exception:
            pass

        # Fallback: monkey-patch the append method on groupchat.messages
        if hasattr(groupchat, "messages") and isinstance(groupchat.messages, list):
            original_append = groupchat.messages.append

            def _patched_append(msg: Any) -> None:
                original_append(msg)
                _on_message(msg)

            groupchat.messages.append = _patched_append  # type: ignore[method-assign]
        else:
            warnings.warn(
                "AutoGenAdapter.register_hook(): could not attach to groupchat. "
                "Use score_from_messages() for post-hoc scoring instead.",
                UserWarning, stacklevel=2,
            )

    def to_trust_object(
        self,
        result: OrchestratorScoringResult,
        answer: Optional[str] = None,
    ) -> "A2ATrustObject":
        """
        Wrap a scoring result into an A2ATrustObject for downstream agent handoff.
        """
        from llm_guard.trust.trust_object import A2ATrustObject
        return A2ATrustObject(
            answer=answer or result.final_answer,
            risk_score=result.risk_score,
            confidence_tier=result.confidence_tier,
            failure_mode=result.chain_result.failure_mode or "UNKNOWN",
            step_count=result.n_steps,
            judge_label=result.chain_result.judge_label,
            downstream_hint="proceed_with_caution" if result.needs_alert else "proceed",
            should_rewrite=result.needs_alert,
            routing_mode="autogen_adapter",
        )


# ── CrewAI Adapter ─────────────────────────────────────────────────────────────

class CrewAIAdapter:
    """
    Converts CrewAI task execution output to AgentGuard confidence scores.

    Supports:
    - ``CrewOutput`` (result of ``Crew.kickoff()``)
    - ``TaskOutput`` (result of a single Task)
    - Raw dict / string outputs

    Parameters
    ----------
    guard : AgentGuard
        Pre-configured AgentGuard instance.
    """

    def __init__(self, guard: "AgentGuard") -> None:
        self._guard = guard

    def score_from_task_output(
        self,
        task_output: Any,
        question: str,
        final_answer: Optional[str] = None,
    ) -> OrchestratorScoringResult:
        """
        Score a CrewAI task or crew output.

        Parameters
        ----------
        task_output : CrewOutput | TaskOutput | dict | str
            The output from ``Crew.kickoff()`` or ``task.execute()``.
        question : str
            The driving question / goal.
        final_answer : str, optional
            Override the final answer. If omitted, extracted from task_output.

        Returns
        -------
        OrchestratorScoringResult
        """
        steps  = _crewai_task_to_steps(task_output)
        answer = final_answer or _extract_crewai_answer(task_output)

        if not steps:
            warnings.warn(
                "CrewAIAdapter: no steps could be extracted from task_output. "
                "Pass a CrewOutput, TaskOutput, or dict with 'output'/'result' key.",
                UserWarning, stacklevel=2,
            )
            steps = [{"thought": "", "action_type": "Finish", "action_arg": answer, "observation": ""}]

        result = self._guard.score_chain(question, steps, answer)
        n_raw = (
            len(task_output.tasks_output)
            if hasattr(task_output, "tasks_output")
            else 1
        )
        return OrchestratorScoringResult(
            chain_result=result,
            framework="crewai",
            n_messages=n_raw,
            n_steps=len(steps),
            question=question,
            final_answer=answer,
        )

    def wrap_task(self, task: Any, question: str) -> Any:
        """
        Return a thin wrapper around a CrewAI Task that auto-scores on completion.

        The wrapped task behaves identically to the original; after ``execute()``
        or ``async_execute()`` the result is passed to ``score_from_task_output()``
        and a UserWarning is emitted if the guard flags it.

        Parameters
        ----------
        task : crewai.Task
            The CrewAI Task to wrap.
        question : str
            The task's driving question.
        """
        adapter = self

        class _WrappedTask:
            def __init__(self, inner: Any) -> None:
                self._inner = inner
                # forward all attribute access
                for attr in dir(inner):
                    if not attr.startswith("_") and not hasattr(self, attr):
                        try:
                            setattr(self, attr, getattr(inner, attr))
                        except Exception:
                            pass

            def execute(self, *args: Any, **kwargs: Any) -> Any:
                output = self._inner.execute(*args, **kwargs)
                _score_and_warn(output)
                return output

            async def async_execute(self, *args: Any, **kwargs: Any) -> Any:
                output = await self._inner.async_execute(*args, **kwargs)
                _score_and_warn(output)
                return output

        def _score_and_warn(output: Any) -> None:
            try:
                r = adapter.score_from_task_output(output, question=question)
                if r.needs_alert:
                    warnings.warn(
                        f"CrewAIAdapter [alert]: task risk={r.risk_score:.3f} "
                        f"tier={r.confidence_tier}. Output may be unreliable.",
                        UserWarning, stacklevel=2,
                    )
            except Exception as exc:
                warnings.warn(
                    f"CrewAIAdapter.wrap_task: scoring failed ({exc}). "
                    "Continuing without guard score.",
                    UserWarning, stacklevel=2,
                )

        return _WrappedTask(task)

    def to_trust_object(
        self,
        result: OrchestratorScoringResult,
        answer: Optional[str] = None,
    ) -> "A2ATrustObject":
        """Wrap a scoring result into an A2ATrustObject for downstream agent handoff."""
        from llm_guard.trust.trust_object import A2ATrustObject
        return A2ATrustObject(
            answer=answer or result.final_answer,
            risk_score=result.risk_score,
            confidence_tier=result.confidence_tier,
            failure_mode=result.chain_result.failure_mode or "UNKNOWN",
            step_count=result.n_steps,
            judge_label=result.chain_result.judge_label,
            downstream_hint="proceed_with_caution" if result.needs_alert else "proceed",
            should_rewrite=result.needs_alert,
            routing_mode="crewai_adapter",
        )


# ── Helpers ────────────────────────────────────────────────────────────────────

def _default_final_answer(messages: List[Dict[str, Any]]) -> str:
    """Extract the last non-system message content as the final answer."""
    for msg in reversed(messages):
        role    = msg.get("role", msg.get("sender", ""))
        content = msg.get("content", "")
        if role not in ("system",) and content:
            return str(content)
    return ""


def _extract_crewai_answer(task_output: Any) -> str:
    """Extract the final answer string from a CrewAI output object."""
    # CrewOutput
    if hasattr(task_output, "raw"):
        return str(task_output.raw or "")
    # TaskOutput list (old API)
    if hasattr(task_output, "tasks_output") and task_output.tasks_output:
        return str(task_output.tasks_output[-1].raw or "")
    # Dict
    if isinstance(task_output, dict):
        return str(task_output.get("output", task_output.get("result", str(task_output))))
    return str(task_output)
