"""
Unit tests for FetchDocumentOperator.
"""

import pytest
import tempfile
import os
from unittest.mock import Mock, MagicMock, patch
from plugins.operators.fetch_document import FetchDocumentOperator


class TestFetchDocumentOperator:
    """Test suite for FetchDocumentOperator."""

    @pytest.fixture
    def mock_context(self):
        """Create a mock Airflow context."""
        context = {
            'ti': Mock(),
        }
        return context

    @pytest.fixture
    def sample_event(self):
        """Sample validated event data."""
        return {
            'bucket': 'test-bucket',
            'object_key': 'documents/sample.xlsx',
            'content_type': 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
            'size_bytes': 2048,
            'document_id': 'test-doc-123',
            'event_id': 'test-event',
        }

    def test_fetch_document_success(self, mock_context, sample_event):
        """Test successful document fetch."""
        mock_context['ti'].xcom_pull.return_value = sample_event

        with patch('plugins.operators.fetch_document.S3Hook') as mock_s3_hook_class:
            # Mock S3 object
            mock_obj = MagicMock()
            mock_obj.get.return_value = {'Body': MagicMock(read=lambda: b'test content')}

            # Mock hook instance
            mock_hook = MagicMock()
            mock_hook.get_key.return_value = mock_obj
            mock_s3_hook_class.return_value = mock_hook

            op = FetchDocumentOperator(
                task_id='test_fetch',
                connection_id='seaweedfs_default',
                validate_event_task_id='validate_event'
            )

            result = op.execute(mock_context)

            # Verify XCom pull was called
            mock_context['ti'].xcom_pull.assert_called_once_with(task_ids='validate_event')

            # Verify S3Hook was initialized with correct connection
            mock_s3_hook_class.assert_called_once_with(aws_conn_id='seaweedfs_default')

            # Verify file was fetched
            mock_hook.get_key.assert_called_once_with(
                key='documents/sample.xlsx',
                bucket_name='test-bucket'
            )

            # Verify result structure
            assert result['local_path'].endswith('sample.xlsx')
            assert result['content_type'] == sample_event['content_type']
            assert result['original_filename'] == 'sample.xlsx'
            assert result['file_extension'] == 'xlsx'
            assert result['document_id'] == 'test-doc-123'

            # Verify file was written
            assert os.path.exists(result['local_path'])

            # Clean up
            temp_dir = os.path.dirname(result['local_path'])
            import shutil
            shutil.rmtree(temp_dir)

    def test_fetch_no_event_raises(self, mock_context):
        """Test that missing event data raises ValueError."""
        mock_context['ti'].xcom_pull.return_value = None

        op = FetchDocumentOperator(task_id='test_fetch')

        with pytest.raises(ValueError, match='No event data found'):
            op.execute(mock_context)

    def test_fetch_object_not_found_raises(self, mock_context, sample_event):
        """Test that missing object in storage raises ValueError."""
        mock_context['ti'].xcom_pull.return_value = sample_event

        with patch('plugins.operators.fetch_document.S3Hook') as mock_s3_hook_class:
            mock_hook = MagicMock()
            mock_hook.get_key.return_value = None  # Object not found
            mock_s3_hook_class.return_value = mock_hook

            op = FetchDocumentOperator(task_id='test_fetch')

            with pytest.raises(ValueError, match='Object not found in storage'):
                op.execute(mock_context)

    def test_fetch_download_failure_cleans_up(self, mock_context, sample_event):
        """Test that download failure cleans up temporary directory."""
        mock_context['ti'].xcom_pull.return_value = sample_event

        with patch('plugins.operators.fetch_document.S3Hook') as mock_s3_hook_class:
            mock_hook = MagicMock()
            mock_hook.get_key.side_effect = Exception('Download failed')
            mock_s3_hook_class.return_value = mock_hook

            op = FetchDocumentOperator(task_id='test_fetch')

            with pytest.raises(Exception, match='Download failed'):
                op.execute(mock_context)

            # Verify cleanup would happen (temp_dir created then removed)
            # Note: In real execution, shutil.rmtree is called on failure

    def test_fetch_different_file_types(self, mock_context, sample_event):
        """Test fetching different file types."""
        file_types = [
            ('document.pdf', 'pdf', 'application/pdf'),
            ('presentation.pptx', 'pptx', 'application/vnd.openxmlformats-officedocument.presentationml.presentation'),
            ('spreadsheet.xls', 'xls', 'application/vnd.ms-excel'),
            ('text.docx', 'docx', 'application/vnd.openxmlformats-officedocument.wordprocessingml.document'),
        ]

        for filename, expected_ext, content_type in file_types:
            event = sample_event.copy()
            event['object_key'] = f'documents/{filename}'
            event['content_type'] = content_type

            mock_context['ti'].xcom_pull.return_value = event

            with patch('plugins.operators.fetch_document.S3Hook') as mock_s3_hook_class:
                mock_obj = MagicMock()
                mock_obj.get.return_value = {'Body': MagicMock(read=lambda: b'content')}
                mock_hook = MagicMock()
                mock_hook.get_key.return_value = mock_obj
                mock_s3_hook_class.return_value = mock_hook

                op = FetchDocumentOperator(task_id='test_fetch')
                result = op.execute(mock_context)

                assert result['original_filename'] == filename
                assert result['file_extension'] == expected_ext
                assert result['content_type'] == content_type

                # Clean up
                temp_dir = os.path.dirname(result['local_path'])
                import shutil
                shutil.rmtree(temp_dir)

    def test_custom_task_id_for_validation(self, mock_context, sample_event):
        """Test using custom validation task ID."""
        mock_context['ti'].xcom_pull.return_value = sample_event

        with patch('plugins.operators.fetch_document.S3Hook') as mock_s3_hook_class:
            mock_obj = MagicMock()
            mock_obj.get.return_value = {'Body': MagicMock(read=lambda: b'test')}
            mock_hook = MagicMock()
            mock_hook.get_key.return_value = mock_obj
            mock_s3_hook_class.return_value = mock_hook

            op = FetchDocumentOperator(
                task_id='test_fetch',
                validate_event_task_id='custom_validate'
            )

            result = op.execute(mock_context)

            # Verify custom task_id was used for XCom pull
            mock_context['ti'].xcom_pull.assert_called_once_with(task_ids='custom_validate')

            # Clean up
            temp_dir = os.path.dirname(result['local_path'])
            import shutil
            shutil.rmtree(temp_dir)