"""
Unit tests for document_extraction_dag.

These tests verify DAG structure, task dependencies, and branching logic
without requiring the full Airflow runtime.
"""

import pytest
from unittest.mock import Mock


class TestDetectFormatFunction:
    """Test the detect_format branching function logic."""

    def _detect_format(self, file_path: str) -> list:
        """
        Replicate the detect_format logic for testing.

        This mirrors the logic in the DAG file without needing to import it.
        Returns a list of task_ids (Excel returns 2 for parallel pipelines).
        """
        ext = file_path.lower().rsplit('.', 1)[-1] if '.' in file_path else ''

        if ext in ('xlsx', 'xls', 'xlsm'):
            # Excel runs BOTH pipelines in parallel
            return ['parse_excel_large_tables', 'parse_excel_detailed']
        elif ext in ('docx', 'doc', 'docm'):
            return ['parse_word']
        elif ext == 'pdf':
            return ['parse_pdf']
        elif ext in ('pptx', 'ppt', 'pptm'):
            return ['parse_powerpoint']
        else:
            return ['unsupported_format']

    def test_detect_excel_xlsx(self):
        """Test detection of xlsx files returns both Excel pipelines."""
        result = self._detect_format('docs/report.xlsx')
        assert result == ['parse_excel_large_tables', 'parse_excel_detailed']

    def test_detect_excel_xls(self):
        """Test detection of xls files."""
        result = self._detect_format('docs/report.xls')
        assert result == ['parse_excel_large_tables', 'parse_excel_detailed']

    def test_detect_excel_xlsm(self):
        """Test detection of xlsm files."""
        result = self._detect_format('docs/report.xlsm')
        assert result == ['parse_excel_large_tables', 'parse_excel_detailed']

    def test_detect_word_docx(self):
        """Test detection of docx files."""
        result = self._detect_format('docs/document.docx')
        assert result == ['parse_word']

    def test_detect_word_doc(self):
        """Test detection of doc files."""
        result = self._detect_format('docs/document.doc')
        assert result == ['parse_word']

    def test_detect_word_docm(self):
        """Test detection of docm files."""
        result = self._detect_format('docs/document.docm')
        assert result == ['parse_word']

    def test_detect_pdf(self):
        """Test detection of pdf files."""
        result = self._detect_format('docs/document.pdf')
        assert result == ['parse_pdf']

    def test_detect_powerpoint_pptx(self):
        """Test detection of pptx files."""
        result = self._detect_format('docs/presentation.pptx')
        assert result == ['parse_powerpoint']

    def test_detect_powerpoint_ppt(self):
        """Test detection of ppt files."""
        result = self._detect_format('docs/presentation.ppt')
        assert result == ['parse_powerpoint']

    def test_detect_powerpoint_pptm(self):
        """Test detection of pptm files."""
        result = self._detect_format('docs/presentation.pptm')
        assert result == ['parse_powerpoint']

    def test_detect_unsupported_format(self):
        """Test detection of unsupported file types."""
        result = self._detect_format('docs/image.png')
        assert result == ['unsupported_format']

    def test_detect_no_extension(self):
        """Test handling of files without extension."""
        result = self._detect_format('docs/noextension')
        assert result == ['unsupported_format']

    def test_detect_case_insensitive(self):
        """Test that extension detection is case-insensitive."""
        result = self._detect_format('docs/REPORT.XLSX')
        assert result == ['parse_excel_large_tables', 'parse_excel_detailed']

    def test_detect_handles_path_with_multiple_dots(self):
        """Test handling of filenames with multiple dots."""
        result = self._detect_format('docs/report.v2.xlsx')
        assert result == ['parse_excel_large_tables', 'parse_excel_detailed']

    def test_detect_txt_unsupported(self):
        """Test that txt files are unsupported."""
        result = self._detect_format('docs/readme.txt')
        assert result == ['unsupported_format']

    def test_detect_csv_unsupported(self):
        """Test that csv files are unsupported."""
        result = self._detect_format('docs/data.csv')
        assert result == ['unsupported_format']

    def test_detect_empty_path(self):
        """Test handling of empty file path."""
        result = self._detect_format('')
        assert result == ['unsupported_format']


