"""
Unit tests for ToolCallDegradationCorrector.
"""

from __future__ import annotations

from typing import Any

import pytest
from langchain_core.messages import AIMessage

from .tool_call_degradation_corrector import ToolCallDegradationCorrector


class _FakeLLM:
    """Simple fake LLM that returns a predefined response."""

    def __init__(self, response: Any) -> None:
        self._response = response

    def invoke(self, input: Any, **kwargs: Any) -> Any:  # noqa: ARG002
        return self._response


def test_corrector_passes_through_when_tool_calls_present() -> None:
    """If AIMessage already has tool_calls, corrector should be a no-op."""
    original = AIMessage(
        content="use tool",
        tool_calls=[
            {
                "name": "read_file",
                "args": {"path": "/foo"},
                "id": "tc-1",
                "type": "tool_call",
            }
        ],
    )
    fake_llm = _FakeLLM(original)
    corrector = ToolCallDegradationCorrector(fake_llm, allowed_tool_names={"read_file"})

    result = corrector.invoke({"messages": []})

    # Should return the same object without modification
    assert result is original
    assert result.tool_calls == original.tool_calls


def test_corrector_restores_degraded_tool_call_with_allowed_tool() -> None:
    """Degraded <tool_call> JSON should be parsed into structured tool_calls."""
    degraded_content = (
        '<tool_call>{"name": "read_file", "args": {"path": "/foo.txt"}}</tool_call>'
    )
    original = AIMessage(content=degraded_content)
    fake_llm = _FakeLLM(original)
    corrector = ToolCallDegradationCorrector(fake_llm, allowed_tool_names={"read_file"})

    result: AIMessage = corrector.invoke({"messages": []})

    # tool_calls should be restored
    assert result.tool_calls
    assert len(result.tool_calls) == 1
    tc = result.tool_calls[0]
    assert tc["name"] == "read_file"
    assert tc["args"] == {"path": "/foo.txt"}
    # content should be cleaned
    assert "<tool_call>" not in result.content
    # finish_reason should be normalized to "tool-calls"
    rm = getattr(result, "response_metadata", {}) or {}
    assert rm.get("finish_reason") == "tool-calls"


def test_corrector_ignores_unregistered_tool_name() -> None:
    """Degraded tool call with name not in whitelist should be ignored."""
    degraded_content = '<tool_call>{"name": "unknown_tool", "args": {"x": 1}}</tool_call>'
    original = AIMessage(content=degraded_content)
    fake_llm = _FakeLLM(original)
    corrector = ToolCallDegradationCorrector(fake_llm, allowed_tool_names={"read_file"})

    result: AIMessage = corrector.invoke({"messages": []})

    # No tool_calls should be added
    assert not getattr(result, "tool_calls", None)
    # Content should remain unchanged (since nothing was parsed)
    assert result.content == degraded_content


@pytest.mark.asyncio
async def test_ainvoke_behaves_the_same_as_invoke() -> None:
    """ainvoke should apply the same correction logic as invoke."""
    degraded_content = (
        '<tool_call>{"name": "read_file", "args": {"path": "/async.txt"}}</tool_call>'
    )
    original = AIMessage(content=degraded_content)

    class _AsyncFakeLLM:
        def __init__(self, response: Any) -> None:
            self._response = response

        async def ainvoke(self, input: Any, **kwargs: Any) -> Any:  # noqa: ARG002
            return self._response

    fake_llm = _AsyncFakeLLM(original)
    corrector = ToolCallDegradationCorrector(fake_llm, allowed_tool_names={"read_file"})

    result: AIMessage = await corrector.ainvoke({"messages": []})

    assert result.tool_calls
    assert result.tool_calls[0]["name"] == "read_file"
    rm = getattr(result, "response_metadata", {}) or {}
    assert rm.get("finish_reason") == "tool-calls"

