"""Tests for RLM cost and timeout limiting features."""

import time
from unittest.mock import MagicMock, patch

import pytest

from rlm.core.exceptions import CompletionTimeoutError, CostLimitExceededError
from rlm.core.rlm import RLM
from rlm.core.types import ModelUsageSummary, UsageSummary


class TestRLMCostLimiting:
    """Test cost limiting functionality in RLM."""

    @patch("rlm.core.rlm.VerbosePrinter")
    def test_max_cost_parameter_stored(self, mock_verbose):
        """Test that max_cost parameter is stored in RLM instance."""
        rlm = RLM(max_cost=10.0)
        assert rlm.max_cost == 10.0

    @patch("rlm.core.rlm.VerbosePrinter")
    def test_max_cost_defaults_to_none(self, mock_verbose):
        """Test that max_cost defaults to None (no limit)."""
        rlm = RLM()
        assert rlm.max_cost is None

    @patch("rlm.core.rlm.LMHandler")
    @patch("rlm.core.rlm.get_environment")
    @patch("rlm.core.rlm.get_client")
    def test_cost_limit_raises_when_exceeded(self, mock_get_client, mock_get_env, mock_lm_handler):
        """Test that CostLimitExceededError is raised when cost limit is exceeded."""
        # Mock the client
        mock_client = MagicMock()
        mock_client.model_name = "gpt-4"
        mock_client.completion.return_value = "```repl\nprint('test')\n```"
        mock_get_client.return_value = mock_client

        # Mock the environment
        mock_env = MagicMock()
        mock_env.execute_code.return_value = MagicMock(
            stdout="test", stderr="", locals={}, execution_time=0.1, rlm_calls=[]
        )
        mock_get_env.return_value = mock_env

        # Mock LMHandler at class level
        mock_handler = MagicMock()
        mock_handler.host = "localhost"
        mock_handler.port = 12345
        mock_handler.completion.return_value = "```repl\nprint('test')\n```"
        mock_handler.get_usage_summary.return_value = UsageSummary(
            model_usage_summaries={
                "gpt-4": ModelUsageSummary(
                    total_calls=10,
                    total_input_tokens=10000,
                    total_output_tokens=10000,
                )
            }
        )
        mock_lm_handler.return_value = mock_handler

        # Create RLM with cost limit of -1 (cost will be 0, but check < 0 triggers)
        rlm = RLM(backend_kwargs={"model_name": "gpt-4"}, max_cost=-1.0)

        # Should raise CostLimitExceededError (0.0 > -1.0)
        with pytest.raises(CostLimitExceededError) as exc_info:
            rlm.completion("test prompt")

        assert exc_info.value.current_cost == 0.0
        assert exc_info.value.cost_limit == -1.0


