"""
Test the CloudZero dry run endpoint functionality
"""

from unittest.mock import AsyncMock, MagicMock, patch

import polars as pl
import pytest


from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger


class TestCloudZeroDryRunEndpoint:
    """Test suite for CloudZero dry run endpoint functionality."""

    @pytest.mark.asyncio
    async def test_dry_run_export_usage_data_returns_data(self):
        """
        Test that dry_run_export_usage_data returns expected data structure
        instead of just logging to console.
        """
        logger = CloudZeroLogger()

        # Mock database data
        mock_usage_data = pl.DataFrame(
            {
                "date": ["2025-01-19", "2025-01-20"],
                "model": ["gpt-4", "gpt-3.5-turbo"],
                "custom_llm_provider": ["openai", "openai"],
                "team_id": ["team1", "team2"],
                "team_alias": ["Team One", "Team Two"],
                "api_key_alias": ["key1", "key2"],
                "user_email": ["one@example.com", None],
                "prompt_tokens": [100, 200],
                "completion_tokens": [50, 100],
                "spend": [0.01, 0.02],
                "successful_requests": [1, 2],
            }
        )

        # Mock CBF transformed data
        mock_cbf_data = pl.DataFrame(
            {
                "time/usage_start": ["2025-01-19T00:00:00Z", "2025-01-20T00:00:00Z"],
                "cost/cost": [0.01, 0.02],
                "usage/amount": [150, 300],
                "resource/service": ["openai", "openai"],
                "resource/account": ["litellm", "litellm"],
                "resource/region": ["us-east-1", "us-east-1"],
                "resource/id": ["gpt-4", "gpt-3.5-turbo"],
                "entity_type": ["user", "user"],
                "entity_id": ["team1", "team2"],
                "resource/tag:team_id": ["team1", "team2"],
                "resource/tag:team_alias": ["Team One", "Team Two"],
                "resource/tag:api_key_alias": ["key1", "key2"],
                "resource/tag:user_email": ["one@example.com", "N/A"],
            }
        )

        with (
            patch(
                "litellm.integrations.cloudzero.database.LiteLLMDatabase"
            ) as mock_db_class,
            patch(
                "litellm.integrations.cloudzero.transform.CBFTransformer"
            ) as mock_transformer_class,
        ):

            # Setup mocks
            mock_db = AsyncMock()
            mock_db.get_usage_data.return_value = mock_usage_data
            mock_db_class.return_value = mock_db

            mock_transformer = MagicMock()
            mock_transformer.transform.return_value = mock_cbf_data
            mock_transformer_class.return_value = mock_transformer

            # Call the method
            result = await logger.dry_run_export_usage_data(limit=1000)

            # Verify the result structure
            assert isinstance(result, dict)
            assert "usage_data" in result
            assert "cbf_data" in result
            assert "summary" in result

            # Verify usage_data
            assert isinstance(result["usage_data"], list)
            assert len(result["usage_data"]) == 2
            assert result["usage_data"][0]["model"] == "gpt-4"
            assert result["usage_data"][1]["model"] == "gpt-3.5-turbo"

            # Verify cbf_data
            assert isinstance(result["cbf_data"], list)
            assert len(result["cbf_data"]) == 2
            assert result["cbf_data"][0]["cost/cost"] == 0.01
            assert result["cbf_data"][1]["cost/cost"] == 0.02
            assert result["cbf_data"][0]["resource/tag:user_email"] == "one@example.com"

            # Verify summary
            summary = result["summary"]
            assert summary["total_records"] == 2
            assert summary["total_cost"] == 0.03
            assert summary["total_tokens"] == 450  # 150 + 300
            assert summary["unique_accounts"] == 1
            assert summary["unique_services"] == 1

    @pytest.mark.asyncio
    async def test_dry_run_export_usage_data_empty_data(self):
        """
        Test that dry_run_export_usage_data handles empty data gracefully.
        """
        logger = CloudZeroLogger()

        # Mock empty database data
        mock_empty_data = pl.DataFrame()

        with patch(
            "litellm.integrations.cloudzero.database.LiteLLMDatabase"
        ) as mock_db_class:

            # Setup mocks
            mock_db = AsyncMock()
            mock_db.get_usage_data.return_value = mock_empty_data
            mock_db_class.return_value = mock_db

            # Call the method
            result = await logger.dry_run_export_usage_data(limit=1000)

            # Verify the result structure for empty data
            assert isinstance(result, dict)
            assert result["usage_data"] == []
            assert result["cbf_data"] == []
            assert result["summary"]["total_records"] == 0
            assert result["summary"]["total_cost"] == 0
            assert result["summary"]["total_tokens"] == 0