class TestDAGStructure:
    """Test DAG structure by parsing the DAG file."""

    @pytest.fixture
    def dag_content(self):
        """Read DAG file content."""
        with open('airflow/dags/document_extraction_dag.py', 'r') as f:
            return f.read()

    def test_expected_task_ids(self, dag_content):
        """Verify all expected tasks are defined."""
        expected_tasks = [
            'validate_event',
            'fetch_document',
            'save_document_metadata',
            'unsupported_format',
            # Excel Large Table pipeline
            'parse_excel_large_tables',
            'detect_headers',
            'chunk_large_table_records',
            'convert_large_tables_to_nl',
            'extract_domain_terms_excel',
            'tag_records_excel',
            'store_large_tables_llamaindex',
            # Excel Detailed pipeline (hierarchical)
            'parse_excel_detailed',
            'upload_images_to_seaweedfs',
            'hierarchical_chunk_excel_detailed',
            'extract_domain_terms_excel_detailed',
            'tag_records_excel_detailed',
            'store_excel_detailed_llamaindex',
            # Text pipeline (hierarchical)
            'parse_word',
            'parse_pdf',
            'parse_powerpoint',
            'hierarchical_chunk_text',
            'extract_domain_terms_text',
            'tag_records_text',
            'store_text_llamaindex',
            # Final
            'update_document_status',
        ]

        for task_id in expected_tasks:
            assert f"task_id='{task_id}'" in dag_content, f"Missing task: {task_id}"

    def test_dag_id(self, dag_content):
        """Verify DAG ID."""
        assert "dag_id='document_extraction_dag'" in dag_content

    def test_dag_schedule(self, dag_content):
        """Verify DAG is event-triggered only."""
        assert "schedule=None" in dag_content

    def test_dag_catchup_disabled(self, dag_content):
        """Verify catchup is disabled."""
        assert "catchup=False" in dag_content

    def test_dag_max_active_runs(self, dag_content):
        """Verify max_active_runs is set."""
        assert "max_active_runs=10" in dag_content

    def test_hierarchical_tag(self, dag_content):
        """Verify hierarchical tag is present."""
        assert "'hierarchical'" in dag_content

    def test_storage_tasks_have_timeout(self, dag_content):
        """Verify storage tasks have execution timeout."""
        assert "execution_timeout=CONFIG.STORE_TIMEOUT" in dag_content

    def test_parse_tasks_have_timeout(self, dag_content):
        """Verify parse tasks have execution timeout."""
        assert "execution_timeout=CONFIG.PARSE_TIMEOUT" in dag_content


class TestDAGDependencies:
    """Test DAG task dependencies."""

    @pytest.fixture
    def dag_content(self):
        """Read DAG file content."""
        with open('airflow/dags/document_extraction_dag.py', 'r') as f:
            return f.read()

    def test_initial_flow_dependency(self, dag_content):
        """Verify validate >> fetch >> save_doc_metadata >> branch dependency chain."""
        assert "validate >> fetch >> save_doc_metadata >> branch" in dag_content

    def test_unsupported_format_branch(self, dag_content):
        """Verify unsupported format branch exists."""
        assert "branch >> unsupported" in dag_content

    def test_excel_large_table_pipeline_dependency(self, dag_content):
        """Verify Excel Large Table pipeline dependencies."""
        assert "branch >> parse_excel_large >> detect_headers >> chunk_large_tables >> nl_chunks >> domain_terms_excel >> tags_excel >> store_large_tables" in dag_content

    def test_excel_detailed_pipeline_dependency(self, dag_content):
        """Verify Excel Detailed pipeline dependencies with hierarchical chunking."""
        assert "branch >> parse_excel_det >> upload_images >> hierarchical_chunk_excel >> domain_terms_excel_det >> tags_excel_det >> store_excel_detailed" in dag_content

    def test_word_path_dependency(self, dag_content):
        """Verify Word path goes to text_parse_result."""
        assert "branch >> parse_word >> text_parse_result" in dag_content

    def test_pdf_path_dependency(self, dag_content):
        """Verify PDF path goes to text_parse_result."""
        assert "branch >> parse_pdf >> text_parse_result" in dag_content

    def test_powerpoint_path_dependency(self, dag_content):
        """Verify PowerPoint path goes to text_parse_result."""
        assert "branch >> parse_powerpoint >> text_parse_result" in dag_content

    def test_text_pipeline_dependency(self, dag_content):
        """Verify text pipeline flow with hierarchical chunking."""
        assert "text_parse_result >> hierarchical_chunk_text >> domain_terms_text >> tags_text >> store_text" in dag_content

    def test_final_status_update_dependency(self, dag_content):
        """Verify status update runs after all storage paths."""
        assert "[store_large_tables, store_excel_detailed, store_text] >> update_doc_status" in dag_content

    def test_cleanup_dependency(self, dag_content):
        """Verify cleanup runs after completion."""
        assert "[update_doc_status, unsupported] >> cleanup" in dag_content


