"""Tests for EscalateAction guardrail failure behavior."""

import json
from unittest.mock import MagicMock, patch

import pytest
from langchain_core.messages import AIMessage, HumanMessage
from langgraph.types import Command
from uipath.platform.guardrails import GuardrailScope
from uipath.runtime.errors import UiPathErrorCode

from uipath_langchain.agent.exceptions import AgentTerminationException
from uipath_langchain.agent.guardrails.actions.escalate_action import EscalateAction
from uipath_langchain.agent.guardrails.types import (
    AgentGuardrailsGraphState,
    ExecutionStage,
)


class TestEscalateAction:
    @pytest.mark.asyncio
    async def test_node_name_pre_llm(self):
        """PreExecution + LLM: name is sanitized correctly."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "My Guardrail v1"
        guardrail.description = "Test description"

        node_name, _ = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.PRE_EXECUTION,
        )

        assert node_name == "my_guardrail_v1_hitl_pre_execution_llm"

    @pytest.mark.asyncio
    async def test_node_name_post_agent(self):
        """PostExecution + AGENT: name is sanitized correctly."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Special-Guardrail@2024"
        guardrail.description = "Test description"

        node_name, _ = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.AGENT,
            execution_stage=ExecutionStage.POST_EXECUTION,
        )

        assert node_name == "special_guardrail_2024_hitl_post_execution_agent"

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_node_interrupts_with_correct_data_pre_llm(self, mock_interrupt):
        """PreExecution + LLM: interrupt is called with correct escalation data."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        # Mock interrupt to return approved escalation
        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {}
        mock_interrupt.return_value = mock_escalation_result

        node_name, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.PRE_EXECUTION,
        )

        state = AgentGuardrailsGraphState(
            messages=[HumanMessage(content="Test message")],
            guardrail_validation_result="Validation failed",
        )

        await node(state)

        # Verify interrupt was called
        assert mock_interrupt.called
        call_args = mock_interrupt.call_args[0][0]

        # Verify escalation data structure
        assert call_args.app_name == "TestApp"
        assert call_args.app_folder_path == "TestFolder"
        assert call_args.title == "TestApp"
        assert call_args.assignee == "test@example.com"
        assert call_args.data["GuardrailName"] == "Test Guardrail"
        assert call_args.data["GuardrailDescription"] == "Test description"
        assert call_args.data["ExecutionStage"] == "PreExecution"
        assert call_args.data["GuardrailResult"] == "Validation failed"
        assert call_args.data["Inputs"] == '"Test message"'

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_node_approval_returns_command(self, mock_interrupt):
        """When escalation is approved, returns Command from _process_escalation_response."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        # Mock interrupt to return approved escalation
        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {}
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.PRE_EXECUTION,
        )

        state = AgentGuardrailsGraphState(
            messages=[HumanMessage(content="Test message")],
        )

        result = await node(state)

        # Should return Command or empty dict (from _process_escalation_response)
        assert isinstance(result, (Command, dict))

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_node_rejection_raises_exception(self, mock_interrupt):
        """When escalation is rejected, raises AgentTerminationException."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        # Mock interrupt to return rejected escalation
        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Reject"
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.PRE_EXECUTION,
        )

        state = AgentGuardrailsGraphState(
            messages=[HumanMessage(content="Test message")],
        )

        with pytest.raises(AgentTerminationException) as excinfo:
            await node(state)

        assert excinfo.value.error_info.title == "Escalation rejected"
        assert (
            excinfo.value.error_info.detail
            == "Action was rejected after reviewing the task created by guardrail [Test Guardrail]. Please contact your administrator."
        )
        assert (
            excinfo.value.error_info.code
            == f"Python.{UiPathErrorCode.EXECUTION_ERROR.value}"
        )

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_node_post_execution_tool_field(self, mock_interrupt):
        """PostExecution: uses ToolOutputs field instead of ToolInputs."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        # Mock interrupt to return approved escalation
        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {}
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.POST_EXECUTION,
        )

        state = AgentGuardrailsGraphState(
            messages=[AIMessage(content="Test response")],
        )

        await node(state)

        # Verify ToolOutputs is used for PostExecution
        call_args = mock_interrupt.call_args[0][0]
        assert call_args.data["Outputs"] == '["Test response"]'
        assert "ToolInputs" not in call_args.data

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_post_execution_ai_message_with_tool_calls_extraction(
        self, mock_interrupt
    ):
        """PostExecution with AIMessage and tool calls: extracts tool calls and content."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {}
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.POST_EXECUTION,
        )

        ai_message = AIMessage(
            content="AI response",
            tool_calls=[
                {
                    "name": "test_tool",
                    "args": {"content": {"input": "test"}},
                    "id": "call_1",
                }
            ],
        )
        state = AgentGuardrailsGraphState(messages=[ai_message])

        await node(state)

        # Verify interrupt was called with tool calls and content in ToolOutputs
        call_args = mock_interrupt.call_args[0][0]
        tool_outputs = call_args.data["Outputs"]
        parsed = json.loads(tool_outputs)
        assert len(parsed) == 2  # Tool call content + message content
        assert parsed[1] == "AI response"

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_pre_execution_with_reviewed_inputs(self, mock_interrupt):
        """PreExecution: updates message content with ReviewedInputs."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        reviewed_content = {"updated": "content"}
        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {"ReviewedInputs": json.dumps(reviewed_content)}
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.PRE_EXECUTION,
        )

        state = AgentGuardrailsGraphState(
            messages=[HumanMessage(content="Original content")],
        )

        result = await node(state)

        # Verify message content was updated
        assert isinstance(result, Command)
        assert result.update is not None
        assert result.update["messages"][0].content == reviewed_content

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_post_execution_human_message_with_reviewed_outputs(
        self, mock_interrupt
    ):
        """PostExecution with HumanMessage: updates content with ReviewedOutputs."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        reviewed_content = ["Updated content"]
        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {"ReviewedOutputs": json.dumps(reviewed_content)}
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.POST_EXECUTION,
        )

        state = AgentGuardrailsGraphState(
            messages=[HumanMessage(content="Original content")],
        )

        result = await node(state)

        # Verify HumanMessage content was updated (ignores tool calls)
        assert isinstance(result, Command)
        assert result.update is not None
        assert result.update["messages"][0].content == "Updated content"

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_post_execution_ai_message_with_reviewed_outputs_and_tool_calls(
        self, mock_interrupt
    ):
        """PostExecution with AIMessage: updates tool calls and content with ReviewedOutputs."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        reviewed_tool_content = {"updated": "tool_content"}
        reviewed_message_content = "Updated message"
        reviewed_outputs = [
            json.dumps(reviewed_tool_content),
            reviewed_message_content,
        ]
        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {"ReviewedOutputs": json.dumps(reviewed_outputs)}
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.POST_EXECUTION,
        )

        ai_message = AIMessage(
            content="Original content",
            tool_calls=[
                {
                    "name": "test_tool",
                    "args": {"content": {"original": "tool_input"}},
                    "id": "call_1",
                }
            ],
        )
        state = AgentGuardrailsGraphState(messages=[ai_message])

        result = await node(state)

        # Verify tool calls and message content were updated
        assert isinstance(result, Command)
        assert result.update is not None
        updated_message = result.update["messages"][0]
        assert updated_message.tool_calls[0]["args"]["content"] == reviewed_tool_content
        assert updated_message.content == reviewed_message_content

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_non_llm_scope_returns_empty_dict(self, mock_interrupt):
        """Non-LLM scope: returns empty dict without processing messages."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {"ReviewedInputs": "test"}
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.AGENT,
            execution_stage=ExecutionStage.PRE_EXECUTION,
        )

        state = AgentGuardrailsGraphState(
            messages=[HumanMessage(content="Test message")],
        )

        result = await node(state)

        # Non-LLM scope should return empty dict
        assert result == {}

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_missing_reviewed_field_returns_empty_dict(self, mock_interrupt):
        """Missing reviewed field: returns empty dict when reviewed field not in result."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {}  # No ReviewedInputs field
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.PRE_EXECUTION,
        )

        state = AgentGuardrailsGraphState(
            messages=[HumanMessage(content="Test message")],
        )

        result = await node(state)

        # Missing reviewed field should return empty dict
        assert result == {}

    @pytest.mark.asyncio
    @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
    async def test_json_parsing_error_raises_exception(self, mock_interrupt):
        """JSON parsing error: raises AgentTerminationException with execution error."""
        action = EscalateAction(
            app_name="TestApp",
            app_folder_path="TestFolder",
            version=1,
            assignee="test@example.com",
        )
        guardrail = MagicMock()
        guardrail.name = "Test Guardrail"
        guardrail.description = "Test description"

        mock_escalation_result = MagicMock()
        mock_escalation_result.action = "Approve"
        mock_escalation_result.data = {
            "ReviewedInputs": "invalid json {"
        }  # Invalid JSON
        mock_interrupt.return_value = mock_escalation_result

        _, node = action.action_node(
            guardrail=guardrail,
            scope=GuardrailScope.LLM,
            execution_stage=ExecutionStage.PRE_EXECUTION,
        )

        state = AgentGuardrailsGraphState(
            messages=[HumanMessage(content="Test message")],
        )

        with pytest.raises(AgentTerminationException) as excinfo:
            await node(state)

        assert (
            excinfo.value.error_info.code
            == f"Python.{UiPathErrorCode.EXECUTION_ERROR.value}"
        )
        assert excinfo.value.error_info.title == "Escalation rejected"

    @pytest.mark.asyncio
    async def test_extract_content_non_llm_scope_returns_empty(self):
        """Non-LLM scope in extraction: returns empty string."""
        from uipath_langchain.agent.guardrails.actions.escalate_action import (
            _extract_escalation_content,
        )

        state = AgentGuardrailsGraphState(
            messages=[HumanMessage(content="Test message")],
        )

        result = _extract_escalation_content(
            state, GuardrailScope.AGENT, ExecutionStage.PRE_EXECUTION
        )

        assert result == ""
