"""Tests for table inference processor (non-LLM functionality)."""

from typing import Optional
from unittest.mock import MagicMock, mock_open, patch

from quber.processors.table_inference import (
    ExecutiveSummary,
    TableContext,
    TableInferenceProcessor,
    TableInfo,
)


class TestTableContext:
    """Tests for TableContext model."""

    def test_table_context_creation(self):
        """Test creating a TableContext instance."""
        context = TableContext(
            executive_summary="Test summary",
            procedural_title="Test Table",
            headers=["Header1", "Header2"],
            descriptive_text="Table description",
            table_preview="| A | B |\n|---|---|\n| 1 | 2 |",
            preceding_text="Some text before",
            page_context="Page context",
            page_number=1,
        )

        assert context.executive_summary == "Test summary"
        assert context.procedural_title == "Test Table"
        assert context.headers == ["Header1", "Header2"]
        assert context.page_number == 1

    def test_table_context_optional_fields(self):
        """Test TableContext with optional fields as None."""
        context = TableContext(
            executive_summary="Summary",
            procedural_title="Title",
            headers=[],
            table_preview="Preview",
            page_number=1,
        )

        assert context.descriptive_text is None
        assert context.preceding_text is None
        assert context.page_context is None


class TestTableInfo:
    """Tests for TableInfo model."""

    def test_table_info_creation(self):
        """Test creating a TableInfo instance."""
        info = TableInfo(title="Financial Data", description="Quarterly revenue breakdown")

        assert info.title == "Financial Data"
        assert info.description == "Quarterly revenue breakdown"


class TestExecutiveSummary:
    """Tests for ExecutiveSummary model."""

    def test_executive_summary_creation(self):
        """Test creating an ExecutiveSummary instance."""
        summary = ExecutiveSummary(summary="This is a test summary of the document.")

        assert summary.summary == "This is a test summary of the document."


class TestTableInferenceProcessorInit:
    """Tests for TableInferenceProcessor initialization."""

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    def test_init_with_defaults(
        self, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging
    ):
        """Test initialization with default parameters."""
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_settings.log_level = "INFO"
        mock_get_settings.return_value = mock_settings

        mock_load_prompt.return_value = "System prompt"

        mock_factory = MagicMock()
        mock_agent_factory.return_value = mock_factory

        processor = TableInferenceProcessor()

        assert processor.settings == mock_settings
        assert processor.provider is None
        assert processor.model is None
        mock_setup_logging.assert_called_once()
        mock_load_prompt.assert_called_once_with("table_inference", "system")

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    def test_init_with_custom_parameters(self, mock_agent_factory, mock_load_prompt, mock_setup_logging):
        """Test initialization with custom parameters."""
        from quber.agents.factory import ModelProvider

        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = True
        mock_settings.log_level = "DEBUG"

        mock_load_prompt.return_value = "System prompt"
        mock_factory = MagicMock()
        mock_agent_factory.return_value = mock_factory

        processor = TableInferenceProcessor(
            provider=ModelProvider.ANTHROPIC, model="claude-3-5-sonnet-20241022", settings=mock_settings
        )

        assert processor.provider == ModelProvider.ANTHROPIC
        assert processor.model == "claude-3-5-sonnet-20241022"
        assert processor.settings == mock_settings


class TestTableInferenceProcessorHelpers:
    """Tests for helper methods."""

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    def test_format_prompt(self, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging):
        """Test format_prompt method."""
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_settings.log_level = "INFO"
        mock_get_settings.return_value = mock_settings
        mock_load_prompt.return_value = "System prompt"

        processor = TableInferenceProcessor()

        context = TableContext(
            executive_summary="Document summary",
            procedural_title="Table 1",
            headers=["Col1", "Col2"],
            table_preview="| A | B |\n| 1 | 2 |",
            page_number=1,
        )

        prompt = processor.format_prompt(context)

        assert "Document summary" in prompt
        assert "Table 1" in prompt
        assert "Col1, Col2" in prompt
        assert "page 1" in prompt
        assert "| A | B |" in prompt

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    def test_get_table_preview(
        self, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging
    ):
        """Test get_table_preview method."""
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_get_settings.return_value = mock_settings
        mock_load_prompt.return_value = "System prompt"

        processor = TableInferenceProcessor()

        markdown = "| A | B |\n| --- | --- |\n| 1 | 2 |\n| 3 | 4 |\n| 5 | 6 |"
        preview = processor.get_table_preview(markdown, max_rows=2)

        # Should include header, separator, and 2 rows
        lines = preview.split("\n")
        assert len(lines) <= 4  # header + separator + 2 rows

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    def test_get_table_preview_empty(
        self, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging
    ):
        """Test get_table_preview with empty markdown."""
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_get_settings.return_value = mock_settings
        mock_load_prompt.return_value = "System prompt"

        processor = TableInferenceProcessor()

        preview = processor.get_table_preview(None)
        assert preview == "No table content available"

        preview = processor.get_table_preview("")
        assert preview == "No table content available"


