"""
Unit tests for Word Document Preprocessing.

Tests cover:
- Image extraction from Word documents
- AI image captioning via GPT-4o vision
- Shape/text box extraction
- Comment extraction
- Inline embedding of all content
"""

import pytest
from pathlib import Path
from unittest.mock import Mock, MagicMock, patch
from docx import Document

from src.extraction_v2.word_preprocessor import WordPreprocessor


class TestWordPreprocessor:
    """Test suite for WordPreprocessor class."""

    @pytest.fixture
    def mock_openai_client(self):
        """Create a mock Azure OpenAI client."""
        client = Mock()
        response = Mock()
        response.choices = [Mock(message=Mock(content="Test image caption"))]
        client.chat.completions.create.return_value = response
        return client

    @pytest.fixture
    def preprocessor(self, mock_openai_client):
        """Create a WordPreprocessor instance with mocked client."""
        return WordPreprocessor(openai_client=mock_openai_client)

    def test_extract_images_empty_document(self, preprocessor):
        """Test image extraction on document with no images."""
        # Create empty document
        doc = Document()
        doc.add_paragraph("This is a test paragraph with no images.")

        images = preprocessor._extract_images(doc)

        assert images == []

    def test_get_paragraph_context(self, preprocessor):
        """Test paragraph context extraction."""
        doc = Document()
        doc.add_paragraph("Paragraph 1")
        doc.add_paragraph("Paragraph 2")
        doc.add_paragraph("Paragraph 3 - Target")
        doc.add_paragraph("Paragraph 4")
        doc.add_paragraph("Paragraph 5")

        # Get context around paragraph 2 (index 2) with window=2
        context = preprocessor._get_paragraph_context(doc, para_idx=2, window=2)

        # Should include paragraphs 0, 1, 2, 3, 4
        assert "Paragraph 1" in context
        assert "Paragraph 2" in context
        assert "Paragraph 3 - Target" in context
        assert "Paragraph 4" in context
        assert "Paragraph 5" in context

    def test_get_paragraph_context_edge_case_start(self, preprocessor):
        """Test paragraph context at document start."""
        doc = Document()
        doc.add_paragraph("First paragraph")
        doc.add_paragraph("Second paragraph")
        doc.add_paragraph("Third paragraph")

        # Get context at start (index 0)
        context = preprocessor._get_paragraph_context(doc, para_idx=0, window=2)

        assert "First paragraph" in context
        assert "Second paragraph" in context
        assert "Third paragraph" in context

    def test_get_paragraph_context_edge_case_end(self, preprocessor):
        """Test paragraph context at document end."""
        doc = Document()
        doc.add_paragraph("First paragraph")
        doc.add_paragraph("Second paragraph")
        doc.add_paragraph("Last paragraph")

        # Get context at end (index 2)
        context = preprocessor._get_paragraph_context(doc, para_idx=2, window=2)

        assert "First paragraph" in context
        assert "Second paragraph" in context
        assert "Last paragraph" in context

    def test_detect_image_type_png(self, preprocessor):
        """Test PNG image type detection."""
        png_data = b'\x89PNG\r\n\x1a\n'
        assert preprocessor._detect_image_type(png_data) == "png"

    def test_detect_image_type_jpeg(self, preprocessor):
        """Test JPEG image type detection."""
        jpeg_data = b'\xFF\xD8\xFF\xE0'
        assert preprocessor._detect_image_type(jpeg_data) == "jpeg"

    def test_detect_image_type_gif(self, preprocessor):
        """Test GIF image type detection."""
        gif_data = b'GIF89a'
        assert preprocessor._detect_image_type(gif_data) == "gif"

    def test_detect_image_type_bmp(self, preprocessor):
        """Test BMP image type detection."""
        bmp_data = b'BM'
        assert preprocessor._detect_image_type(bmp_data) == "bmp"

    def test_detect_image_type_unknown(self, preprocessor):
        """Test unknown image type defaults to PNG."""
        unknown_data = b'UNKNOWN'
        assert preprocessor._detect_image_type(unknown_data) == "png"

    def test_generate_image_captions_with_mock(self, preprocessor, mock_openai_client):
        """Test AI image caption generation with mock client."""
        images = [
            {
                'data': b'\x89PNG test data',
                'context': 'Test context around image',
                'para_idx': 0
            }
        ]

        captions = preprocessor._generate_image_captions(images)

        assert len(captions) == 1
        assert captions[0] == "Test image caption"
        assert mock_openai_client.chat.completions.create.called

    def test_generate_image_captions_error_handling(self, preprocessor):
        """Test caption generation error handling."""
        # Create preprocessor with client that raises error
        error_client = Mock()
        error_client.chat.completions.create.side_effect = Exception("API Error")
        preprocessor._client = error_client

        images = [
            {
                'data': b'\x89PNG test data',
                'context': 'Test context',
                'para_idx': 0
            }
        ]

        captions = preprocessor._generate_image_captions(images)

        assert len(captions) == 1
        assert "[Image: Caption unavailable]" in captions[0]

    def test_embed_image_captions(self, preprocessor):
        """Test inline embedding of image captions."""
        doc = Document()
        para = doc.add_paragraph("Text before image. ")

        images = [
            {
                'paragraph': para,
                'para_idx': 0
            }
        ]
        captions = ["A beautiful sunset over the ocean"]

        preprocessor._embed_image_captions(doc, images, captions)

        # Check that caption was added to paragraph
        para_text = para.text
        assert "[IMAGE: A beautiful sunset over the ocean]" in para_text

    def test_extract_shapes_empty_document(self, preprocessor):
        """Test shape extraction on document with no shapes."""
        doc = Document()
        doc.add_paragraph("Simple text, no shapes.")

        shapes = preprocessor._extract_shapes(doc)

        assert shapes == []

    def test_extract_textbox_text_no_textbox(self, preprocessor):
        """Test textbox text extraction on element without textbox."""
        doc = Document()
        para = doc.add_paragraph("Regular paragraph")

        # Element without textbox content
        text = preprocessor._extract_textbox_text(para._element)

        # Should return None or empty string
        assert text is None or text == ""

    def test_extract_comments_empty_document(self, preprocessor):
        """Test comment extraction on document with no comments."""
        doc = Document()
        doc.add_paragraph("Text without comments")

        comments = preprocessor._extract_comments(doc)

        assert comments == []

    def test_apply_comment_markers_empty_list(self, preprocessor):
        """Test applying comment markers with empty list."""
        doc = Document()
        initial_para_count = len(doc.paragraphs)

        preprocessor._apply_comment_markers(doc, [])

        # Should not add anything if no comments
        assert len(doc.paragraphs) == initial_para_count

    def test_apply_comment_markers_with_comments(self, preprocessor):
        """Test applying comment markers with actual comments."""
        doc = Document()
        initial_para_count = len(doc.paragraphs)

        comments = [
            {
                'id': '1',
                'author': 'John Doe',
                'text': 'This is a test comment',
                'date': '2025-12-01'
            },
            {
                'id': '2',
                'author': 'Jane Smith',
                'text': 'Another comment here',
                'date': '2025-12-01'
            }
        ]

        preprocessor._apply_comment_markers(doc, comments)

        # Should add:
        # 1 blank paragraph + 1 "DOCUMENT COMMENTS:" + 2 comment paragraphs = 4 new paragraphs
        assert len(doc.paragraphs) == initial_para_count + 4

        # Check comment markers were added
        doc_text = "\n".join([p.text for p in doc.paragraphs])
        assert "DOCUMENT COMMENTS:" in doc_text
        assert "[COMMENT_1]" in doc_text
        assert "John Doe" in doc_text
        assert "This is a test comment" in doc_text
        assert "[COMMENT_2]" in doc_text
        assert "Jane Smith" in doc_text
        assert "Another comment here" in doc_text

    def test_embed_shape_descriptions(self, preprocessor):
        """Test inline embedding of shape descriptions."""
        doc = Document()
        para = doc.add_paragraph("Text near shape. ")

        shapes = [
            {
                'type': 'textbox',
                'text': 'Important note in text box',
                'para_idx': 0,
                'paragraph': para
            }
        ]

        preprocessor._embed_shape_descriptions(doc, shapes)

        # Check that shape description was added to paragraph
        para_text = para.text
        assert "[SHAPE: Important note in text box]" in para_text

    @patch('src.extraction_v2.word_preprocessor.Document')
    def test_preprocess_full_workflow(self, mock_document_class, preprocessor, tmp_path):
        """Test full preprocessing workflow."""
        # Create a test file
        test_file = tmp_path / "test_doc.docx"
        test_file.touch()

        # Mock Document
        mock_doc = MagicMock()
        mock_doc.paragraphs = []
        mock_document_class.return_value = mock_doc

        # Mock the extraction methods to return empty lists
        preprocessor._extract_images = Mock(return_value=[])
        preprocessor._extract_shapes = Mock(return_value=[])
        preprocessor._extract_comments = Mock(return_value=[])

        # Run preprocessing
        output_path = preprocessor.preprocess(test_file)

        # Verify preprocessing steps were called
        assert preprocessor._extract_images.called
        assert preprocessor._extract_shapes.called
        assert preprocessor._extract_comments.called

        # Verify output path
        assert output_path is not None
        assert "_preprocessed" in str(output_path)

    def test_lazy_client_initialization(self):
        """Test that OpenAI client is initialized lazily."""
        preprocessor = WordPreprocessor(openai_client=None)
        assert preprocessor._client is None

        # Client should be created on first use
        with patch('openai.AzureOpenAI') as mock_azure:
            with patch('src.core.config.settings') as mock_settings:
                mock_settings.azure_openai_api_key = "test_key"
                mock_settings.azure_openai_api_version = "2024-02-15-preview"
                mock_settings.azure_openai_endpoint = "https://test.openai.azure.com/"

                client = preprocessor._get_client()

                assert client is not None
                assert mock_azure.called


class TestWordPreprocessorIntegration:
    """Integration tests for WordPreprocessor with real documents."""

    def test_preprocess_with_real_empty_document(self, tmp_path):
        """Test preprocessing on a real but empty document."""
        # Create a simple document
        doc = Document()
        doc.add_paragraph("This is a simple test document.")
        doc.add_paragraph("It has no images, shapes, or comments.")

        # Save it
        test_file = tmp_path / "empty_test.docx"
        doc.save(str(test_file))

        # Preprocess it
        mock_client = Mock()
        preprocessor = WordPreprocessor(openai_client=mock_client)

        output_path = preprocessor.preprocess(test_file)

        # Verify output exists
        assert output_path.exists()
        assert output_path.name.endswith("_preprocessed.docx")

        # Verify the preprocessed document can be opened
        preprocessed_doc = Document(str(output_path))
        assert len(preprocessed_doc.paragraphs) >= 2


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
