"""
ToolCallExtractor — StepExtractor for function-calling / tool-use agents.
==========================================================================

Supports three input formats (auto-detected in priority order):
  1. OpenAI tool_calls format: {"tool_calls": [...], "content": ..., "role": ...}
  2. LangChain AgentAction objects: hasattr(step, "tool")
  3. Generic dict: {"tool": str, "args": dict|str, "output": str, "error": str|None}

All three normalize to the generic dict before feature extraction.

10 features (7 step-level + 3 chain-level, computed in aggregate()):
  has_error, output_size_norm, arg_complexity_norm, latency_norm,
  is_retry (computed in extract() when aggregate() injects _prev_tool/_prev_args context),
  tool_idx_norm (computed in extract() when aggregate() injects _step_idx/_chain_len context),
  error_follows_error (computed in extract() from injected _prev_has_error context),
  tool_diversity_ratio (chain-level: unique tools / total steps),
  error_recovery_rate (chain-level: same-tool recoveries / total errors),
  total_error_rate (chain-level: error steps / total steps)
Purpose:    Supports OpenAI tool_calls, LangChain AgentAction, and generic dict formats. Extracts 7 step-level and 3 chain-level features including error rate, tool diversity, and retry detection.
Exports:    ToolCallExtractor
Added:      v0.33.0
Status:     Validated
"""
from __future__ import annotations

import json
import re
import warnings
from typing import Any, Dict, List, Optional

import numpy as np

from llm_guard.features.step_extractor import StepExtractor

# Tool names validated in exp_tool_call_validate (HP chains, AUROC 0.6526).
# Scores on tools NOT in this set are extrapolations — no empirical AUROC backing.
_VALIDATED_TOOL_RE = re.compile(
    r"(tavily_search|duckduckgo_search|wikipedia|retriev|search|lookup)",
    re.IGNORECASE,
)


