import json
import os
from unittest.mock import MagicMock, patch

import pytest

from minisweagent.models import GLOBAL_MODEL_STATS
from minisweagent.models.portkey_model import PortkeyModel, PortkeyModelConfig
from minisweagent.models.utils.actions_toolcall import BASH_TOOL


def test_portkey_model_missing_api_key():
    """Test that PortkeyModel raises ValueError when no API key is provided."""
    with patch("minisweagent.models.portkey_model.Portkey"):
        with patch.dict(os.environ, {}, clear=True):
            with pytest.raises(ValueError, match="Portkey API key is required"):
                PortkeyModel(model_name="gpt-4o")


def test_portkey_model_config():
    """Test PortkeyModelConfig creation."""
    config = PortkeyModelConfig(model_name="gpt-4o", model_kwargs={"temperature": 0.7})
    assert config.model_name == "gpt-4o"
    assert config.model_kwargs == {"temperature": 0.7}


def test_portkey_model_initialization():
    """Test PortkeyModel initialization with mocked Portkey."""
    mock_portkey_class = MagicMock()
    mock_client = MagicMock()
    mock_portkey_class.return_value = mock_client

    with patch("minisweagent.models.portkey_model.Portkey", mock_portkey_class):
        with patch.dict(os.environ, {"PORTKEY_API_KEY": "test-key", "PORTKEY_VIRTUAL_KEY": "test-virtual"}):
            model = PortkeyModel(model_name="gpt-4o")

            assert model.config.model_name == "gpt-4o"

            # Verify Portkey was called with correct parameters
            mock_portkey_class.assert_called_once_with(api_key="test-key", virtual_key="test-virtual")


def test_portkey_model_query():
    """Test PortkeyModel.query method with mocked response."""
    mock_portkey_class = MagicMock()
    mock_client = MagicMock()
    mock_response = MagicMock()
    mock_choice = MagicMock()
    mock_message = MagicMock()
    mock_tool_call = MagicMock()

    # Response uses tool_calls
    mock_tool_call.id = "call_123"
    mock_tool_call.function.name = "bash"
    mock_tool_call.function.arguments = json.dumps({"command": "echo 'Hello!'"})
    mock_message.tool_calls = [mock_tool_call]
    mock_message.content = None
    mock_message.model_dump.return_value = {
        "role": "assistant",
        "content": None,
        "tool_calls": [{"id": "call_123", "function": {"name": "bash", "arguments": '{"command": "echo \'Hello!\'"}'}}],
    }
    mock_choice.message = mock_message
    mock_response.choices = [mock_choice]
    mock_response.model_dump.return_value = {"test": "response"}

    mock_client.chat.completions.create.return_value = mock_response
    mock_portkey_class.return_value = mock_client

    with patch("minisweagent.models.portkey_model.Portkey", mock_portkey_class):
        with patch.dict(os.environ, {"PORTKEY_API_KEY": "test-key"}):
            with patch("minisweagent.models.portkey_model.litellm.cost_calculator.completion_cost") as mock_cost:
                mock_cost.return_value = 0.01

                model = PortkeyModel(model_name="gpt-4o")

                messages = [{"role": "user", "content": "Hello!"}]
                result = model.query(messages)

                assert result["extra"]["actions"] == [{"command": "echo 'Hello!'", "tool_call_id": "call_123"}]
                assert result["extra"]["response"] == {"test": "response"}
                assert result["extra"]["cost"] == 0.01

                # Verify the API was called correctly with tools
                mock_client.chat.completions.create.assert_called_once_with(
                    model="gpt-4o", messages=messages, tools=[BASH_TOOL]
                )
                # Verify cost calculation was called
                mock_cost.assert_called_once_with(mock_response.model_copy(), model=None)