class TestTableInferenceProcessorSaveMethods:
    """Tests for save methods."""

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    @patch("builtins.open", new_callable=mock_open)
    def test_save_json(
        self, mock_file, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging
    ):
        """Test save_json method."""
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_get_settings.return_value = mock_settings
        mock_load_prompt.return_value = "System prompt"

        processor = TableInferenceProcessor()

        data = {"document": "test.pdf", "total_tables": 2}
        processor.save_json(data, "/output/test.json")

        mock_file.assert_called_once_with("/output/test.json", "w", encoding="utf-8")
        # Verify json.dump was called with the data
        handle = mock_file()
        written_content = "".join(call.args[0] for call in handle.write.call_args_list)
        # Should contain the data
        assert "test.pdf" in written_content or handle.write.called

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    @patch("builtins.open", new_callable=mock_open)
    def test_save_summary(
        self, mock_file, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging
    ):
        """Test save_summary method."""
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_get_settings.return_value = mock_settings
        mock_load_prompt.return_value = "System prompt"

        processor = TableInferenceProcessor()

        data = {
            "document": "test.pdf",
            "extraction_date": "2024-01-01",
            "model_provider": "anthropic",
            "model": "claude-3-5-sonnet-20241022",
            "total_tables": 1,
            "tables": [
                {
                    "table_id": 0,
                    "page": 1,
                    "procedural_title": "Table 1",
                    "headers": ["Col1", "Col2"],
                    "descriptive_text": "Test description",
                    "llm_title": "Financial Data",
                    "llm_description": "Quarterly revenue",
                    "table_markdown": "| A | B |\n| 1 | 2 |",
                    "metadata": {"rows": 2, "cols": 2, "first_cell": "A"},
                }
            ],
        }

        processor.save_summary(data, "/output/test_summary.md")

        mock_file.assert_called_once_with("/output/test_summary.md", "w", encoding="utf-8")
        handle = mock_file()
        # Verify write was called
        assert handle.write.called

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    @patch("builtins.open", new_callable=mock_open)
    def test_save_summary_content(
        self, mock_file, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging
    ):
        """Test save_summary generates correct content."""
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_get_settings.return_value = mock_settings
        mock_load_prompt.return_value = "System prompt"

        processor = TableInferenceProcessor()

        data = {
            "document": "test.pdf",
            "extraction_date": "2024-01-01",
            "model_provider": "anthropic",
            "model": "claude-3-5-sonnet-20241022",
            "total_tables": 1,
            "tables": [
                {
                    "table_id": 0,
                    "page": 1,
                    "procedural_title": "Table 1",
                    "llm_title": "Financial Data",
                    "llm_description": "Revenue breakdown",
                    "table_markdown": "| Q1 | Q2 |\n| 100 | 200 |",
                    "metadata": {"rows": 2, "cols": 2},
                }
            ],
        }

        # Capture what gets written
        written_lines = []

        def write_side_effect(content):
            written_lines.append(content)

        mock_file().write.side_effect = write_side_effect

        processor.save_summary(data, "/output/summary.md")

        written_content = "".join(written_lines)

        # Verify key content is included
        assert "test.pdf" in written_content
        assert "Table 1" in written_content or "Financial Data" in written_content


class FakeProv:
    """Stand-in for a docling ProvenanceItem: just carries a 1-indexed page_no."""

    def __init__(self, page_no: int):
        self.page_no = page_no


class FakeItem:
    """Stand-in for a docling node item as yielded by iterate_items().

    Mirrors the real contract the code depends on: the page lives in a `prov`
    list (not a `_page` attribute), text items have `text`, and header items
    additionally have a `level`. iterate_items() yields (item, level) tuples,
    so tests wrap these in tuples."""

    def __init__(self, page: int, text: str = "", level: Optional[int] = None):
        self.prov = [FakeProv(page)]
        self.text = text
        if level is not None:
            self.level = level


def _items(*items: FakeItem):
    """Wrap items as the (item, tree-level) tuples iterate_items() yields."""
    return [(item, 0) for item in items]


class TestExtractPageContext:
    """Tests for extract_page_context method."""

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    def test_extract_page_context(
        self, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging
    ):
        """Page text and headers on the table's page are collected; other pages excluded.

        Drives iterate_items() with the real contract (tuples of (item, level);
        page read from item.prov[*].page_no), so it fails against the prior code
        that looked for a non-existent `_page` attribute on the tuple.
        """
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_get_settings.return_value = mock_settings
        mock_load_prompt.return_value = "System prompt"

        processor = TableInferenceProcessor()

        mock_document = MagicMock()
        mock_document.iterate_items.return_value = _items(
            FakeItem(page=1, text="Header text", level=1),
            FakeItem(page=1, text="Body text"),
            FakeItem(page=2, text="Other page text"),
        )

        context = processor.extract_page_context(mock_document, page_number=1)

        assert "Header text" in context
        assert "Body text" in context
        assert "[Header L1]" in context  # header labeled with its level
        assert "Other page text" not in context  # different page is excluded

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    def test_extract_page_context_includes_previous_page_tail(
        self, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging
    ):
        """When the table's page is sparse, a tail of the previous page is added."""
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_get_settings.return_value = mock_settings
        mock_load_prompt.return_value = "System prompt"

        processor = TableInferenceProcessor()

        mock_document = MagicMock()
        mock_document.iterate_items.return_value = _items(
            FakeItem(page=1, text="Intro on the prior page"),
            FakeItem(page=2, text="Short"),
        )

        context = processor.extract_page_context(mock_document, page_number=2)

        assert "=== Previous Page Context ===" in context
        assert "Intro on the prior page" in context
        assert "Short" in context

    @patch("quber.processors.table_inference.setup_logging")
    @patch("quber.processors.table_inference.load_prompt")
    @patch("quber.processors.table_inference.AgentFactory")
    @patch("quber.processors.table_inference.get_settings")
    def test_extract_page_context_no_content(
        self, mock_get_settings, mock_agent_factory, mock_load_prompt, mock_setup_logging
    ):
        """Test extract_page_context with no content."""
        mock_settings = MagicMock()
        mock_settings.obs.enable_logfire = False
        mock_get_settings.return_value = mock_settings
        mock_load_prompt.return_value = "System prompt"

        processor = TableInferenceProcessor()

        mock_document = MagicMock()
        mock_document.iterate_items.return_value = []

        context = processor.extract_page_context(mock_document, page_number=1)

        assert context == "No page context available"