class TestDAGDocumentation:
    """Test DAG documentation."""

    @pytest.fixture
    def dag_content(self):
        """Read DAG file content."""
        with open('airflow/dags/document_extraction_dag.py', 'r') as f:
            return f.read()

    def test_dag_has_docmd(self, dag_content):
        """Verify DAG has doc_md documentation."""
        assert "doc_md=" in dag_content
        assert "## Document Extraction DAG" in dag_content

    def test_dag_has_description(self, dag_content):
        """Verify DAG has description."""
        assert "description='Multi-format document extraction pipeline'" in dag_content

    def test_dag_has_tags(self, dag_content):
        """Verify DAG has tags for filtering."""
        assert "'extraction'" in dag_content
        assert "'hierarchical'" in dag_content

    def test_dag_doc_mentions_trigger_format(self, dag_content):
        """Verify DAG doc explains trigger format."""
        assert '"document_id"' in dag_content
        assert '"object_key"' in dag_content

    def test_dag_doc_mentions_hierarchical_chunking(self, dag_content):
        """Verify DAG doc explains hierarchical chunking."""
        assert "Hierarchical Chunking" in dag_content or "hierarchical chunking" in dag_content
        assert "HierarchicalNodeParser" in dag_content

    def test_dag_doc_mentions_llamaindex(self, dag_content):
        """Verify DAG doc mentions LlamaIndex storage."""
        assert "LlamaIndex" in dag_content
        assert "MilvusVectorStore" in dag_content
        assert "PostgresDocumentStore" in dag_content


class TestDAGImports:
    """Test DAG imports are correct."""

    @pytest.fixture
    def dag_content(self):
        """Read DAG file content."""
        with open('airflow/dags/document_extraction_dag.py', 'r') as f:
            return f.read()

    def test_imports_operators(self, dag_content):
        """Verify operators are imported."""
        operators = [
            'ValidateEventOperator',
            'FetchDocumentOperator',
        ]
        for op in operators:
            assert op in dag_content, f"Missing import: {op}"

    def test_imports_parse_tasks(self, dag_content):
        """Verify parse tasks are imported."""
        tasks = [
            'parse_excel_large_tables',
            'parse_word_document',
            'parse_pdf_document',
            'parse_powerpoint_document',
            'parse_excel_detailed',
        ]
        for task in tasks:
            assert task in dag_content, f"Missing import: {task}"

    def test_imports_hierarchical_tasks(self, dag_content):
        """Verify hierarchical chunking tasks are imported."""
        tasks = [
            'hierarchical_chunk_text_task',
            'hierarchical_chunk_excel_detailed_task',
            'extract_domain_terms_nodes_task',
            'tag_records_nodes_task',
            'store_to_llamaindex_task',
            'store_large_table_to_llamaindex_task',
        ]
        for task in tasks:
            assert task in dag_content, f"Missing import: {task}"

    def test_imports_flat_record_tasks(self, dag_content):
        """Verify flat record tasks are imported for Large Table pipeline."""
        tasks = [
            'detect_headers_task',
            'chunk_large_table_records',
            'convert_large_tables_to_natural_language',
            'extract_domain_terms_task',
            'tag_records_task',
        ]
        for task in tasks:
            assert task in dag_content, f"Missing import: {task}"

    def test_imports_empty_operator(self, dag_content):
        """Verify EmptyOperator is imported."""
        assert "EmptyOperator" in dag_content

    def test_imports_config(self, dag_content):
        """Verify config imports."""
        assert "from config import CONFIG, default_args" in dag_content


class TestDetectFormatSync:
    """Verify test helper stays in sync with actual DAG implementation."""

    @pytest.fixture
    def dag_content(self):
        """Read DAG file content."""
        with open('airflow/dags/document_extraction_dag.py', 'r') as f:
            return f.read()

    def test_excel_extensions_match_dag(self, dag_content):
        """Verify Excel extensions in test match DAG."""
        test_extensions = ('xlsx', 'xls', 'xlsm')
        for ext in test_extensions:
            assert f"'{ext}'" in dag_content, f"Extension {ext} not in DAG"

    def test_word_extensions_match_dag(self, dag_content):
        """Verify Word extensions in test match DAG."""
        test_extensions = ('docx', 'doc', 'docm')
        for ext in test_extensions:
            assert f"'{ext}'" in dag_content, f"Extension {ext} not in DAG"

    def test_powerpoint_extensions_match_dag(self, dag_content):
        """Verify PowerPoint extensions in test match DAG."""
        test_extensions = ('pptx', 'ppt', 'pptm')
        for ext in test_extensions:
            assert f"'{ext}'" in dag_content, f"Extension {ext} not in DAG"

    def test_pdf_extension_in_dag(self, dag_content):
        """Verify PDF extension is in DAG."""
        assert "'pdf'" in dag_content

    def test_excel_returns_both_pipelines(self, dag_content):
        """Verify Excel branch returns both pipeline task IDs."""
        assert "return ['parse_excel_large_tables', 'parse_excel_detailed']" in dag_content


