"""Tests for ToolCallExtractor — format normalization and feature extraction."""
import sys
import numpy as np
import pytest
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from llm_guard.features.tool_call_extractor import ToolCallExtractor


# ── Fixtures ──────────────────────────────────────────────────────────────────

OPENAI_STEPS = [
    {
        "tool_calls": [{"function": {"name": "execute_sql", "arguments": '{"query": "SELECT * FROM users"}'}, "id": "call_1"}],
        "content": None, "role": "assistant"
    },
    {
        "tool_calls": [{"function": {"name": "execute_sql", "arguments": '{"query": "SELECT * FROM orders"}'}, "id": "call_2"}],
        "content": None, "role": "assistant"
    },
]

GENERIC_STEPS = [
    {"tool": "read_file", "args": {"path": "data.csv"}, "output": "col1,col2\n1,2", "error": None},
    {"tool": "execute_sql", "args": {"query": "SELECT COUNT(*) FROM t"}, "output": "42", "error": None},
    {"tool": "execute_sql", "args": {"query": "SELECT COUNT(*) FROM t"}, "output": "", "error": "Table not found"},
]

GENERIC_STEPS_WITH_RETRY = [
    {"tool": "call_api", "args": {"url": "http://x.com"}, "output": "", "error": "timeout"},
    {"tool": "call_api", "args": {"url": "http://x.com"}, "output": "200 OK", "error": None},
]


# ── Format detection and normalization ────────────────────────────────────────

def test_openai_format_detected():
    ext = ToolCallExtractor()
    norm = ext._normalize(OPENAI_STEPS[0])
    assert norm["tool"] == "execute_sql"
    assert "query" in str(norm["args"])
    assert norm["error"] is None


def test_generic_format_detected():
    ext = ToolCallExtractor()
    norm = ext._normalize(GENERIC_STEPS[0])
    assert norm["tool"] == "read_file"
    assert norm["error"] is None


def test_langchain_format_detected():
    """LangChain AgentAction-like objects with .tool attribute."""
    class FakeAgentAction:
        def __init__(self, tool, tool_input, output):
            self.tool = tool
            self.tool_input = tool_input
            self.return_values = {"output": output}
    step = FakeAgentAction("execute_sql", '{"query": "SELECT 1"}', "[(1,)]")
    ext = ToolCallExtractor()
    norm = ext._normalize(step)
    assert norm["tool"] == "execute_sql"
    assert norm["error"] is None


# ── feature_names contract ────────────────────────────────────────────────────

def test_feature_names_count():
    ext = ToolCallExtractor()
    assert len(ext.feature_names) == 10


def test_extract_returns_all_feature_names():
    ext = ToolCallExtractor()
    step = {"tool": "read_file", "args": {}, "output": "data", "error": None}
    feats = ext.extract(step)
    for name in ext.feature_names:
        assert name in feats, f"Missing feature: {name}"


def test_extract_values_in_range():
    ext = ToolCallExtractor()
    step = {"tool": "read_file", "args": {"path": "f.csv"}, "output": "a,b,c", "error": None}
    feats = ext.extract(step)
    for name, val in feats.items():
        assert 0.0 <= val <= 1.0, f"{name}={val} out of [0,1]"


def test_error_step_has_error_1():
    ext = ToolCallExtractor()
    step = {"tool": "execute_sql", "args": {}, "output": "", "error": "Table not found"}
    feats = ext.extract(step)
    assert feats["has_error"] == 1.0


def test_no_error_step_has_error_0():
    ext = ToolCallExtractor()
    step = {"tool": "execute_sql", "args": {}, "output": "42", "error": None}
    feats = ext.extract(step)
    assert feats["has_error"] == 0.0


# ── aggregate() chain-level features ─────────────────────────────────────────

def test_aggregate_shape():
    ext = ToolCallExtractor()
    vec = ext.aggregate(GENERIC_STEPS, final_answer="42")
    assert vec.shape == (len(ext.feature_names),)


def test_aggregate_values_in_range():
    ext = ToolCallExtractor()
    vec = ext.aggregate(GENERIC_STEPS, final_answer="42")
    assert np.all(vec >= 0.0) and np.all(vec <= 1.0)


def test_total_error_rate_computed():
    ext = ToolCallExtractor()
    # 1 error out of 3 steps
    vec = ext.aggregate(GENERIC_STEPS, final_answer="42")
    names = ext.feature_names
    err_idx = names.index("total_error_rate")
    assert abs(vec[err_idx] - 1/3) < 0.01


def test_tool_diversity_ratio_all_unique():
    ext = ToolCallExtractor()
    steps = [
        {"tool": "read_file", "args": {}, "output": "a", "error": None},
        {"tool": "execute_sql", "args": {}, "output": "b", "error": None},
    ]
    vec = ext.aggregate(steps, final_answer="result")
    names = ext.feature_names
    div_idx = names.index("tool_diversity_ratio")
    assert vec[div_idx] == 1.0  # 2 unique / 2 steps


def test_is_retry_detected():
    ext = ToolCallExtractor()
    # Two identical consecutive calls → retry
    vec = ext.aggregate(GENERIC_STEPS_WITH_RETRY, final_answer="200 OK")
    names = ext.feature_names
    retry_idx = names.index("is_retry")
    assert vec[retry_idx] > 0.0


def test_empty_steps_returns_zeros():
    ext = ToolCallExtractor()
    vec = ext.aggregate([], final_answer="")
    assert vec.shape == (len(ToolCallExtractor().feature_names),)
    assert np.all(vec == 0.0)
