import pytest
import os
import tempfile
import importlib.util
from unittest.mock import patch, MagicMock
from litellm.proxy.types_utils.utils import (
    get_instance_fn,
    _load_instance_from_remote_storage,
)


class TestCustomLoggerS3GCS:
    """Test custom logger loading from S3/GCS using URL prefixes"""

    @pytest.fixture
    def sample_custom_logger_content(self):
        """Sample custom logger file content"""
        return """
from litellm.integrations.custom_logger import CustomLogger

class TestCustomLogger(CustomLogger):
    def __init__(self):
        super().__init__()
        self.initialized = True
    
    def log_pre_api_call(self, model, messages, kwargs):
        print(f"Pre-API call to {model}")

# Instance to be imported
test_logger_instance = TestCustomLogger()
"""

    @pytest.fixture
    def temp_config_dir(self):
        """Create a temporary directory for config files"""
        with tempfile.TemporaryDirectory() as temp_dir:
            yield temp_dir

    def test_local_file_loading_still_works(
        self, temp_config_dir, sample_custom_logger_content
    ):
        """Test that local file loading continues to work (no URL prefix)"""
        # Create a local custom logger file
        custom_logger_path = os.path.join(temp_config_dir, "test_custom_logger.py")
        with open(custom_logger_path, "w") as f:
            f.write(sample_custom_logger_content)

        # Create a dummy config file
        config_path = os.path.join(temp_config_dir, "config.yaml")
        with open(config_path, "w") as f:
            f.write("model_list: []")

        # Test loading the custom logger (traditional way)
        instance = get_instance_fn(
            "test_custom_logger.test_logger_instance", config_path
        )

        assert instance is not None
        assert hasattr(instance, "initialized")
        assert instance.initialized is True

    def test_s3_url_parsing(self):
        """Test S3 URL parsing"""
        test_url = "s3://my-bucket/loggers/custom_callbacks.proxy_handler_instance"

        # Mock the download function to avoid actual S3 calls
        with patch(
            "litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3"
        ) as mock_download:
            mock_download.return_value = (
                False  # Will cause failure, but we just want to test parsing
            )

            with pytest.raises(ImportError, match="Failed to download"):
                _load_instance_from_remote_storage(test_url)

            # Verify the download was called with correct parameters
            mock_download.assert_called_once()
            call_args = mock_download.call_args
            assert call_args.kwargs["bucket_name"] == "my-bucket"
            assert call_args.kwargs["object_key"] == "loggers/custom_callbacks.py"

    def test_gcs_url_parsing(self):
        """Test GCS URL parsing"""
        test_url = "gcs://my-bucket/custom_logger.my_instance"

        # Mock the download function
        with patch(
            "litellm.proxy.types_utils.utils._download_gcs_file_wrapper"
        ) as mock_download:
            mock_download.return_value = False  # Will cause failure

            with pytest.raises(ImportError, match="Failed to download"):
                _load_instance_from_remote_storage(test_url)

            # Verify the download was called with correct parameters
            mock_download.assert_called_once()
            call_args = mock_download.call_args
            assert (
                call_args[0][0] == "my-bucket"
            )  # bucket_name (positional for _download_gcs_file_wrapper)
            assert call_args[0][1] == "custom_logger.py"  # object_key

    @patch("litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3")
    def test_s3_download_success(self, mock_s3_download, sample_custom_logger_content):
        """Test successful S3 download and loading"""

        # Configure S3 download to succeed and create the file
        def mock_download(bucket_name, object_key, local_file_path):
            with open(local_file_path, "w") as f:
                f.write(sample_custom_logger_content)
            return True

        mock_s3_download.side_effect = mock_download

        # Test loading with S3 URL — pass config_file_path to indicate
        # this is a startup config-file load (the documented operator
        # flow that the runtime gate preserves).
        test_url = "s3://test-bucket/test_custom_logger.test_logger_instance"
        instance = get_instance_fn(test_url, config_file_path="/any/path")

        assert instance is not None
        assert hasattr(instance, "initialized")
        assert instance.initialized is True

        # Verify S3 download was called with correct parameters
        mock_s3_download.assert_called_once()
        call_args = mock_s3_download.call_args
        assert call_args.kwargs["bucket_name"] == "test-bucket"
        assert call_args.kwargs["object_key"] == "test_custom_logger.py"

    @patch("litellm.proxy.types_utils.utils._download_gcs_file_wrapper")
    def test_gcs_download_success(
        self, mock_gcs_download, sample_custom_logger_content
    ):
        """Test successful GCS download and loading"""

        # Configure GCS download to succeed and create the file
        def mock_download(bucket_name, object_key, local_file_path):
            with open(local_file_path, "w") as f:
                f.write(sample_custom_logger_content)
            return True

        mock_gcs_download.side_effect = mock_download

        # Test loading with GCS URL (startup config-file load path).
        test_url = "gcs://test-bucket/test_custom_logger.test_logger_instance"
        instance = get_instance_fn(test_url, config_file_path="/any/path")

        assert instance is not None
        assert hasattr(instance, "initialized")
        assert instance.initialized is True

    def test_nested_path_parsing(self):
        """Test parsing of nested paths in URLs"""
        test_url = "s3://my-bucket/loggers/production/advanced_logger.handler_instance"

        with patch(
            "litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3"
        ) as mock_download:
            mock_download.return_value = False

            with pytest.raises(ImportError):
                _load_instance_from_remote_storage(test_url)

            # Verify correct object key was generated
            call_args = mock_download.call_args
            assert (
                call_args.kwargs["object_key"]
                == "loggers/production/advanced_logger.py"
            )

    def test_invalid_url_schemes(self):
        """Test error handling for invalid URL schemes"""
        # URLs that look like URLs but aren't s3:// or gcs:// will be treated as module names
        # and fail with regular ImportError
        with pytest.raises(ImportError):
            get_instance_fn("http://bucket/module.instance")

        with pytest.raises(ImportError):
            get_instance_fn("ftp://bucket/module.instance")

    def test_invalid_url_format(self):
        """Test error handling for invalid URL formats (config-file load path)."""
        # Missing bucket
        with pytest.raises(ImportError, match="Invalid URL format"):
            get_instance_fn("s3://", config_file_path="/any/path")

        # Missing path
        with pytest.raises(ImportError, match="Invalid URL format"):
            get_instance_fn("s3://bucket-only", config_file_path="/any/path")

        # Missing instance name
        with pytest.raises(ImportError, match="Invalid module specification"):
            get_instance_fn("s3://bucket/module-only", config_file_path="/any/path")

        # Including .py extension (common mistake)
        with pytest.raises(
            ImportError,
            match="Don't include '\\.py' extension and you must specify the instance name",
        ):
            get_instance_fn(
                "s3://bucket/custom_guardrail.py", config_file_path="/any/path"
            )

    @patch("litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3")
    def test_download_failure_handling(self, mock_s3_download):
        """Test handling of download failures"""
        mock_s3_download.return_value = False

        test_url = "s3://test-bucket/failing_logger.instance"

        with pytest.raises(ImportError, match="Failed to download"):
            get_instance_fn(test_url, config_file_path="/any/path")

    @patch("litellm.proxy.common_utils.load_config_utils.download_python_file_from_s3")
    def test_file_cleanup(self, mock_s3_download, sample_custom_logger_content):
        """Test that temporary files are cleaned up"""
        created_files = []

        def mock_download(bucket_name, object_key, local_file_path):
            created_files.append(local_file_path)
            with open(local_file_path, "w") as f:
                f.write(sample_custom_logger_content)
            return True

        mock_s3_download.side_effect = mock_download

        test_url = "s3://test-bucket/test_custom_logger.test_logger_instance"
        instance = get_instance_fn(test_url, config_file_path="/any/path")

        assert instance is not None

        # Verify file was created and then cleaned up
        assert len(created_files) == 1
        temp_file = created_files[0]
        assert not os.path.exists(
            temp_file
        ), f"Temporary file {temp_file} was not cleaned up"

    def test_no_url_prefix_fallback(self, temp_config_dir):
        """Test fallback when no URL prefix is used and local file doesn't exist"""
        config_path = os.path.join(temp_config_dir, "config.yaml")
        with open(config_path, "w") as f:
            f.write("model_list: []")

        # Test that it tries local loading when no URL prefix is used
        with pytest.raises(
            ImportError, match="Could not import instance from nonexistent_logger"
        ):
            get_instance_fn("nonexistent_logger.instance", config_path)