class TestHierarchicalChunkingTasks:
    """Test hierarchical chunking tasks are available."""

    @pytest.fixture
    def processing_tasks_content(self):
        """Read processing_tasks.py content."""
        with open('airflow/dags/tasks/processing_tasks.py', 'r') as f:
            return f.read()

    def test_hierarchical_chunk_text_task_exists(self, processing_tasks_content):
        """Verify hierarchical_chunk_text_task is defined."""
        assert "def hierarchical_chunk_text_task(" in processing_tasks_content

    def test_hierarchical_chunk_excel_detailed_task_exists(self, processing_tasks_content):
        """Verify hierarchical_chunk_excel_detailed_task is defined."""
        assert "def hierarchical_chunk_excel_detailed_task(" in processing_tasks_content

    def test_extract_domain_terms_nodes_task_exists(self, processing_tasks_content):
        """Verify extract_domain_terms_nodes_task is defined."""
        assert "def extract_domain_terms_nodes_task(" in processing_tasks_content

    def test_tag_records_nodes_task_exists(self, processing_tasks_content):
        """Verify tag_records_nodes_task is defined."""
        assert "def tag_records_nodes_task(" in processing_tasks_content

    def test_store_to_llamaindex_task_exists(self, processing_tasks_content):
        """Verify store_to_llamaindex_task is defined."""
        assert "def store_to_llamaindex_task(" in processing_tasks_content

    def test_store_large_table_to_llamaindex_task_exists(self, processing_tasks_content):
        """Verify store_large_table_to_llamaindex_task is defined."""
        assert "def store_large_table_to_llamaindex_task(" in processing_tasks_content

    def test_hierarchical_chunk_uses_llamaindex(self, processing_tasks_content):
        """Verify hierarchical chunking uses LlamaIndex."""
        assert "from llama_index.core import Document" in processing_tasks_content
        assert "HierarchicalNodeParser" in processing_tasks_content

    def test_store_task_uses_milvus_vector_store(self, processing_tasks_content):
        """Verify store task uses LlamaIndex MilvusVectorStore."""
        assert "from llama_index.vector_stores.milvus import MilvusVectorStore" in processing_tasks_content

    def test_store_task_uses_postgres_docstore(self, processing_tasks_content):
        """Verify store task uses LlamaIndex PostgresDocumentStore."""
        assert "from llama_index.storage.docstore.postgres import PostgresDocumentStore" in processing_tasks_content

    def test_hierarchical_chunk_sizes(self, processing_tasks_content):
        """Verify default chunk sizes are [2048, 512, 256]."""
        assert "chunk_sizes = [2048, 512, 256]" in processing_tasks_content

    def test_uses_tiktoken_tokenizer(self, processing_tasks_content):
        """Verify tiktoken is used for tokenization."""
        assert "import tiktoken" in processing_tasks_content
        assert 'tiktoken.get_encoding("cl100k_base")' in processing_tasks_content


class TestHierarchicalStorageArchitecture:
    """Test hierarchical storage architecture is correctly implemented."""

    @pytest.fixture
    def processing_tasks_content(self):
        """Read processing_tasks.py content."""
        with open('airflow/dags/tasks/processing_tasks.py', 'r') as f:
            return f.read()

    def test_all_nodes_to_docstore(self, processing_tasks_content):
        """Verify all nodes are added to docstore."""
        assert "docstore.add_documents(nodes)" in processing_tasks_content

    def test_leaf_nodes_to_milvus(self, processing_tasks_content):
        """Verify only leaf nodes are indexed to Milvus."""
        assert "get_leaf_nodes" in processing_tasks_content
        assert "VectorStoreIndex" in processing_tasks_content

    def test_storage_context_setup(self, processing_tasks_content):
        """Verify StorageContext is used with both stores."""
        assert "StorageContext.from_defaults" in processing_tasks_content
        assert "vector_store=vector_store" in processing_tasks_content
        assert "docstore=docstore" in processing_tasks_content

    def test_azure_embedding_configuration(self, processing_tasks_content):
        """Verify Azure OpenAI embedding is configured."""
        assert "AzureOpenAIEmbedding" in processing_tasks_content
        assert "text-embedding-3-large" in processing_tasks_content
        assert "dim=3072" in processing_tasks_content