class ToolCallExtractor(StepExtractor):
    """StepExtractor for tool-calling agents. Auto-detects input format."""

    _FEATURE_NAMES = [
        "has_error",
        "output_size_norm",
        "arg_complexity_norm",
        "latency_norm",
        "is_retry",
        "tool_idx_norm",
        "error_follows_error",
        "tool_diversity_ratio",
        "error_recovery_rate",
        "total_error_rate",
    ]

    @property
    def feature_names(self) -> List[str]:
        return list(self._FEATURE_NAMES)

    # ------------------------------------------------------------------
    # Format normalization
    # ------------------------------------------------------------------

    def _normalize(self, step: Any) -> Dict:
        """Normalize any supported format to generic dict."""
        # OpenAI tool_calls format
        if isinstance(step, dict) and "tool_calls" in step:
            tc = step["tool_calls"]
            if tc and isinstance(tc, list):
                fn = tc[0].get("function", {})
                args_raw = fn.get("arguments", "{}")
                try:
                    args = json.loads(args_raw) if isinstance(args_raw, str) else args_raw
                except (json.JSONDecodeError, TypeError):
                    args = {"_raw": str(args_raw)}
                return {
                    "tool": fn.get("name", "unknown"),
                    "args": args,
                    "output": step.get("content") or "",
                    "error": None,
                    "latency_ms": step.get("latency_ms"),
                }

        # LangChain AgentAction-like objects
        if hasattr(step, "tool"):
            args_raw = getattr(step, "tool_input", {})
            if isinstance(args_raw, str):
                try:
                    args_raw = json.loads(args_raw)
                except (json.JSONDecodeError, TypeError):
                    args_raw = {"_raw": args_raw}
            rv = getattr(step, "return_values", {})
            output = rv.get("output", "") if isinstance(rv, dict) else str(rv)
            return {
                "tool": step.tool,
                "args": args_raw,
                "output": output,
                "error": None,
                "latency_ms": getattr(step, "latency_ms", None),
            }

        # Generic dict (fallback)
        if isinstance(step, dict):
            args = step.get("args", {})
            if isinstance(args, str):
                try:
                    args = json.loads(args)
                except (json.JSONDecodeError, TypeError):
                    args = {"_raw": args}
            return {
                "tool": step.get("tool", "unknown"),
                "args": args,
                "output": str(step.get("output", "") or ""),
                "error": step.get("error") or None,
                "latency_ms": step.get("latency_ms"),
            }

        return {"tool": "unknown", "args": {}, "output": "", "error": None, "latency_ms": None}

    # ------------------------------------------------------------------
    # Per-step feature extraction
    # ------------------------------------------------------------------

    def extract(self, step: Any) -> Dict[str, float]:
        """Extract per-step features. Placeholders filled by aggregate()."""
        norm = self._normalize(step) if not isinstance(step, dict) or "_normalized" not in step else step

        has_error = float(bool(norm.get("error")))
        output_size_norm = min(len(str(norm.get("output", "")).split()) / 100.0, 1.0)

        try:
            arg_str = json.dumps(norm.get("args", {}))
        except (TypeError, ValueError):
            arg_str = str(norm.get("args", ""))
        arg_complexity_norm = min(len(arg_str) / 200.0, 1.0)

        latency_ms = norm.get("latency_ms")
        latency_norm = min(float(latency_ms) / 5000.0, 1.0) if latency_ms is not None else 0.0

        # Injected context keys (set by aggregate() pre-processing)
        prev_has_error = norm.get("_prev_has_error", 0.0)
        prev_tool = norm.get("_prev_tool")
        prev_args = norm.get("_prev_args")
        step_idx = norm.get("_step_idx", 0)
        chain_len = norm.get("_chain_len", 1)

        # is_retry: same tool + same args as previous step
        is_retry = 0.0
        if prev_tool is not None and prev_tool == norm.get("tool"):
            if prev_args is not None:
                try:
                    is_retry = float(json.dumps(prev_args, sort_keys=True) ==
                                     json.dumps(norm.get("args", {}), sort_keys=True))
                except (TypeError, ValueError):
                    is_retry = float(str(prev_args) == str(norm.get("args", {})))

        tool_idx_norm = step_idx / max(chain_len - 1, 1) if chain_len > 1 else 0.0
        error_follows_error = float(bool(prev_has_error) and bool(has_error))

        return {
            "has_error":            has_error,
            "output_size_norm":     output_size_norm,
            "arg_complexity_norm":  arg_complexity_norm,
            "latency_norm":         latency_norm,
            "is_retry":             is_retry,
            "tool_idx_norm":        float(tool_idx_norm),
            "error_follows_error":  error_follows_error,
            "tool_diversity_ratio": 0.0,   # filled in aggregate()
            "error_recovery_rate":  0.0,   # filled in aggregate()
            "total_error_rate":     0.0,   # filled in aggregate()
        }

    # ------------------------------------------------------------------
    # Chain-level aggregation
    # ------------------------------------------------------------------

    def aggregate(self, steps: List[Any], final_answer: str = "") -> np.ndarray:
        if not steps:
            return np.zeros(len(self.feature_names), dtype=np.float32)

        n = len(steps)

        # Warn if any tool names are outside the validated set.
        # ToolCallExtractor was validated on Search-based HotpotQA chains only
        # (exp_tool_call_validate, AUROC 0.6526). Non-search tools (code execution,
        # API calls, file I/O, etc.) have no empirical AUROC backing.
        unmapped = set()
        for s in steps:
            norm = self._normalize(s)
            tool = norm.get("tool", "unknown")
            if tool and not _VALIDATED_TOOL_RE.search(tool):
                unmapped.add(tool)
        if unmapped:
            warnings.warn(
                f"ToolCallExtractor.aggregate(): tool(s) {sorted(unmapped)!r} are not in the "
                f"validated set (tavily_search, duckduckgo_search, wikipedia, retriev*, search). "
                f"Scores on non-search tools have no empirical AUROC backing "
                f"(validated only on HotpotQA Search-based chains, AUROC 0.6526). "
                f"Do not publish reliability estimates for these tool types.",
                UserWarning,
                stacklevel=2,
            )

        # Pre-process: inject context keys into each step dict copy
        enriched = []
        for i, s in enumerate(steps):
            norm = self._normalize(s)
            norm["_normalized"] = True
            norm["_step_idx"] = i
            norm["_chain_len"] = n
            if i > 0:
                prev = self._normalize(steps[i - 1])
                norm["_prev_tool"] = prev.get("tool")
                norm["_prev_args"] = prev.get("args")
                norm["_prev_has_error"] = float(bool(prev.get("error")))
            else:
                norm["_prev_tool"] = None
                norm["_prev_args"] = None
                norm["_prev_has_error"] = 0.0
            enriched.append(norm)

        rows = [self.extract(s) for s in enriched]
        names = self.feature_names
        matrix = np.array([[r.get(name, 0.0) for name in names] for r in rows], dtype=np.float32)
        agg = matrix.mean(axis=0).copy()

        idx = {name: i for i, name in enumerate(names)}

        # Chain-level: tool_diversity_ratio
        tools = [e.get("tool", "unknown") for e in enriched]
        agg[idx["tool_diversity_ratio"]] = len(set(tools)) / max(n, 1)

        # Chain-level: total_error_rate
        errors = [float(bool(e.get("error"))) for e in enriched]
        n_errors = sum(errors)
        agg[idx["total_error_rate"]] = n_errors / max(n, 1)

        # Chain-level: error_recovery_rate
        # A recovery = error at step i, then success (no error) at step i+1 on the SAME tool.
        # Cross-tool recovery (error on tool A, success on tool B) is not counted.
        n_recoveries = 0
        for i in range(1, n):
            prev_norm = enriched[i - 1]
            curr_norm = enriched[i]
            if (prev_norm.get("error") and not curr_norm.get("error")
                    and prev_norm.get("tool") == curr_norm.get("tool")):
                n_recoveries += 1
        agg[idx["error_recovery_rate"]] = n_recoveries / max(n_errors, 1)

        return agg.astype(np.float32)