def test_portkey_model_get_template_vars():
    """Test PortkeyModel.get_template_vars method."""
    mock_portkey_class = MagicMock()
    mock_client = MagicMock()
    mock_portkey_class.return_value = mock_client

    with patch("minisweagent.models.portkey_model.Portkey", mock_portkey_class):
        with patch.dict(os.environ, {"PORTKEY_API_KEY": "test-key"}):
            model = PortkeyModel(model_name="gpt-4o", model_kwargs={"temperature": 0.7})

            template_vars = model.get_template_vars()

            assert template_vars["model_name"] == "gpt-4o"
            assert template_vars["model_kwargs"] == {"temperature": 0.7}


def test_portkey_model_cost_tracking_ignore_errors():
    """Test that models work with cost_tracking='ignore_errors'."""
    mock_portkey_class = MagicMock()
    mock_client = MagicMock()
    mock_response = MagicMock()
    mock_choice = MagicMock()
    mock_message = MagicMock()
    mock_tool_call = MagicMock()

    # Response uses tool_calls
    mock_tool_call.id = "call_456"
    mock_tool_call.function.name = "bash"
    mock_tool_call.function.arguments = json.dumps({"command": "echo test"})
    mock_message.tool_calls = [mock_tool_call]
    mock_message.content = None
    mock_message.model_dump.return_value = {
        "role": "assistant",
        "content": None,
        "tool_calls": [{"id": "call_456", "function": {"name": "bash", "arguments": '{"command": "echo test"}'}}],
    }
    mock_choice.message = mock_message
    mock_response.choices = [mock_choice]
    mock_response.model_dump.return_value = {"test": "response"}

    mock_client.chat.completions.create.return_value = mock_response
    mock_portkey_class.return_value = mock_client

    with patch("minisweagent.models.portkey_model.Portkey", mock_portkey_class):
        with patch.dict(os.environ, {"PORTKEY_API_KEY": "test-key"}):
            model = PortkeyModel(model_name="gpt-4o", cost_tracking="ignore_errors")

            initial_cost = GLOBAL_MODEL_STATS.cost

            with patch(
                "minisweagent.models.portkey_model.litellm.cost_calculator.completion_cost",
                side_effect=ValueError("Model not found"),
            ):
                messages = [{"role": "user", "content": "test"}]
                result = model.query(messages)

                assert result["extra"]["actions"] == [{"command": "echo test", "tool_call_id": "call_456"}]
                assert result["extra"]["cost"] == 0.0
                assert GLOBAL_MODEL_STATS.cost == initial_cost


def test_portkey_model_cost_validation_error():
    """Test that cost calculation errors raise RuntimeError when cost tracking is enabled."""
    mock_portkey_class = MagicMock()
    mock_client = MagicMock()
    mock_response = MagicMock()
    mock_choice = MagicMock()
    mock_message = MagicMock()
    mock_usage = MagicMock()
    mock_tool_call = MagicMock()

    mock_tool_call.id = "call_789"
    mock_tool_call.function.name = "bash"
    mock_tool_call.function.arguments = json.dumps({"command": "echo test"})
    mock_message.tool_calls = [mock_tool_call]
    mock_message.content = None
    mock_choice.message = mock_message
    mock_response.choices = [mock_choice]
    mock_response.model_dump.return_value = {"test": "response"}
    mock_response.model_copy.return_value = mock_response
    mock_response.usage = mock_usage
    mock_usage.prompt_tokens = 10
    mock_usage.completion_tokens = 20
    mock_usage.total_tokens = 30

    mock_client.chat.completions.create.return_value = mock_response
    mock_portkey_class.return_value = mock_client

    with patch("minisweagent.models.portkey_model.Portkey", mock_portkey_class):
        with patch.dict(os.environ, {"PORTKEY_API_KEY": "test-key"}):
            model = PortkeyModel(model_name="gpt-4o")

            with patch("minisweagent.models.portkey_model.litellm.cost_calculator.completion_cost") as mock_cost:
                mock_cost.side_effect = ValueError("Model not found")

                messages = [{"role": "user", "content": "test"}]

                with pytest.raises(RuntimeError) as exc_info:
                    model.query(messages)

                assert "Error calculating cost" in str(exc_info.value)
                assert "MSWEA_COST_TRACKING='ignore_errors'" in str(exc_info.value)