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

import pytest

sys.path.insert(0, os.path.abspath("../.."))

from litellm.integrations.langsmith import LangsmithLogger


class TestLangsmithLoggerInit:
    """Test cases for LangSmith logger initialization, particularly sampling rate handling.

    These tests verify that the sampling_rate attribute is set during initialization.
    Note: The current implementation has some edge cases in the sampling rate logic.
    """

    @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
    def test_langsmith_sampling_rate_parameter_respected_with_valid_env(self):
        """Test that langsmith_sampling_rate parameter is properly set when env var condition is met."""
        sampling_rate = 0.5
        logger = LangsmithLogger(
            langsmith_api_key="test-key",
            langsmith_project="test-project",
            langsmith_sampling_rate=sampling_rate,
        )

        assert (
            logger.sampling_rate == sampling_rate
        ), f"Expected sampling_rate to be {sampling_rate}, got {logger.sampling_rate}"

    @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
    def test_langsmith_sampling_rate_zero_parameter_falls_back_to_env(self):
        """Test that 0.0 parameter falls back to env var due to falsy value."""
        logger = LangsmithLogger(
            langsmith_api_key="test-key",
            langsmith_project="test-project",
            langsmith_sampling_rate=0.0,
        )

        assert (
            logger.sampling_rate == 1.0
        ), f"Expected sampling_rate to fall back to 1.0 from env, got {logger.sampling_rate}"

    @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
    def test_langsmith_sampling_rate_from_integer_env_var(self):
        """Test that sampling rate uses environment variable when parameter not provided and env var is integer."""
        logger = LangsmithLogger(
            langsmith_api_key="test-key", langsmith_project="test-project"
        )

        assert (
            logger.sampling_rate == 1.0
        ), f"Expected sampling_rate to be 1.0 from env var, got {logger.sampling_rate}"

    @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "0.8"}, clear=False)
    def test_langsmith_sampling_rate_decimal_env_var_ignored(self):
        """Test that decimal environment variables are ignored due to isdigit() check."""
        logger = LangsmithLogger(
            langsmith_api_key="test-key", langsmith_project="test-project"
        )

        assert (
            logger.sampling_rate == 1.0
        ), f"Expected sampling_rate to default to 1.0 (decimal env ignored), got {logger.sampling_rate}"

    @patch.dict(os.environ, {}, clear=True)
    def test_langsmith_sampling_rate_default_value(self):
        """Test that sampling rate defaults to 1.0 when no parameter or env var provided."""
        logger = LangsmithLogger(
            langsmith_api_key="test-key", langsmith_project="test-project"
        )

        assert (
            logger.sampling_rate == 1.0
        ), f"Expected default sampling_rate to be 1.0, got {logger.sampling_rate}"

    @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "invalid"}, clear=False)
    def test_langsmith_sampling_rate_invalid_env_var_defaults(self):
        """Test that invalid environment variable falls back to default value."""
        logger = LangsmithLogger(
            langsmith_api_key="test-key", langsmith_project="test-project"
        )

        assert (
            logger.sampling_rate == 1.0
        ), f"Expected sampling_rate to default to 1.0 with invalid env var, got {logger.sampling_rate}"

    @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": ""}, clear=False)
    def test_langsmith_sampling_rate_empty_env_var_defaults(self):
        """Test that empty environment variable falls back to default value."""
        logger = LangsmithLogger(
            langsmith_api_key="test-key", langsmith_project="test-project"
        )

        assert (
            logger.sampling_rate == 1.0
        ), f"Expected sampling_rate to default to 1.0 with empty env var, got {logger.sampling_rate}"

    def test_langsmith_sampling_rate_attribute_exists(self):
        """Test that the sampling_rate attribute is always set on the logger instance."""
        logger = LangsmithLogger(
            langsmith_api_key="test-key", langsmith_project="test-project"
        )

        assert hasattr(
            logger, "sampling_rate"
        ), "LangsmithLogger should have sampling_rate attribute"
        assert isinstance(
            logger.sampling_rate, float
        ), f"sampling_rate should be a float, got {type(logger.sampling_rate)}"
        assert (
            logger.sampling_rate >= 0.0
        ), f"sampling_rate should be non-negative, got {logger.sampling_rate}"

    @patch.object(LangsmithLogger, "_start_periodic_flush_task", return_value=None)
    def test_langsmith_init_skips_periodic_flush_without_running_loop(
        self, mock_start_periodic_flush_task
    ):
        """Test that sync initialization leaves the periodic flush task unset."""
        logger = LangsmithLogger(
            langsmith_api_key="test-key", langsmith_project="test-project"
        )

        assert logger is not None
        mock_start_periodic_flush_task.assert_called_once()
        assert logger._flush_task is None

    @patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop"))
    def test_start_periodic_flush_task_returns_none_without_running_loop(
        self, mock_get_running_loop
    ):
        """Test that helper returns None when no running event loop exists."""
        with patch.object(LangsmithLogger, "_start_periodic_flush_task", return_value=None):
            logger = LangsmithLogger(
                langsmith_api_key="test-key",
                langsmith_project="test-project",
            )

        mock_get_running_loop.reset_mock()

        assert logger._start_periodic_flush_task() is None
        mock_get_running_loop.assert_called_once()

    @patch("asyncio.get_running_loop")
    def test_langsmith_init_starts_periodic_flush_with_running_loop(
        self, mock_get_running_loop
    ):
        """Test that init schedules periodic flush when a running loop exists."""
        mock_loop = MagicMock()
        mock_task = MagicMock()
        mock_loop.create_task.return_value = mock_task
        mock_get_running_loop.return_value = mock_loop

        logger = LangsmithLogger(
            langsmith_api_key="test-key", langsmith_project="test-project"
        )

        assert logger._flush_task == mock_task
        mock_loop.create_task.assert_called_once()
        scheduled_coro = mock_loop.create_task.call_args.args[0]
        scheduled_coro.close()

    @pytest.mark.asyncio
    async def test_async_log_success_event_lazily_starts_periodic_flush(self):
        """Test that async logging lazily starts periodic flush after sync init."""
        with patch.object(LangsmithLogger, "_start_periodic_flush_task", return_value=None):
            logger = LangsmithLogger(
                langsmith_api_key="test-key",
                langsmith_project="test-project",
            )
        logger._get_sampling_rate_to_use_for_request = MagicMock(return_value=1.0)
        logger._get_credentials_to_use_for_request = MagicMock(
            return_value=logger.default_credentials
        )
        logger._prepare_log_data = MagicMock(return_value={"id": "run-id"})
        logger._start_periodic_flush_task = MagicMock(return_value=MagicMock())

        await logger.async_log_success_event({}, {}, None, None)

        logger._start_periodic_flush_task.assert_called_once()
        assert len(logger.log_queue) == 1

    @pytest.mark.asyncio
    async def test_async_log_failure_event_lazily_starts_periodic_flush(self):
        """Test that async failure logging lazily starts periodic flush after sync init."""
        with patch.object(LangsmithLogger, "_start_periodic_flush_task", return_value=None):
            logger = LangsmithLogger(
                langsmith_api_key="test-key",
                langsmith_project="test-project",
            )
        logger._get_sampling_rate_to_use_for_request = MagicMock(return_value=1.0)
        logger._get_credentials_to_use_for_request = MagicMock(
            return_value=logger.default_credentials
        )
        logger._prepare_log_data = MagicMock(return_value={"id": "run-id"})
        logger._start_periodic_flush_task = MagicMock(return_value=MagicMock())

        await logger.async_log_failure_event({}, {}, None, None)

        logger._start_periodic_flush_task.assert_called_once()
        assert len(logger.log_queue) == 1

