"""
Tests for runpod.s3_client module.

Tests binary upload, log streaming, and result download.
"""

import pytest
from unittest.mock import MagicMock, patch, Mock
from pathlib import Path
from botocore.exceptions import ClientError, BotoCoreError
import io

from runpod.s3_client import S3Client, S3Error


class TestS3Client:
    """Tests for S3Client class."""

    @pytest.fixture
    def s3_config(self):
        """S3 configuration for testing."""
        return {
            "bucket_name": "se3zdnb5o4",
            "endpoint_url": "https://s3api-eur-is-1.runpod.io",
            "aws_access_key_id": "test_access_key",
            "aws_secret_access_key": "test_secret_key",
            "region_name": "eur-is-1"
        }

    def test_s3_client_initialization(self, s3_config):
        """Test S3 client initializes correctly."""
        with patch('boto3.client') as mock_boto:
            mock_boto.return_value = MagicMock()

            client = S3Client(**s3_config)

            assert client.bucket_name == "se3zdnb5o4"
            assert client.endpoint_url == "https://s3api-eur-is-1.runpod.io"
            mock_boto.assert_called_once()

    def test_s3_client_initialization_failure(self, s3_config):
        """Test S3 client initialization failure."""
        with patch('boto3.client') as mock_boto:
            mock_boto.side_effect = BotoCoreError()

            with pytest.raises(S3Error, match="Failed to initialize"):
                S3Client(**s3_config)

    def test_upload_file_success(self, s3_config, temp_directory):
        """Test successful file upload."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)

            # Create test file
            test_file = temp_directory / "test.bin"
            test_file.write_bytes(b"test data")

            result = client.upload_file(test_file, "binaries/test.bin")

            assert result is True
            mock_client.upload_file.assert_called_once()

    def test_upload_file_not_found(self, s3_config, temp_directory):
        """Test upload with non-existent file."""
        with patch('boto3.client') as mock_boto:
            mock_boto.return_value = MagicMock()

            client = S3Client(**s3_config)
            nonexistent = temp_directory / "nonexistent.bin"

            with pytest.raises(FileNotFoundError):
                client.upload_file(nonexistent, "binaries/test.bin")

    def test_upload_file_s3_error(self, s3_config, temp_directory):
        """Test upload with S3 error."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.upload_file.side_effect = ClientError(
                {"Error": {"Code": "NoSuchBucket", "Message": "Bucket not found"}},
                "upload_file"
            )
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)

            test_file = temp_directory / "test.bin"
            test_file.write_bytes(b"test data")

            with pytest.raises(S3Error, match="Failed to upload"):
                client.upload_file(test_file, "binaries/test.bin")

    def test_upload_file_with_extra_args(self, s3_config, temp_directory):
        """Test upload with extra arguments."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)

            test_file = temp_directory / "test.bin"
            test_file.write_bytes(b"test data")

            extra_args = {"ContentType": "application/octet-stream"}
            client.upload_file(test_file, "binaries/test.bin", extra_args)

            call_args = mock_client.upload_file.call_args
            assert call_args[1]["ExtraArgs"] == extra_args

    def test_upload_binaries_success(self, s3_config, sample_binary_files):
        """Test uploading directory of binaries."""
        binaries_dir, binary_files = sample_binary_files

        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            uploaded = client.upload_binaries(binaries_dir, "binaries/")

            # Should upload all binaries
            assert len(uploaded) == 4
            assert all(key.startswith("binaries/") for key in uploaded)
            assert mock_client.upload_file.call_count == 4

    def test_upload_binaries_not_directory(self, s3_config, temp_directory):
        """Test upload_binaries with non-directory."""
        with patch('boto3.client') as mock_boto:
            mock_boto.return_value = MagicMock()

            client = S3Client(**s3_config)
            not_a_dir = temp_directory / "file.txt"
            not_a_dir.write_text("test")

            with pytest.raises(NotADirectoryError):
                client.upload_binaries(not_a_dir)

    def test_download_file_success(self, s3_config, temp_directory):
        """Test successful file download."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            local_path = temp_directory / "downloaded.bin"

            result = client.download_file("models/test.safetensors", local_path)

            assert result is True
            mock_client.download_file.assert_called_once()

    def test_download_file_creates_directories(self, s3_config, temp_directory):
        """Test download creates parent directories."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            local_path = temp_directory / "nested" / "dir" / "file.bin"

            client.download_file("models/test.bin", local_path)

            assert local_path.parent.exists()

    def test_download_file_s3_error(self, s3_config, temp_directory):
        """Test download with S3 error."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.download_file.side_effect = ClientError(
                {"Error": {"Code": "NoSuchKey", "Message": "Key not found"}},
                "download_file"
            )
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            local_path = temp_directory / "test.bin"

            with pytest.raises(S3Error, match="Failed to download"):
                client.download_file("nonexistent.bin", local_path)

    def test_stream_logs_from_start(self, s3_config, mock_log_content):
        """Test streaming logs from start."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_response = {
                'Body': io.BytesIO(mock_log_content.encode('utf-8'))
            }
            mock_client.get_object.return_value = mock_response
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            content = client.stream_logs("logs/test.log", start_byte=0)

            assert "Starting training" in content
            mock_client.get_object.assert_called_once()

    def test_stream_logs_with_offset(self, s3_config):
        """Test streaming logs from specific byte offset."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            partial_content = b"Partial log content"
            mock_response = {
                'Body': io.BytesIO(partial_content)
            }
            mock_client.get_object.return_value = mock_response
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            content = client.stream_logs("logs/test.log", start_byte=1000)

            assert content == "Partial log content"
            # Verify Range header was used
            call_args = mock_client.get_object.call_args
            assert call_args[1]["Range"] == "bytes=1000-"

    def test_stream_logs_tail_bytes(self, s3_config):
        """Test streaming last N bytes of logs."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()

            # Mock head_object to return file size
            mock_client.head_object.return_value = {'ContentLength': 5000}

            # Mock get_object for range request
            tail_content = b"Last 1KB of logs"
            mock_response = {
                'Body': io.BytesIO(tail_content)
            }
            mock_client.get_object.return_value = mock_response
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            content = client.stream_logs("logs/test.log", tail_bytes=1024)

            assert content == "Last 1KB of logs"
            # Should request from byte 3976 (5000 - 1024)
            call_args = mock_client.get_object.call_args
            assert call_args[1]["Range"] == "bytes=3976-"

    def test_stream_logs_nonexistent_file(self, s3_config):
        """Test streaming logs from nonexistent file."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.get_object.side_effect = ClientError(
                {"Error": {"Code": "NoSuchKey", "Message": "Key not found"}},
                "get_object"
            )
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            content = client.stream_logs("logs/nonexistent.log")

            # Should return empty string for nonexistent logs
            assert content == ""

    def test_stream_logs_s3_error(self, s3_config):
        """Test streaming logs with S3 error."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.get_object.side_effect = ClientError(
                {"Error": {"Code": "AccessDenied", "Message": "Access denied"}},
                "get_object"
            )
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)

            with pytest.raises(S3Error, match="Failed to stream logs"):
                client.stream_logs("logs/test.log")

    def test_list_files_success(self, s3_config, mock_s3_file_list):
        """Test listing files in S3."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.list_objects_v2.return_value = mock_s3_file_list
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            files = client.list_files(prefix="models/")

            assert len(files) == 3
            assert files[0]["Key"] == "models/tft_best.safetensors"
            assert files[0]["Size"] == 1024000

    def test_list_files_empty(self, s3_config):
        """Test listing files with no results."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.list_objects_v2.return_value = {}  # No Contents field
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            files = client.list_files(prefix="nonexistent/")

            assert files == []

    def test_list_files_with_max_keys(self, s3_config):
        """Test listing files with max_keys parameter."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.list_objects_v2.return_value = {"Contents": []}
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            client.list_files(prefix="models/", max_keys=500)

            call_args = mock_client.list_objects_v2.call_args
            assert call_args[1]["MaxKeys"] == 500

    def test_list_files_s3_error(self, s3_config):
        """Test list_files with S3 error."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.list_objects_v2.side_effect = ClientError(
                {"Error": {"Code": "AccessDenied", "Message": "Access denied"}},
                "list_objects_v2"
            )
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)

            with pytest.raises(S3Error, match="Failed to list files"):
                client.list_files(prefix="models/")

    def test_download_results_success(self, s3_config, temp_directory):
        """Test downloading training results."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            # Mock only returns model files (not logs)
            mock_client.list_objects_v2.return_value = {
                "Contents": [
                    {
                        "Key": "models/tft_best.safetensors",
                        "Size": 1024000,
                        "LastModified": "2025-10-29T12:30:00Z"
                    },
                    {
                        "Key": "models/tft_config.json",
                        "Size": 1024,
                        "LastModified": "2025-10-29T12:30:01Z"
                    }
                ]
            }
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            downloaded = client.download_results("models/", temp_directory)

            # Should download 2 model files
            assert len(downloaded) == 2
            assert mock_client.download_file.call_count == 2

    def test_download_results_with_patterns(self, s3_config, mock_s3_file_list, temp_directory):
        """Test downloading results with file patterns."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.list_objects_v2.return_value = mock_s3_file_list
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            downloaded = client.download_results(
                "models/",
                temp_directory,
                file_patterns=["*.safetensors"]
            )

            # Should only download .safetensors files
            assert len(downloaded) == 1
            assert downloaded[0].name == "tft_best.safetensors"

    def test_download_results_preserves_structure(self, s3_config, temp_directory):
        """Test download preserves directory structure."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.list_objects_v2.return_value = {
                "Contents": [
                    {"Key": "models/subdir/file.bin", "Size": 1024, "LastModified": "2025-10-29"}
                ]
            }
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            downloaded = client.download_results("models/", temp_directory)

            # Should preserve subdirectory
            assert len(downloaded) == 1
            assert "subdir" in str(downloaded[0])

    def test_delete_file_success(self, s3_config):
        """Test deleting a file."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            result = client.delete_file("models/old_model.bin")

            assert result is True
            mock_client.delete_object.assert_called_once()

    def test_delete_file_error(self, s3_config):
        """Test delete with S3 error."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()
            mock_client.delete_object.side_effect = ClientError(
                {"Error": {"Code": "AccessDenied", "Message": "Access denied"}},
                "delete_object"
            )
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)

            with pytest.raises(S3Error, match="Failed to delete"):
                client.delete_file("models/test.bin")

    @pytest.mark.parametrize("prefix,expected_count", [
        ("models/", 2),  # 2 model files
        ("logs/", 1),    # 1 log file
        ("", 3),         # All files
    ])
    def test_list_files_with_different_prefixes(self, s3_config, mock_s3_file_list, prefix, expected_count):
        """Test listing files with different prefixes."""
        with patch('boto3.client') as mock_boto:
            mock_client = MagicMock()

            # Filter mock data based on prefix
            filtered_contents = [
                item for item in mock_s3_file_list["Contents"]
                if item["Key"].startswith(prefix)
            ]
            mock_client.list_objects_v2.return_value = {"Contents": filtered_contents}
            mock_boto.return_value = mock_client

            client = S3Client(**s3_config)
            files = client.list_files(prefix=prefix)

            assert len(files) == expected_count