"""
FunctionCallGuard — AgentGuard wrapper for native tool-call message lists.
==========================================================================

Adapts OpenAI and Claude native tool_calls / tool_use message lists into
llm-guard step dicts, then scores them with AgentGuard behavioral SC_OLD.

ToolCallExtractor remains separately available for supervised 10-feature scoring.

Quick start::

    from llm_guard import FunctionCallGuard

    guard = FunctionCallGuard()
    result = guard.score(messages, question="Who was Einstein?")
    print(result.risk_score, result.confidence_tier)
Purpose:    Adapts OpenAI and Claude native tool_calls/tool_use message lists into llm-guard step dicts and scores them with AgentGuard behavioral SC_OLD. ToolCallExtractor remains separately available for supervised scoring.
Exports:    FunctionCallGuard
Added:      v0.33.0
Status:     Stable
"""

from __future__ import annotations

import json
import os
from typing import List, Optional


class FunctionCallGuard:
    """
    Score OpenAI or Claude native tool-call message lists with AgentGuard.

    Parameters
    ----------
    api_key : str, optional
        Anthropic API key. Defaults to ANTHROPIC_API_KEY env var.
    use_judge : bool, default False
        If True, adds Sonnet judge signal (same as AgentGuard(use_judge=True)).
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        use_judge: bool = False,
    ) -> None:
        self._api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        self._use_judge = use_judge
        self._guard = None
        self._tce = None  # ToolCallExtractor (fitted, optional supervised mode)

    def _get_guard(self):
        if self._guard is None:
            from llm_guard.scoring.agent_guard import AgentGuard
            self._guard = AgentGuard(
                api_key=self._api_key if self._api_key else None,
                use_judge=self._use_judge,
            )
        return self._guard

    @staticmethod
    def _parse_tool_calls(messages: List[dict]) -> List[dict]:
        """
        Convert OpenAI or Claude message list → llm-guard step dicts.

        OpenAI assistant: {"role":"assistant","tool_calls":[{"function":{"name":..,"arguments":..}}]}
        OpenAI tool result: {"role":"tool","content":...}
        Claude assistant: {"role":"assistant","content":[{"type":"tool_use","name":..,"input":..}]}
        Claude tool result: {"role":"user","content":[{"type":"tool_result","content":...}]}

        Returns List[{"thought", "action_type", "action_arg", "observation"}]
        """
        steps = []
        pending: Optional[dict] = None  # step dict awaiting its observation

        for msg in messages:
            role = msg.get("role", "")
            content = msg.get("content", "")

            # ── OpenAI assistant turn with tool_calls ────────────────────────
            if role == "assistant" and "tool_calls" in msg and msg["tool_calls"]:
                tc = msg["tool_calls"][0]  # take first tool call
                fn = tc.get("function", {})
                name = fn.get("name", "tool")
                args_raw = fn.get("arguments", "{}")
                try:
                    args_obj = json.loads(args_raw)
                    action_arg = str(args_obj.get("query", args_obj.get("input",
                                     args_obj.get("answer", args_raw))))
                except (json.JSONDecodeError, TypeError):
                    action_arg = str(args_raw)
                # thought from text content before tool_calls
                thought = ""
                if isinstance(content, str) and content:
                    thought = content.strip()
                pending = {
                    "thought": thought,
                    "action_type": name,
                    "action_arg": action_arg,
                    "observation": "",
                }

            # ── OpenAI tool result ───────────────────────────────────────────
            elif role == "tool":
                obs = content if isinstance(content, str) else str(content)
                if pending is not None:
                    pending["observation"] = obs[:500]
                    steps.append(pending)
                    pending = None

            # ── Claude assistant turn with tool_use blocks ───────────────────
            elif role == "assistant" and isinstance(content, list):
                thought = ""
                tool_block = None
                for block in content:
                    if isinstance(block, dict):
                        if block.get("type") == "text":
                            thought = block.get("text", "").strip()
                        elif block.get("type") == "tool_use":
                            tool_block = block

                if tool_block is not None:
                    inp = tool_block.get("input", {})
                    if isinstance(inp, dict):
                        action_arg = str(inp.get("query", inp.get("input", str(inp))))
                    else:
                        action_arg = str(inp)
                    pending = {
                        "thought": thought,
                        "action_type": tool_block.get("name", "tool"),
                        "action_arg": action_arg,
                        "observation": "",
                    }

            # ── Claude tool_result block (comes as user role) ────────────────
            elif role == "user" and isinstance(content, list):
                for block in content:
                    if isinstance(block, dict) and block.get("type") == "tool_result":
                        obs_raw = block.get("content", "")
                        obs = obs_raw if isinstance(obs_raw, str) else str(obs_raw)
                        if pending is not None:
                            pending["observation"] = obs[:500]
                            steps.append(pending)
                            pending = None
                        break

        # Flush any pending step without observation
        if pending is not None:
            steps.append(pending)

        return steps

    def score(self, messages: List[dict], question: str) -> object:
        """
        Parse native tool-call messages and score with AgentGuard.score_chain().

        Returns ChainTrustResult (same as AgentGuard.score_chain()).
        """
        guard = self._get_guard()
        steps = self._parse_tool_calls(messages)

        # Extract final answer from last assistant message or last step arg
        final_answer = ""
        for msg in reversed(messages):
            if msg.get("role") == "assistant":
                c = msg.get("content", "")
                if isinstance(c, str) and c and "tool_calls" not in msg:
                    final_answer = c.strip()
                    break
        if not final_answer and steps:
            final_answer = steps[-1].get("action_arg", "")

        # Add Finish step if missing
        if steps and steps[-1].get("action_type", "").lower() not in ("finish",):
            steps.append({
                "thought": "",
                "action_type": "Finish",
                "action_arg": final_answer,
                "observation": "",
            })

        if not steps:
            # No tool calls parsed — return a neutral result
            from llm_guard.scoring.agent_guard import ChainTrustResult
            return ChainTrustResult(
                risk_score=0.5,
                confidence_tier="MEDIUM",
                needs_alert=False,
                failure_mode=None,
                judge_label=None,
                step_count=0,
                behavioral_score=0.5,
            )

        result = guard.score_chain(question, steps, final_answer)

        # Blend supervised tce_risk when fit() has been called
        if self._tce is not None:
            tce_obj, pipe = self._tce
            search_steps = [s for s in steps if s.get("action_type", "").lower() != "finish"]
            if search_steps:
                try:
                    from dataclasses import replace as _dc_replace
                    feats = tce_obj.aggregate(search_steps)
                    tce_risk = float(pipe.predict_proba([feats.tolist()])[0][1])
                    blended = float(max(0.0, min(1.0, 0.5 * result.risk_score + 0.5 * tce_risk)))
                    result.behavioral_components["tce_risk"] = tce_risk
                    tier = "HIGH" if blended < 0.50 else ("MEDIUM" if blended < 0.70 else "LOW")
                    result = _dc_replace(
                        result,
                        risk_score=blended,
                        confidence_tier=tier,
                        behavioral_components=result.behavioral_components,
                    )
                except Exception:
                    pass

        return result

    def fit(self, labeled_chains: List[dict]) -> "FunctionCallGuard":
        """
        Optional supervised mode: train ToolCallExtractor LogReg.

        Parameters
        ----------
        labeled_chains : list of dicts
            Each: {"messages": [...], "label": int (1=wrong, 0=correct)}.
            Requires >= 20 chains.
        """
        if len(labeled_chains) < 20:
            raise ValueError(
                f"fit() requires >= 20 labeled chains, got {len(labeled_chains)}"
            )

        from llm_guard.features.tool_call_extractor import ToolCallExtractor
        from sklearn.linear_model import LogisticRegression
        from sklearn.pipeline import Pipeline
        from sklearn.preprocessing import StandardScaler
        import numpy as np

        tce = ToolCallExtractor()
        X, y = [], []
        for chain in labeled_chains:
            steps = self._parse_tool_calls(chain.get("messages", []))
            if not steps:
                continue
            try:
                feats = tce.aggregate(steps)
                X.append(feats.tolist())
                y.append(int(chain.get("label", 0)))
            except Exception:
                continue

        if len(X) < 20 or len(set(y)) < 2:
            raise ValueError("Not enough valid chains or only one class present after parsing.")

        pipe = Pipeline([("scaler", StandardScaler()),
                         ("lr", LogisticRegression(max_iter=500))])
        pipe.fit(np.array(X), np.array(y))
        self._tce = (tce, pipe)
        return self