class TestLangsmithPrepareLogData:
    """Regression test for #24001: _prepare_log_data must inject
    usage_metadata into outputs so LangSmith's Cost column is populated."""

    @patch("asyncio.create_task")
    @patch.dict(os.environ, {"LANGSMITH_SAMPLING_RATE": "1"}, clear=False)
    def test_outputs_contain_usage_metadata(self, mock_create_task):
        logger = LangsmithLogger(
            langsmith_api_key="test-key",
            langsmith_project="test-project",
        )

        payload = {
            "id": "test-id",
            "response": {"choices": [{"message": {"content": "hi"}}]},
            "metadata": {},
            "startTime": 1.0,
            "endTime": 2.0,
            "request_tags": [],
            "error_str": None,
            "status": "success",
            "response_cost": 0.0042,
            "prompt_tokens": 100,
            "completion_tokens": 50,
            "total_tokens": 150,
        }

        kwargs = {
            "litellm_params": {"metadata": {}},
            "standard_logging_object": payload,
        }

        credentials = {
            "LANGSMITH_API_KEY": "test-key",
            "LANGSMITH_PROJECT": "test-project",
            "LANGSMITH_BASE_URL": "https://api.smith.langchain.com",
        }

        data = logger._prepare_log_data(
            kwargs=kwargs,
            response_obj=None,
            start_time=1.0,
            end_time=2.0,
            credentials=credentials,
        )

        assert "usage_metadata" in data["outputs"]
        um = data["outputs"]["usage_metadata"]
        assert um["total_cost"] == 0.0042
        assert um["input_tokens"] == 100
        assert um["output_tokens"] == 50
        assert um["total_tokens"] == 150