class TestRLMTimeoutPolicies:
    """Test timeout policies in RLM."""

    @patch("rlm.core.rlm.VerbosePrinter")
    def test_timeout_seconds_parameter_stored(self, mock_verbose):
        """Test that timeout_seconds parameter is stored."""
        rlm = RLM(timeout_seconds=120.0)
        assert rlm.timeout_seconds == 120.0

    @patch("rlm.core.rlm.VerbosePrinter")
    def test_iteration_timeout_seconds_parameter_stored(self, mock_verbose):
        """Test that iteration_timeout_seconds parameter is stored."""
        rlm = RLM(iteration_timeout_seconds=30.0)
        assert rlm.iteration_timeout_seconds == 30.0

    @patch("rlm.core.rlm.VerbosePrinter")
    def test_timeouts_default_to_none(self, mock_verbose):
        """Test that timeout parameters default to None (no limit)."""
        rlm = RLM()
        assert rlm.timeout_seconds is None
        assert rlm.iteration_timeout_seconds is None

    @patch("rlm.core.rlm.LMHandler")
    @patch("rlm.core.rlm.get_environment")
    @patch("rlm.core.rlm.get_client")
    @patch("rlm.core.rlm.time.perf_counter")
    def test_completion_timeout_raises_when_exceeded(
        self, mock_perf_counter, mock_get_client, mock_get_env, mock_lm_handler
    ):
        """Test that CompletionTimeoutError is raised when total timeout is exceeded."""
        # Mock time to simulate timeout
        # Need enough values: start, iter_start1, elapsed_check1, iter_start2, elapsed_check2
        mock_perf_counter.side_effect = [0.0, 0.1, 0.2, 150.0, 150.1]

        # Mock the client
        mock_client = MagicMock()
        mock_client.model_name = "gpt-4"
        mock_client.completion.return_value = "Thinking..."
        mock_get_client.return_value = mock_client

        # Mock the environment
        mock_env = MagicMock()
        mock_env.execute_code.return_value = MagicMock(
            stdout="", stderr="", locals={}, execution_time=0.1, rlm_calls=[]
        )
        mock_env.get_context_count.return_value = 1
        mock_env.get_history_count.return_value = 0
        mock_get_env.return_value = mock_env

        # Mock LMHandler at class level
        mock_handler = MagicMock()
        mock_handler.host = "localhost"
        mock_handler.port = 12345
        mock_handler.completion.return_value = "Thinking about the problem..."
        mock_handler.get_usage_summary.return_value = UsageSummary(
            model_usage_summaries={}
        )
        mock_lm_handler.return_value = mock_handler

        # Create RLM with short timeout
        rlm = RLM(backend_kwargs={"model_name": "gpt-4"}, timeout_seconds=120.0)

        # Should raise CompletionTimeoutError before completing
        with pytest.raises(CompletionTimeoutError) as exc_info:
            rlm.completion("test prompt")

        assert exc_info.value.timeout_seconds == 120.0
        assert exc_info.value.elapsed_seconds == 150.1


class TestRLMParameterDefaults:
    """Test that new parameters don't break existing usage."""

    @patch("rlm.core.rlm.VerbosePrinter")
    def test_rlm_works_without_limits(self, mock_verbose):
        """Test that RLM can be instantiated without any limits."""
        rlm = RLM()
        assert rlm.max_cost is None
        assert rlm.timeout_seconds is None
        assert rlm.iteration_timeout_seconds is None

    @patch("rlm.core.rlm.VerbosePrinter")
    def test_rlm_works_with_partial_limits(self, mock_verbose):
        """Test that RLM works with only some limits set."""
        rlm1 = RLM(max_cost=10.0)
        assert rlm1.max_cost == 10.0
        assert rlm1.timeout_seconds is None

        rlm2 = RLM(timeout_seconds=120.0)
        assert rlm2.max_cost is None
        assert rlm2.timeout_seconds == 120.0

    @patch("rlm.core.rlm.VerbosePrinter")
    def test_rlm_works_with_all_limits(self, mock_verbose):
        """Test that RLM works with all limits set."""
        rlm = RLM(max_cost=10.0, timeout_seconds=120.0, iteration_timeout_seconds=30.0)
        assert rlm.max_cost == 10.0
        assert rlm.timeout_seconds == 120.0
        assert rlm.iteration_timeout_seconds == 30.0


class TestExceptionAttributes:
    """Test that exceptions contain the expected attributes."""

    def test_cost_limit_error_attributes(self):
        """Test CostLimitExceededError contains cost information."""
        error = CostLimitExceededError(
            "Cost $15.00 exceeded limit $10.00", current_cost=15.0, cost_limit=10.0
        )
        assert error.current_cost == 15.0
        assert error.cost_limit == 10.0
        assert "$15" in str(error)
        assert "$10" in str(error)

    def test_timeout_error_attributes(self):
        """Test CompletionTimeoutError contains timing information."""
        error = CompletionTimeoutError(
            "Timeout after 120s", timeout_seconds=120.0, elapsed_seconds=150.5
        )
        assert error.timeout_seconds == 120.0
        assert error.elapsed_seconds == 150.5
        assert "120" in str(error)
