"""
Integration tests for V2 extraction pipeline with raw text handling.

Tests the full pipeline behavior when processing mixed content
(tables and raw text).
"""

import pytest
from unittest.mock import Mock, patch, MagicMock, AsyncMock
from uuid import uuid4

from src.extraction_v2.pipeline import ExtractionPipelineV2, ExtractionResult
from src.extraction_v2.html_table_extractor import HtmlTable, HtmlTableExtractor
from src.extraction_v2.record_builder import RecordBuilder
from src.extraction_v2.llm_header_detector import HeaderStructure


class TestPipelineRawTextIntegration:
    """Integration tests for pipeline raw text handling."""

    @pytest.fixture
    def mock_converter(self):
        """Create a mock converter."""
        converter = Mock()
        converter.convert_to_html.return_value = "<html></html>"
        converter.get_sheet_name_for_table.return_value = "TestSheet"
        return converter

    @pytest.fixture
    def mock_detector(self):
        """Create a mock LLM header detector."""
        detector = Mock()
        detector.detect_headers.return_value = HeaderStructure(
            header_rows=[0],
            headers=[["Col1", "Col2"]],
        )
        return detector

    @pytest.mark.asyncio
    async def test_pipeline_processes_mixed_content(
        self, mock_converter, mock_detector
    ):
        """Test pipeline correctly processes both valid tables and raw text."""
        # Create mixed HTML content
        mixed_html = """
        <html>
        <body>
            <table>
                <tr><td>Raw text content line 1</td></tr>
                <tr><td>Raw text content line 2</td></tr>
            </table>
            <table>
                <tr><th>Name</th><th>Value</th></tr>
                <tr><td>Item1</td><td>100</td></tr>
            </table>
        </body>
        </html>
        """
        mock_converter.convert_to_html.return_value = mixed_html

        pipeline = ExtractionPipelineV2(
            converter=mock_converter,
            detector=mock_detector,
        )

        document_id = uuid4()
        result = await pipeline.process_document(document_id, "test.xlsx")

        assert result.success is True
        # Should have records from both tables
        assert result.record_count > 0
        # Table count should include both types
        assert result.table_count == 2

        # Check that we have both types of records
        raw_text_records = [
            r for r in result.records if r.get("_source", {}).get("row") == -1
        ]
        table_records = [
            r for r in result.records if r.get("_source", {}).get("row") != -1
        ]

        assert len(raw_text_records) >= 1  # At least one raw text record
        assert raw_text_records[0]["content"]["text"] == "Raw text content line 1\nRaw text content line 2"

    @pytest.mark.asyncio
    async def test_pipeline_handles_only_raw_text(self, mock_converter, mock_detector):
        """Test pipeline when document contains only raw text tables."""
        raw_text_only_html = """
        <html>
        <body>
            <table>
                <tr><td>First paragraph text</td></tr>
                <tr><td>Second paragraph text</td></tr>
            </table>
            <table>
                <tr><td>Another raw text block</td></tr>
            </table>
        </body>
        </html>
        """
        mock_converter.convert_to_html.return_value = raw_text_only_html

        pipeline = ExtractionPipelineV2(
            converter=mock_converter,
            detector=mock_detector,
        )

        document_id = uuid4()
        result = await pipeline.process_document(document_id, "test.xlsx")

        assert result.success is True
        assert result.table_count == 2
        # All records should be raw text (row = -1)
        for record in result.records:
            assert record["_source"]["row"] == -1

        # LLM detector should NOT have been called (raw text skips LLM)
        mock_detector.detect_headers.assert_not_called()

    @pytest.mark.asyncio
    async def test_pipeline_handles_only_tables(self, mock_converter, mock_detector):
        """Test pipeline when document contains only valid data tables."""
        tables_only_html = """
        <html>
        <body>
            <table>
                <tr><th>A</th><th>B</th></tr>
                <tr><td>1</td><td>2</td></tr>
            </table>
            <table>
                <tr><th>X</th><th>Y</th><th>Z</th></tr>
                <tr><td>a</td><td>b</td><td>c</td></tr>
            </table>
        </body>
        </html>
        """
        mock_converter.convert_to_html.return_value = tables_only_html

        pipeline = ExtractionPipelineV2(
            converter=mock_converter,
            detector=mock_detector,
        )

        document_id = uuid4()
        result = await pipeline.process_document(document_id, "test.xlsx")

        assert result.success is True
        assert result.table_count == 2
        # No records should have row = -1 (no raw text)
        for record in result.records:
            assert record.get("_source", {}).get("row", 0) != -1

        # LLM detector should have been called for each table
        assert mock_detector.detect_headers.call_count == 2

    @pytest.mark.asyncio
    async def test_pipeline_no_llm_for_raw_text(self, mock_converter, mock_detector):
        """Test that LLM is NOT called for raw text tables (cost savings)."""
        # One raw text table
        html = """
        <table>
            <tr><td>Just some text content</td></tr>
        </table>
        """
        mock_converter.convert_to_html.return_value = html

        pipeline = ExtractionPipelineV2(
            converter=mock_converter,
            detector=mock_detector,
        )

        document_id = uuid4()
        result = await pipeline.process_document(document_id, "test.xlsx")

        assert result.success is True
        # LLM should NOT be called for raw text
        mock_detector.detect_headers.assert_not_called()

    @pytest.mark.asyncio
    async def test_pipeline_llm_only_for_valid_tables(
        self, mock_converter, mock_detector
    ):
        """Test that LLM is only called for valid data tables."""
        html = """
        <table><tr><td>Raw text 1</td></tr></table>
        <table><tr><td>A</td><td>B</td></tr></table>
        <table><tr><td>Raw text 2</td></tr></table>
        <table><tr><td>X</td><td>Y</td></tr></table>
        """
        mock_converter.convert_to_html.return_value = html

        pipeline = ExtractionPipelineV2(
            converter=mock_converter,
            detector=mock_detector,
        )

        document_id = uuid4()
        result = await pipeline.process_document(document_id, "test.xlsx")

        assert result.success is True
        # LLM should only be called for the 2 valid tables (not 4 total)
        assert mock_detector.detect_headers.call_count == 2

    @pytest.mark.asyncio
    async def test_pipeline_empty_document(self, mock_converter, mock_detector):
        """Test pipeline with no tables at all."""
        mock_converter.convert_to_html.return_value = "<html><body></body></html>"

        pipeline = ExtractionPipelineV2(
            converter=mock_converter,
            detector=mock_detector,
        )

        document_id = uuid4()
        result = await pipeline.process_document(document_id, "test.xlsx")

        assert result.success is True
        assert result.record_count == 0
        assert result.table_count == 0

    @pytest.mark.asyncio
    async def test_raw_text_record_structure(self, mock_converter, mock_detector):
        """Test the structure of raw text records."""
        html = """
        <table>
            <tr><td>Line 1</td></tr>
            <tr><td>Line 2</td></tr>
        </table>
        """
        mock_converter.convert_to_html.return_value = html

        pipeline = ExtractionPipelineV2(
            converter=mock_converter,
            detector=mock_detector,
        )

        document_id = uuid4()
        result = await pipeline.process_document(document_id, "test.xlsx")

        assert len(result.records) == 1
        record = result.records[0]

        # Check content structure
        assert "content" in record
        assert "text" in record["content"]
        assert record["content"]["text"] == "Line 1\nLine 2"

        # Check source metadata
        assert "_source" in record
        assert record["_source"]["document_id"] == str(document_id)
        assert record["_source"]["row"] == -1
        # Combined raw text uses table_ids (list) instead of table_id
        assert record["_source"]["table_ids"] == [0]


class TestPipelineBackwardCompatibility:
    """Test backward compatibility of extract_tables method."""

    def test_extract_tables_still_works(self):
        """Test that existing extract_tables method still returns only valid tables."""
        html = """
        <table><tr><td>Raw text</td></tr></table>
        <table><tr><td>A</td><td>B</td></tr></table>
        """

        extractor = HtmlTableExtractor()
        tables = extractor.extract_tables(html)

        # Should only return the valid table
        assert len(tables) == 1
        assert tables[0].col_count == 2


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