"""
Processing tasks using @task.external_python decorator.

These tasks run in the uv virtual environment at /opt/airflow/.venv
which has all the required dependencies.
"""

from airflow.decorators import task
from typing import Dict, Any, List, Optional

# Path to the uv virtual environment Python
EXTERNAL_PYTHON = '/opt/airflow/.venv/bin/python'


@task.external_python(python=EXTERNAL_PYTHON)
def detect_headers_task(parse_result: Dict[str, Any]) -> Dict[str, Any]:
    """
    Detect table headers using LLM (Excel documents only).

    Args:
        parse_result: Result from parse_excel task (contains data_file path)

    Returns:
        DetectedHeaders dict with header info for each table
    """
    from src.extraction_v2.llm_header_detector import LlmHeaderDetector
    import logging
    import pickle
    import os

    log = logging.getLogger(__name__)

    # Skip if no parse result
    if not parse_result:
        log.info("No parse result provided")
        return {'tables_processed': 0, 'header_data': []}

    # Load data from pickle file if data_file path is provided
    data_file = parse_result.get('data_file')
    if data_file and os.path.exists(data_file):
        log.info(f"Loading parse data from {data_file}")
        with open(data_file, 'rb') as f:
            parse_result = pickle.load(f)

    if parse_result.get('doc_type') != 'excel':
        log.info(f"Not an Excel document (type: {parse_result.get('doc_type')})")
        return {'tables_processed': 0, 'header_data': []}

    # Get tables from parse result
    tables = parse_result.get('tables', [])
    if not tables:
        log.info("No tables found in parsed document")
        return {'tables_processed': 0, 'header_data': []}

    log.info(f"Detecting headers for {len(tables)} tables")

    # Initialize detector
    detector = LlmHeaderDetector()

    # Detect headers for each table
    all_headers = []
    for i, table in enumerate(tables):
        try:
            table_html = table.get('html', '')
            if not table_html:
                log.warning(f"Table {i} has no HTML content, using fallback")
                all_headers.append({
                    'table_index': i,
                    'header_rows': [0],
                    'headers': [],
                })
                continue

            log.info(f"Detecting headers for table {i}")
            header_info = detector.detect_headers(table_html)

            all_headers.append({
                'table_index': i,
                'header_rows': header_info.header_rows,
                'headers': header_info.headers,
            })

            log.info(
                f"Table {i}: {len(header_info.header_rows)} header row(s), "
                f"{len(header_info.headers)} column(s)"
            )

        except Exception as e:
            log.warning(f"Header detection failed for table {i}: {e}. Using fallback.")
            all_headers.append({
                'table_index': i,
                'header_rows': [0],
                'headers': [],
            })

    log.info(f"Header detection complete: {len(all_headers)} tables processed")

    return {
        'tables_processed': len(all_headers),
        'header_data': all_headers,
    }


@task.external_python(python=EXTERNAL_PYTHON)
def chunk_large_table_records(
    document_id: str,
    parse_result: Dict[str, Any],
    headers_result: Dict[str, Any],
    tenant_id: str,
    node_id: str,
    correlation_id: str,
) -> Dict[str, Any]:
    """
    Chunk large Excel tables into records using RecordBuilder.

    This task processes LARGE TABLES (>1000 rows) from the NEW pipeline.

    Args:
        document_id: Document UUID string
        parse_result: Result from parse_excel_large_tables task (contains data_file path)
        headers_result: Result from detect_headers task
        tenant_id: Tenant UUID for pipeline-event routing
        node_id: Node UUID for pipeline-event routing
        correlation_id: Airflow dag_run_id, for event correlation

    Returns:
        ContentChunks dict with records
    """
    from uuid import UUID
    from src.extraction_v2.record_builder import RecordBuilder
    from src.extraction_v2.llm_header_detector import HeaderStructure
    from src.extraction_v2.html_table_extractor import HtmlTable
    from src.services.pipeline_events import publish_stage_event
    import logging
    import pickle
    import os

    log = logging.getLogger(__name__)

    publish_stage_event(
        "chunking", "started",
        document_id=document_id,
        tenant_id=tenant_id,
        node_id=node_id,
        correlation_id=correlation_id,
    )

    doc_uuid = UUID(document_id)

    if not parse_result:
        log.warning("No parse result provided")
        return {'chunks': [], 'chunk_count': 0, 'chunking_strategy': 'table'}

    # Load data from pickle file if data_file path is provided
    data_file = parse_result.get('data_file')
    if data_file and os.path.exists(data_file):
        log.info(f"Loading parse data from {data_file}")
        with open(data_file, 'rb') as f:
            parse_result = pickle.load(f)
        # Clean up the temp file after loading
        try:
            os.remove(data_file)
            log.info(f"Cleaned up parse data file: {data_file}")
        except Exception as e:
            log.warning(f"Failed to clean up parse data file {data_file}: {e}")

    # Get filename from parse result
    original_filename = parse_result.get('original_filename', '')

    # Get header data
    header_data = headers_result.get('header_data', []) if headers_result else []

    # Get tables from parse result (NEW pipeline only processes large tables)
    large_table_records = parse_result.get('large_table_records', [])

    if not large_table_records:
        log.info("No large table records found (NEW pipeline only processes large tables)")
        return {'chunks': [], 'chunk_count': 0, 'chunking_strategy': 'table'}

    log.info(
        f"Chunking {len(large_table_records)} pre-processed large table records"
    )

    builder = RecordBuilder()
    all_records = []

    # NEW PIPELINE: Only process large table records
    # Normal tables, raw text tables, images, shapes, and comments are handled by DETAILED pipeline
    log.info(f"Including {len(large_table_records)} pre-processed large table records")
    all_records.extend(large_table_records)

    log.info(f"Total: {len(all_records)} records from Excel chunking")

    # Add filename and col_range to each record's _source (required by ExtractedRecord)
    for record in all_records:
        if isinstance(record, dict) and '_source' in record:
            record['_source']['filename'] = original_filename
            if 'col_range' not in record['_source']:
                record['_source']['col_range'] = ''

    # Save large data to pickle file to avoid XCom size limits
    chunk_data = {
        'chunks': all_records,
        'chunk_count': len(all_records),
        'chunking_strategy': 'table',
        'doc_type': 'excel',
        'original_filename': original_filename,
    }

    data_file = f"/opt/airflow/logs/chunks_excel_{document_id}.pkl"
    with open(data_file, 'wb') as f:
        pickle.dump(chunk_data, f)

    log.info(f"Saved chunk result to {data_file}")

    # Return only small metadata via XCom
    return {
        'data_file': data_file,
        'chunk_count': len(all_records),
        'chunking_strategy': 'table',
        'doc_type': 'excel',
        'original_filename': original_filename,
        'document_id': document_id,
    }


@task.external_python(python=EXTERNAL_PYTHON)
def convert_large_tables_to_natural_language(
    chunks_result: Dict[str, Any]
) -> Dict[str, Any]:
    """
    Convert large table records to natural language using rule-based formatting.

    Takes flat record dictionaries from chunk_excel_tables and converts them
    to natural language descriptions with proper structure for downstream tasks.

    Args:
        chunks_result: Result from chunk_excel_tables (contains data_file path)

    Returns:
        Dict with data_file path containing structured records with natural language
    """
    import logging
    import pickle
    import os
    from uuid import uuid4

    log = logging.getLogger(__name__)

    if not chunks_result:
        log.warning("No chunks result provided")
        return {'chunks': [], 'chunk_count': 0, 'doc_type': 'excel'}

    # Load data from pickle file
    data_file = chunks_result.get('data_file')
    if not data_file or not os.path.exists(data_file):
        log.warning("No data file found")
        return chunks_result

    log.info(f"Loading chunk data from {data_file}")
    with open(data_file, 'rb') as f:
        chunk_data = pickle.load(f)

    # Clean up input file
    try:
        os.remove(data_file)
        log.info(f"Cleaned up input file: {data_file}")
    except Exception as e:
        log.warning(f"Failed to clean up input file {data_file}: {e}")

    chunks = chunk_data.get('chunks', [])
    if not chunks:
        log.info("No chunks to convert")
        return chunks_result

    log.info(f"Converting {len(chunks)} large table records to natural language")

    # Get original filename for metadata
    original_filename = chunk_data.get('original_filename', '')

    from src.extraction_v2.record_to_natural_language import RecordToNaturalLanguage

    converter = RecordToNaturalLanguage()

    # Prepare batch data
    records = []
    sheet_names = []
    row_numbers = []

    for idx, chunk in enumerate(chunks):
        if isinstance(chunk, dict):
            sheet_name = chunk.get('_sheet_name', 'Table')
            row_number = chunk.get('_row_number', idx + 1)

            records.append(chunk)
            sheet_names.append(sheet_name)
            row_numbers.append(row_number)

    # Convert to natural language (rule-based, no LLM)
    natural_language_texts = converter.convert_batch(
        records, sheet_names, row_numbers
    )

    # Group 3-4 records into one chunk, respecting 8k token limit for embedding
    MAX_TOKENS = 8000
    TARGET_GROUP_SIZE = 4
    from src.llm.factory import get_token_counter
    token_len = get_token_counter()

    from langchain_text_splitters import RecursiveCharacterTextSplitter
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=MAX_TOKENS,
        chunk_overlap=200,
        length_function=token_len,
        separators=["\n\n", "\n", " ", ""],
    )

    def make_chunk(text, sheet, row):
        return {
            'content': {'text': text},
            'text': text,
            '_source': {
                'filename': original_filename,
                'sheet': sheet,
                'row': row,
            },
            'headers': ['text'],
        }

    structured_chunks = []
    i = 0
    while i < len(natural_language_texts):
        group_texts = []
        group_rows = []
        group_sheet = sheet_names[i]

        while i < len(natural_language_texts) and len(group_texts) < TARGET_GROUP_SIZE:
            nl_text = natural_language_texts[i]
            if not nl_text:
                i += 1
                continue

            # Check token count of the joined result (including separators)
            candidate = "\n\n".join(group_texts + [nl_text])
            if token_len(candidate) > MAX_TOKENS and group_texts:
                break

            group_texts.append(nl_text)
            group_rows.append(row_numbers[i])
            i += 1

        if not group_texts:
            continue

        grouped_text = "\n\n".join(group_texts)

        # If a single record exceeds the limit, split it into multiple chunks
        if token_len(grouped_text) > MAX_TOKENS:
            splits = text_splitter.split_text(grouped_text)
            for split_text in splits:
                structured_chunks.append(make_chunk(split_text, group_sheet, group_rows[0]))
        else:
            structured_chunks.append(make_chunk(grouped_text, group_sheet, group_rows[0]))

    log.info(
        f"Grouped {len(natural_language_texts)} records into {len(structured_chunks)} chunks "
        f"(target {TARGET_GROUP_SIZE} records/chunk, max {MAX_TOKENS} tokens)"
    )

    # Save to new pickle file
    output_data = {
        'chunks': structured_chunks,
        'chunk_count': len(structured_chunks),
        'chunking_strategy': chunk_data.get('chunking_strategy', 'table'),
        'doc_type': 'excel',
        'original_filename': chunk_data.get('original_filename', ''),
    }

    nl_doc_id = chunks_result.get('document_id', uuid4().hex)
    output_file = f"/opt/airflow/logs/nl_chunks_{nl_doc_id}_{uuid4().hex[:8]}.pkl"
    with open(output_file, 'wb') as f:
        pickle.dump(output_data, f)

    log.info(f"Saved structured chunks to {output_file}")

    result = {
        'data_file': output_file,
        'chunk_count': len(structured_chunks),
        'chunking_strategy': 'table',
        'doc_type': 'excel',
        'original_filename': chunk_data.get('original_filename', ''),
    }
    # Propagate document_id for downstream tasks and scoped cleanup
    if chunks_result.get('document_id'):
        result['document_id'] = chunks_result['document_id']
    return result


@task.external_python(python=EXTERNAL_PYTHON)
def hierarchical_chunk_text_task(
    document_id: str,
    parse_result: Dict[str, Any] = None,
    chunk_sizes: List[int] = None,
    fetch_result_md: Dict[str, Any] = None,
) -> Dict[str, Any]:
    """
    Hierarchically chunk text documents using LlamaIndex HierarchicalNodeParser.

    Creates 2-level hierarchy (MarkdownNodeParser -> SentenceSplitter(512)) with
    parent-child relationships for Word/PDF/PPTX documents. For markdown (when
    ``fetch_result_md`` is provided), reads the .md file directly and uses the
    same mermaid-aware Excel parser map (MarkdownNodeParser ->
    MermaidAwareLineSplitter), so ```mermaid fences stay inline with their
    surrounding content as atomic line-groups. A leaf containing a fence may
    grow up to ``chunker_max = embed_max - EMBED_SAFETY_BUFFER`` so the diagram
    stays adjacent to its explanatory prose; oversize single fences fall
    through to ``split_oversize_mermaid``.

    Args:
        document_id: Document UUID string
        parse_result: Result from parse task (contains data_file path with DoclingDocument).
            Required for non-markdown formats; ignored when ``fetch_result_md`` is set.
        chunk_sizes: Legacy kwarg, kept for backward compatibility. Ignored by both
            branches (the parser map is fixed: markdown -> 512).
        fetch_result_md: Result from fetch task for the .md file. When set, the task
            takes the markdown branch; otherwise uses the existing parse_result path.

    Returns:
        Dict with nodes_file path and metadata
    """
    import logging
    import pickle
    import os
    from uuid import UUID

    from llama_index.core import Document
    from llama_index.core.node_parser import HierarchicalNodeParser

    from src.llm.factory import get_token_counter, get_tokenizer

    log = logging.getLogger(__name__)

    # Legacy kwarg, retained for backward compatibility; the parser map is fixed.
    if chunk_sizes is None:
        chunk_sizes = [512]

    doc_uuid = UUID(document_id)

    # Provider-aware tokenizer (Azure -> tiktoken cl100k_base; vLLM -> HF tokenizer)
    tokenizer = get_tokenizer()
    token_count = get_token_counter()

    # ===== Markdown branch =====
    # The DAG always wires fetch_result_md (it can't conditionally omit kwargs at
    # XCom resolution time), so we gate on the actual filename extension here.
    _md_filename = (fetch_result_md or {}).get('original_filename') or ''
    if fetch_result_md is not None and _md_filename.lower().endswith('.md'):
        from llama_index.core.node_parser import get_leaf_nodes
        from llama_index.core.schema import NodeRelationship, RelatedNodeInfo
        from src.core.config import Settings
        from src.extraction_v2.mermaid_chunking import (
            TEXT_PARSER_IDS,
            build_excel_parser_map,
            prune_header_only_nodes,
        )

        embed_max_tokens = Settings().embedding_model_context_tokens
        # Match the Excel branch: leave headroom against the embed ceiling for
        # retained metadata (~100 tok template + filename/document_id) plus
        # Azure-tokenizer variance vs local cl100k_base counts.
        EMBED_SAFETY_BUFFER = 400
        chunker_max = embed_max_tokens - EMBED_SAFETY_BUFFER

        local_path = fetch_result_md.get('local_path')
        original_filename = _md_filename or 'document.md'
        log.info(
            f"Markdown chunking: {original_filename} (path={local_path}, "
            f"embed_max={embed_max_tokens} chunker_max={chunker_max} buffer={EMBED_SAFETY_BUFFER})"
        )

        with open(local_path, 'rb') as fh:
            content = fh.read().decode('utf-8')  # let UnicodeDecodeError propagate

        if not content.strip():
            raise ValueError("Markdown file is empty")

        # Mermaid-aware 2-level hierarchy (same parser map as the Excel branch):
        # MarkdownNodeParser carves sections by ATX headers, then
        # MermaidAwareLineSplitter produces leaves where any ```mermaid fence
        # stays atomic and sits inline with its surrounding prose if the leaf
        # fits under chunker_max. Oversize single fences fall through to
        # split_oversize_mermaid (body-line split only — fence boundary preserved).
        md_node_parser_map = build_excel_parser_map(
            tokenize_fn=token_count,
            embed_max_tokens=chunker_max,
        )
        md_parser = HierarchicalNodeParser(
            node_parser_ids=TEXT_PARSER_IDS,
            node_parser_map=md_node_parser_map,
        )
        doc = Document(
            doc_id=str(doc_uuid),
            text=content,
            metadata={
                "document_id": str(doc_uuid),
                "filename": original_filename,
                "document_type": "markdown",
                "page_number": 1,
            },
        )
        nodes = md_parser.get_nodes_from_documents([doc])
        nodes = list(prune_header_only_nodes(nodes))

        leaf_nodes = get_leaf_nodes(nodes)
        mermaid_bearing = sum(1 for n in leaf_nodes if n.metadata.get("has_mermaid"))
        log.info(
            f"Markdown: {len(nodes)} total nodes, {len(leaf_nodes)} leaves "
            f"({mermaid_bearing} mermaid-bearing)"
        )

        # Common post-processing: ref_doc_id, token_count metadata, save pickle.
        doc_id_str = str(doc_uuid)
        for node in nodes:
            if not node.ref_doc_id:
                node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=doc_id_str)
            node.metadata["token_count"] = token_count(node.text)
            node.metadata["char_count"] = len(node.text)

        nodes_dir = '/opt/airflow/logs/hierarchical_nodes'
        os.makedirs(nodes_dir, exist_ok=True)
        nodes_file = os.path.join(nodes_dir, f'nodes_{document_id}.pkl')
        nodes_data = {
            'nodes': nodes,
            'node_count': len(nodes),
            'leaf_count': len(leaf_nodes),
            'doc_type': 'markdown',
            'original_filename': original_filename,
            'document_id': document_id,
            'chunk_sizes': [512],
            'mermaid_bearing_leaves': mermaid_bearing,
        }
        with open(nodes_file, 'wb') as f:
            pickle.dump(nodes_data, f)
        log.info(f"Saved {len(nodes)} markdown nodes to {nodes_file}")
        return {
            'nodes_file': nodes_file,
            'node_count': len(nodes),
            'leaf_count': len(leaf_nodes),
            'doc_type': 'markdown',
            'original_filename': original_filename,
            'document_id': document_id,
        }
    # ===== End markdown branch =====

    if not parse_result:
        log.warning("No parse result provided")
        return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0}

    # Load data from pickle file if data_file path is provided
    data_file = parse_result.get('data_file')
    if data_file and os.path.exists(data_file):
        log.info(f"Loading parse data from {data_file}")
        with open(data_file, 'rb') as f:
            parse_result = pickle.load(f)
        # Clean up the temp file after loading
        try:
            os.remove(data_file)
            log.info(f"Cleaned up parse data file: {data_file}")
        except Exception as e:
            log.warning(f"Failed to clean up parse data file {data_file}: {e}")

    original_filename = parse_result.get('original_filename', 'document')
    doc_type = parse_result.get('doc_type', 'word')

    log.info(f"Hierarchical chunking document: {original_filename} (type: {doc_type})")

    # Non-md branch uses the unified text parser map:
    #   MarkdownNodeParser -> SentenceSplitter(512, "\n\n")
    # `chunk_sizes` is a legacy kwarg ignored by this branch.
    from src.extraction_v2.mermaid_chunking import (
        TEXT_PARSER_IDS,
        build_text_parser_map,
        prune_header_only_nodes,
    )

    if chunk_sizes != [512]:
        log.info(
            f"chunk_sizes={chunk_sizes} ignored for non-md branch; "
            "using unified text parser map (markdown, 512)."
        )

    node_parser_map = build_text_parser_map(tokenizer)
    node_parser = HierarchicalNodeParser(
        node_parser_ids=TEXT_PARSER_IDS,
        node_parser_map=node_parser_map,
    )

    # Create LlamaIndex Documents from pre-serialized per-page markdown
    documents = []
    pages = parse_result.get('pages', [])

    if pages:
        total_pages = len(pages)
        for page_data in pages:
            page_number = page_data.get('page_number', 1)
            markdown = page_data.get('markdown', '')
            if not markdown or not markdown.strip():
                continue
            doc = Document(
                doc_id=str(doc_uuid),
                text=markdown,
                metadata={
                    "document_id": str(doc_uuid),
                    "filename": original_filename,
                    "document_type": doc_type,
                    "page_number": page_number,
                    "total_pages": total_pages,
                },
            )
            documents.append(doc)
    else:
        # Legacy fallback: use raw text content
        log.warning("No pages found, using fallback text extraction")
        text_content = parse_result.get('markdown_content', '') or parse_result.get('text', '') or parse_result.get('content', '')
        if text_content:
            doc = Document(
                doc_id=str(doc_uuid),
                text=text_content,
                metadata={
                    "document_id": str(doc_uuid),
                    "filename": original_filename,
                    "document_type": doc_type,
                    "page_number": 1,
                },
            )
            documents.append(doc)

    if not documents:
        log.warning("No content to chunk")
        return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0}

    log.info(f"Created {len(documents)} LlamaIndex Documents")

    # Apply hierarchical parsing
    nodes = node_parser.get_nodes_from_documents(documents)
    log.info(f"Created {len(nodes)} total nodes (before header-only prune)")

    nodes = list(prune_header_only_nodes(nodes))
    log.info(f"After header-only prune: {len(nodes)} nodes")

    # Count leaf nodes
    from llama_index.core.node_parser import get_leaf_nodes
    leaf_nodes = get_leaf_nodes(nodes)
    log.info(f"Leaf nodes: {len(leaf_nodes)}")

    # Ensure ALL nodes have ref_doc_id set (used as doc_id in Milvus)
    from llama_index.core.schema import NodeRelationship, RelatedNodeInfo
    doc_id_str = str(doc_uuid)
    for node in nodes:
        if not node.ref_doc_id:
            node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=doc_id_str)

    # Add token count to each node's metadata
    for node in nodes:
        node.metadata["token_count"] = token_count(node.text)
        node.metadata["char_count"] = len(node.text)

    # Save nodes to pickle file (they can be large)
    nodes_dir = '/opt/airflow/logs/hierarchical_nodes'
    os.makedirs(nodes_dir, exist_ok=True)
    nodes_file = os.path.join(nodes_dir, f'nodes_{document_id}.pkl')

    nodes_data = {
        'nodes': nodes,
        'node_count': len(nodes),
        'leaf_count': len(leaf_nodes),
        'doc_type': doc_type,
        'original_filename': original_filename,
        'document_id': document_id,
        'parser_ids': TEXT_PARSER_IDS,
    }

    with open(nodes_file, 'wb') as f:
        pickle.dump(nodes_data, f)

    log.info(f"Saved {len(nodes)} nodes to {nodes_file}")

    return {
        'nodes_file': nodes_file,
        'node_count': len(nodes),
        'leaf_count': len(leaf_nodes),
        'doc_type': doc_type,
        'original_filename': original_filename,
        'document_id': document_id,
    }


@task.external_python(python=EXTERNAL_PYTHON)
def chunk_text_task(
    document_id: str,
    parse_result: Dict[str, Any],
    max_tokens: int = 8000
) -> Dict[str, Any]:
    """
    Chunk text documents semantically using SemanticChunker.

    Args:
        document_id: Document UUID string
        parse_result: Result from parse task (word/pdf/powerpoint) - contains data_file path
        max_tokens: Maximum tokens per chunk

    Returns:
        ContentChunks dict with semantic chunks
    """
    from uuid import UUID
    from src.extraction_v2.semantic_chunker import SemanticChunker
    from src.extraction_v2.document_converter import DocumentToHtmlConverter
    import logging
    import pickle
    import os

    log = logging.getLogger(__name__)

    doc_uuid = UUID(document_id)

    if not parse_result:
        log.warning("No parse result provided")
        return {'chunks': [], 'chunk_count': 0, 'chunking_strategy': 'semantic'}

    # Load data from pickle file if data_file path is provided
    data_file = parse_result.get('data_file')
    if data_file and os.path.exists(data_file):
        log.info(f"Loading parse data from {data_file}")
        with open(data_file, 'rb') as f:
            parse_result = pickle.load(f)
        # Clean up the temp file after loading
        try:
            os.remove(data_file)
            log.info(f"Cleaned up parse data file: {data_file}")
        except Exception as e:
            log.warning(f"Failed to clean up parse data file {data_file}: {e}")

    local_path = parse_result.get('local_path')
    if not local_path:
        log.warning("No local_path in parse result")
        return {'chunks': [], 'chunk_count': 0, 'chunking_strategy': 'semantic'}

    original_filename = parse_result.get('original_filename', 'document')
    doc_type = parse_result.get('doc_type', 'word')

    log.info(f"Semantic chunking document: {original_filename} (type: {doc_type})")

    # Initialize variables that may be used in finally block
    converted_path = None
    converted_path_from_parse = None

    try:
        # USE STORED DOCLING DOCUMENT - NO RE-CONVERSION!
        docling_doc = parse_result.get('docling_document')

        if docling_doc is None:
            # Fallback for backward compatibility (if parse task didn't store DoclingDocument)
            log.warning("No stored DoclingDocument, falling back to re-conversion")
            from pathlib import Path
            import os
            from docling.document_converter import DocumentConverter

            converted_path_from_parse = parse_result.get('converted_path', '')
            file_path = Path(local_path)
            converted_path = None
            file_to_convert = local_path

            log.info(f"Original file path: {local_path}")
            log.info(f"Original file exists: {os.path.exists(local_path)}")
            log.info(f"Converted path from parse: {converted_path_from_parse}")
            if converted_path_from_parse:
                log.info(f"Converted file exists: {os.path.exists(converted_path_from_parse)}")

            # Use pre-converted file if available (avoids double conversion)
            if converted_path_from_parse and os.path.exists(converted_path_from_parse):
                log.info(f"Using pre-converted .docx from parse task: {converted_path_from_parse}")
                file_to_convert = converted_path_from_parse
                converted_path = converted_path_from_parse
            # Otherwise, check if we need to convert legacy Word formats
            elif file_path.suffix.lower() in ['.doc', '.docm']:
                log.info(f"Converting {file_path.suffix} to .docx before chunking...")
                log.info(f"File path before conversion: {file_path}")
                log.info(f"File exists before conversion: {os.path.exists(file_path)}")

                if not os.path.exists(file_path):
                    raise ValueError(f"Original file not found: {file_path}. The temporary directory may have been cleaned up.")

                html_converter = DocumentToHtmlConverter()
                converted_path = html_converter._convert_doc_to_docx(file_path)

                if converted_path is None:
                    raise ValueError(f"Failed to convert {file_path.suffix} to .docx: LibreOffice conversion failed")

                file_to_convert = converted_path
                log.info(f"Converted to .docx: {file_to_convert}")

            # Re-convert document for full structure
            converter = DocumentConverter()
            result = converter.convert(file_to_convert)
            docling_doc = result.document

        # Initialize chunker and chunk document
        chunker = SemanticChunker(max_tokens=max_tokens)
        chunks = chunker.chunk_document(
            doc=docling_doc,
            file_id=doc_uuid,
            original_filename=original_filename,
        )

        # Add doc_type to each chunk for downstream tasks
        for chunk in chunks:
            chunk['doc_type'] = doc_type
            chunk['original_filename'] = original_filename

        log.info(f"Created {len(chunks)} semantic chunks")

        # Save large data to pickle file to avoid XCom size limits
        chunk_data = {
            'chunks': chunks,
            'chunk_count': len(chunks),
            'chunking_strategy': 'semantic',
            'doc_type': doc_type,
            'original_filename': original_filename,
        }

        data_file = f"/opt/airflow/logs/chunks_text_{document_id}.pkl"
        with open(data_file, 'wb') as f:
            pickle.dump(chunk_data, f)

        log.info(f"Saved chunk result to {data_file}")

        # Return only small metadata via XCom
        return {
            'data_file': data_file,
            'chunk_count': len(chunks),
            'chunking_strategy': 'semantic',
            'doc_type': doc_type,
            'original_filename': original_filename,
            'document_id': document_id,
        }

    except Exception as e:
        log.error(f"Semantic chunking failed: {e}")
        raise
    finally:
        # Clean up temporary converted file ONLY if we created it
        # Don't delete pre-converted files from parse task
        if converted_path and converted_path != converted_path_from_parse and os.path.exists(converted_path):
            try:
                import shutil
                temp_dir = os.path.dirname(converted_path)
                if temp_dir and os.path.exists(temp_dir) and 'tmp' in temp_dir:
                    shutil.rmtree(temp_dir)
                    log.debug(f"Cleaned up temporary directory: {temp_dir}")
            except Exception as cleanup_error:
                log.warning(f"Failed to clean up temp file {converted_path}: {cleanup_error}")


@task.external_python(python=EXTERNAL_PYTHON)
def extract_domain_terms_task(chunks_result: Dict[str, Any], tenant_id: str, node_id: str) -> Dict[str, Any]:
    """
    Extract domain terms from chunks for domain-aware RAG optimization.

    Analyzes chunk content and extracts domain-specific terms and categories
    to enable efficient pre-filtering during vector search.

    Args:
        chunks_result: Result from chunk task (contains data_file path)

    Returns:
        Dict with file path to chunk data enriched with domain terms
    """
    from src.extraction_v2.domain_term_extractor import DomainTermExtractor
    import asyncio
    import logging
    import pickle
    import os
    from uuid import uuid4

    log = logging.getLogger(__name__)

    if not chunks_result:
        log.warning("No chunks result provided")
        return {
            'chunk_count': 0,
            'doc_type': 'excel',
            'original_filename': '',
        }

    # Load data from pickle file if data_file path is provided
    # Preserve original metadata before loading pickle
    original_metadata = {k: v for k, v in chunks_result.items() if k != 'data_file'}
    data_file = chunks_result.get('data_file')
    if data_file and os.path.exists(data_file):
        log.info(f"Loading chunk data from {data_file}")
        with open(data_file, 'rb') as f:
            chunks_result = pickle.load(f)
        # Restore important metadata from XCom (like 'adapted' flag)
        for key in ['adapted', 'document_id', 'filename']:
            if key in original_metadata and key not in chunks_result:
                chunks_result[key] = original_metadata[key]
        # Clean up the temp file after loading
        try:
            os.remove(data_file)
            log.info(f"Cleaned up chunk data file: {data_file}")
        except Exception as e:
            log.warning(f"Failed to clean up chunk data file {data_file}: {e}")

    # Handle both NEW pipeline format ('chunks') and OLD pipeline format ('records')
    chunks = chunks_result.get('chunks') or chunks_result.get('records', [])
    if not chunks:
        log.info("No chunks to process")
        return {
            'chunk_count': 0,
            'doc_type': chunks_result.get('doc_type', 'excel'),
            'original_filename': chunks_result.get('original_filename', ''),
        }

    log.info(f"Extracting domain terms from {len(chunks)} chunks")

    # Initialize extractor
    extractor = DomainTermExtractor()

    # Extract domain terms concurrently
    doc_type = chunks_result.get('doc_type', 'excel')
    domain_data = asyncio.run(
        extractor.extract_from_chunks_async(chunks, doc_type=doc_type)
    )

    log.info(f"Domain term extraction complete: {len(domain_data)} chunks processed")

    # Store enriched data to temporary file (include document_id for scoped cleanup)
    data_dir = '/opt/airflow/logs/domain_data'
    os.makedirs(data_dir, exist_ok=True)
    doc_id = chunks_result.get('document_id', uuid4().hex)
    output_file = os.path.join(data_dir, f'domain_terms_{doc_id}_{uuid4().hex[:8]}.pkl')

    # Determine the correct key name: 'records' for OLD pipeline, 'chunks' for NEW pipeline
    # Also preserve embeddings if they exist (from OLD pipeline)
    chunk_key = 'records' if chunks_result.get('adapted') or chunks_result.get('records') else 'chunks'

    enriched_data = {
        chunk_key: chunks,
        'domain_data': domain_data,
        'chunk_count': len(chunks),
        'doc_type': doc_type,
        'original_filename': chunks_result.get('original_filename', ''),
    }

    # Preserve embeddings if they exist (from OLD pipeline which already has them)
    if 'embeddings' in chunks_result:
        enriched_data['embeddings'] = chunks_result['embeddings']

    with open(output_file, 'wb') as f:
        pickle.dump(enriched_data, f)

    log.info(f"Saved enriched data to {output_file}")

    # Persist new terms to domain dictionary
    all_terms = set()
    term_confidences: dict[str, float] = {}
    for data in domain_data:
        all_terms.update(data['domain_terms'])
        # Max across chunks: same term appearing in multiple chunks keeps the
        # strongest confidence. Rule-only chunks emit 0.95; LLM chunks emit
        # exp(min(token_logprobs)) or 0.70 on alignment failure.
        for t, c in data.get('domain_confidences', {}).items():
            term_confidences[t] = max(term_confidences.get(t, 0.0), c)

    # Infer dominant category from chunk categories
    from collections import Counter
    category_counts_list = Counter(data['domain_category'] for data in domain_data)
    dominant_category = category_counts_list.most_common(1)[0][0] if category_counts_list else 'general'

    # Get document_id from chunks_result metadata (passed via XCom)
    document_id = chunks_result.get('document_id')  # May be None for some pipelines
    if document_id:
        from uuid import UUID
        document_id = UUID(document_id) if isinstance(document_id, str) else document_id

    # Persist new terms to dictionary (skip duplicates, translate via LLM)
    if all_terms:
        try:
            new_count = extractor.persist_new_terms(
                terms=list(all_terms),
                domain_category=dominant_category,
                document_id=document_id,
                tenant_id=tenant_id,
                node_id=node_id,
                confidence_by_term=term_confidences,
            )
            log.info(f"Persisted {new_count} new terms to domain dictionary")
        except Exception as e:
            log.warning(f"Failed to persist terms to dictionary: {e}. Continuing with pipeline.")

    # Return file path via XCom (preserve 'adapted' flag and document_id)
    result = {
        'data_file': output_file,
        'chunk_count': len(chunks),
        'doc_type': doc_type,
        'original_filename': chunks_result.get('original_filename', ''),
    }
    # Preserve 'adapted' flag from OLD pipeline if present
    if chunks_result.get('adapted'):
        result['adapted'] = True
    # Propagate document_id for downstream tasks and scoped cleanup
    if chunks_result.get('document_id'):
        result['document_id'] = chunks_result['document_id']
    return result


@task.external_python(python=EXTERNAL_PYTHON)
def tag_records_task(chunks_result: Dict[str, Any]) -> Dict[str, Any]:
    """
    Assign semantic tags to records using SemanticTagger.

    This task runs AFTER domain_terms extraction and BEFORE embedding generation.
    Tags are added to each record for downstream filtering in Milvus.

    Tags include:
    - req:functional, req:non_functional, req:constraint, etc. (ISO 29148)
    - risk:safety, risk:security, risk:assumption
    - info:background, info:scope
    - topic:{keyword} (LLM-extracted domain topics)

    Args:
        chunks_result: Result from domain_terms task (contains data_file path)

    Returns:
        Dict with file path to tagged data and metadata
    """
    from src.extraction_v2.semantic_tagger import SemanticTagger
    import asyncio
    import logging
    import pickle
    import os
    from uuid import uuid4

    log = logging.getLogger(__name__)

    if not chunks_result:
        log.warning("No chunks result provided")
        return {
            'chunk_count': 0,
            'doc_type': 'excel',
            'original_filename': '',
        }

    # Preserve original metadata before loading pickle
    original_metadata = {k: v for k, v in chunks_result.items() if k != 'data_file'}
    data_file = chunks_result.get('data_file')

    if data_file and os.path.exists(data_file):
        log.info(f"Loading chunk data from {data_file}")
        with open(data_file, 'rb') as f:
            chunks_result = pickle.load(f)
        # Restore important metadata from XCom
        for key in ['adapted', 'document_id', 'filename']:
            if key in original_metadata and key not in chunks_result:
                chunks_result[key] = original_metadata[key]
        # Clean up the temp file after loading
        try:
            os.remove(data_file)
            log.info(f"Cleaned up chunk data file: {data_file}")
        except Exception as e:
            log.warning(f"Failed to clean up chunk data file {data_file}: {e}")

    # Handle both NEW pipeline format ('chunks') and OLD pipeline format ('records')
    chunks = chunks_result.get('chunks') or chunks_result.get('records', [])
    if not chunks:
        log.info("No chunks to tag")
        return {
            'chunk_count': 0,
            'doc_type': chunks_result.get('doc_type', 'excel'),
            'original_filename': chunks_result.get('original_filename', ''),
        }

    log.info(f"Tagging {len(chunks)} chunks with semantic tags")

    try:
        # Initialize tagger
        tagger = SemanticTagger()

        # Prepare records for tagging (tagger expects specific format)
        records_for_tagging = []
        for idx, chunk in enumerate(chunks):
            record = {
                'id': idx,
                'text': '',
                'parent_header': '',
                'sheet_name': '',
            }

            # Extract text content
            if isinstance(chunk, dict):
                content = chunk.get('content') or chunk.get('text', '')
                if isinstance(content, dict):
                    # Excel record: format as key: value pairs
                    record['text'] = ' | '.join(f"{k}: {v}" for k, v in content.items() if v)
                else:
                    record['text'] = str(content) if content else ''

                # Extract context from metadata
                source = chunk.get('_source', {})
                metadata = chunk.get('metadata', {})
                record['parent_header'] = metadata.get('parent_header', '')
                record['sheet_name'] = source.get('sheet', '')
            else:
                record['text'] = str(chunk)

            if record['text'].strip():
                records_for_tagging.append(record)

        if not records_for_tagging:
            log.info("No valid text content to tag")
        else:
            # Tag records concurrently (modifies in-place)
            asyncio.run(tagger.tag_records_async(records_for_tagging))

            # Copy tags back to original chunks
            for record in records_for_tagging:
                idx = record['id']
                tags = record.get('tags', [])
                # TEXTIQ-569: SemanticTagger also emits verified_tag_names
                # post-inheritance; carry it forward to Milvus.
                verified_tag_names = record.get('verified_tag_names', [])
                if idx < len(chunks):
                    if isinstance(chunks[idx], dict):
                        chunks[idx]['tags'] = tags
                        chunks[idx]['verified_tag_names'] = verified_tag_names
                    elif hasattr(chunks[idx], 'tags'):
                        chunks[idx].tags = tags
                        if hasattr(chunks[idx], 'verified_tag_names'):
                            chunks[idx].verified_tag_names = verified_tag_names

            # Count tags assigned
            total_tags = sum(len(r.get('tags', [])) for r in records_for_tagging)
            log.info(f"Assigned {total_tags} total tags to {len(records_for_tagging)} chunks")

    except Exception as e:
        log.warning(f"Tagging failed: {e}. Continuing with empty tags.")
        # On failure, ensure all chunks have empty tags list
        for chunk in chunks:
            if isinstance(chunk, dict) and 'tags' not in chunk:
                chunk['tags'] = []

    # Save tagged data to temporary file (include document_id for scoped cleanup)
    data_dir = '/opt/airflow/logs/tag_data'
    os.makedirs(data_dir, exist_ok=True)
    doc_id = chunks_result.get('document_id', uuid4().hex)
    output_file = os.path.join(data_dir, f'tagged_{doc_id}_{uuid4().hex[:8]}.pkl')

    # Determine the correct key name
    chunk_key = 'records' if chunks_result.get('adapted') or chunks_result.get('records') else 'chunks'

    tagged_data = {
        chunk_key: chunks,
        'domain_data': chunks_result.get('domain_data', []),
        'embeddings': chunks_result.get('embeddings', []),
        'chunk_count': len(chunks),
        'doc_type': chunks_result.get('doc_type', 'excel'),
        'original_filename': chunks_result.get('original_filename', ''),
    }

    with open(output_file, 'wb') as f:
        pickle.dump(tagged_data, f)

    log.info(f"Saved tagged data to {output_file}")

    # Return file path via XCom
    result = {
        'data_file': output_file,
        'chunk_count': len(chunks),
        'doc_type': chunks_result.get('doc_type', 'excel'),
        'original_filename': chunks_result.get('original_filename', ''),
    }
    # Preserve 'adapted' flag if present
    if chunks_result.get('adapted'):
        result['adapted'] = True
    # Propagate document_id for downstream tasks and scoped cleanup
    if chunks_result.get('document_id'):
        result['document_id'] = chunks_result['document_id']
    return result


@task.external_python(python=EXTERNAL_PYTHON)
def extract_domain_terms_nodes_task(
    nodes_result: Dict[str, Any],
    tenant_id: str,
    node_id: str,
) -> Dict[str, Any]:
    """
    Extract domain terms from LlamaIndex nodes for domain-aware RAG.

    Works with hierarchical nodes from hierarchical_chunk_*_task.
    Enriches node.metadata with domain_category and domain_terms.

    Args:
        nodes_result: Result from hierarchical_chunk task (contains nodes_file path)

    Returns:
        Dict with enriched nodes_file path and metadata
    """
    from src.extraction_v2.domain_term_extractor import DomainTermExtractor
    import asyncio
    import logging
    import pickle
    import os
    from uuid import UUID
    from collections import Counter

    log = logging.getLogger(__name__)

    if not nodes_result:
        log.warning("No nodes result provided")
        return {'nodes_file': '', 'node_count': 0}

    # Load nodes from pickle file
    nodes_file = nodes_result.get('nodes_file')
    if not nodes_file or not os.path.exists(nodes_file):
        log.warning(f"Nodes file not found: {nodes_file}")
        return {'nodes_file': '', 'node_count': 0}

    log.info(f"Loading nodes from {nodes_file}")
    with open(nodes_file, 'rb') as f:
        nodes_data = pickle.load(f)

    nodes = nodes_data.get('nodes', [])
    if not nodes:
        log.info("No nodes to process")
        return nodes_result

    log.info(f"Extracting domain terms from {len(nodes)} nodes concurrently")

    # Initialize extractor and run async extraction
    extractor = DomainTermExtractor()
    doc_type = nodes_data.get('doc_type', 'word')

    all_terms, category_counts, term_confidences = asyncio.run(
        extractor.extract_from_nodes_async(nodes, doc_type=doc_type)
    )

    log.info(f"Domain term extraction complete: {len(nodes)} nodes processed")
    log.info(f"Category distribution: {category_counts}")

    # Persist new terms to dictionary
    dominant_category = Counter(category_counts).most_common(1)[0][0] if category_counts else 'general'
    document_id = nodes_data.get('document_id')
    if document_id:
        document_id = UUID(document_id) if isinstance(document_id, str) else document_id

    if all_terms:
        try:
            new_count = extractor.persist_new_terms(
                terms=list(all_terms),
                domain_category=dominant_category,
                document_id=document_id,
                tenant_id=tenant_id,
                node_id=node_id,
                confidence_by_term=term_confidences,
            )
            log.info(f"Persisted {new_count} new terms to domain dictionary")
        except Exception as e:
            log.warning(f"Failed to persist terms: {e}")

    # Save enriched nodes back
    nodes_data['nodes'] = nodes
    with open(nodes_file, 'wb') as f:
        pickle.dump(nodes_data, f)

    log.info(f"Saved enriched nodes to {nodes_file}")

    return {
        'nodes_file': nodes_file,
        'node_count': len(nodes),
        'leaf_count': nodes_data.get('leaf_count', 0),
        'doc_type': doc_type,
        'original_filename': nodes_data.get('original_filename', ''),
        'document_id': nodes_data.get('document_id', ''),
    }


@task.external_python(python=EXTERNAL_PYTHON)
def tag_records_nodes_task(
    nodes_result: Dict[str, Any],
) -> Dict[str, Any]:
    """
    Assign semantic tags to LlamaIndex nodes using SemanticTagger.

    Works with hierarchical nodes from hierarchical_chunk_*_task.
    Enriches node.metadata with tags list.

    Tags include:
    - req:functional, req:non_functional, req:constraint, etc. (ISO 29148)
    - risk:safety, risk:security, risk:assumption
    - info:background, info:scope
    - topic:{keyword} (LLM-extracted domain topics)

    Args:
        nodes_result: Result from extract_domain_terms_nodes_task

    Returns:
        Dict with tagged nodes_file path and metadata
    """
    from src.extraction_v2.semantic_tagger import SemanticTagger
    import asyncio
    import logging
    import pickle
    import os

    log = logging.getLogger(__name__)

    if not nodes_result:
        log.warning("No nodes result provided")
        return {'nodes_file': '', 'node_count': 0}

    # Load nodes from pickle file
    nodes_file = nodes_result.get('nodes_file')
    if not nodes_file or not os.path.exists(nodes_file):
        log.warning(f"Nodes file not found: {nodes_file}")
        return {'nodes_file': '', 'node_count': 0}

    log.info(f"Loading nodes from {nodes_file}")
    with open(nodes_file, 'rb') as f:
        nodes_data = pickle.load(f)

    nodes = nodes_data.get('nodes', [])
    if not nodes:
        log.info("No nodes to tag")
        return nodes_result

    log.info(f"Tagging {len(nodes)} nodes with semantic tags concurrently")

    try:
        # Initialize tagger
        tagger = SemanticTagger()

        # Prepare records for tagging
        records_for_tagging = []
        for idx, node in enumerate(nodes):
            text = node.text if hasattr(node, 'text') else str(node)
            metadata = node.metadata if hasattr(node, 'metadata') else {}

            if text.strip():
                records_for_tagging.append({
                    'id': idx,
                    'text': text,
                    'parent_header': metadata.get('parent_header', ''),
                    'sheet_name': metadata.get('sheet_name', ''),
                })

        if records_for_tagging:
            # Tag records concurrently (modifies in-place)
            asyncio.run(tagger.tag_records_async(records_for_tagging))

            # Copy tags back to nodes (convert TagAssignment objects to dicts for JSON serialization)
            for record in records_for_tagging:
                idx = record['id']
                tags = record.get('tags', [])
                if idx < len(nodes):
                    # TagAssignment is a Pydantic model - convert to dict
                    serialized_tags = [
                        t.model_dump() if hasattr(t, 'model_dump') else t
                        for t in tags
                    ] if tags else []
                    nodes[idx].metadata['tags'] = serialized_tags
                    # TEXTIQ-569: SemanticTagger emits `verified_tag_names` on
                    # each record post-inheritance; carry it onto the node so
                    # LlamaIndex writes it into the Milvus dynamic envelope.
                    nodes[idx].metadata['verified_tag_names'] = record.get(
                        'verified_tag_names', []
                    )

            total_tags = sum(len(r.get('tags', [])) for r in records_for_tagging)
            log.info(f"Assigned {total_tags} total tags to {len(records_for_tagging)} nodes")
        else:
            log.info("No valid text content to tag")

    except Exception as e:
        log.warning(f"Tagging failed: {e}. Continuing with empty tags.")
        for node in nodes:
            if 'tags' not in node.metadata:
                node.metadata['tags'] = []
            if 'verified_tag_names' not in node.metadata:
                node.metadata['verified_tag_names'] = []

    # Save tagged nodes back
    nodes_data['nodes'] = nodes
    with open(nodes_file, 'wb') as f:
        pickle.dump(nodes_data, f)

    log.info(f"Saved tagged nodes to {nodes_file}")

    return {
        'nodes_file': nodes_file,
        'node_count': len(nodes),
        'leaf_count': nodes_data.get('leaf_count', 0),
        'doc_type': nodes_data.get('doc_type', 'word'),
        'original_filename': nodes_data.get('original_filename', ''),
        'document_id': nodes_data.get('document_id', ''),
    }


@task.external_python(python=EXTERNAL_PYTHON)
def generate_embeddings_task(
    chunks_result: Dict[str, Any],
    batch_size: int = 16
) -> Dict[str, Any]:
    """
    Generate embeddings for content chunks.

    Stores large data (chunks + embeddings + domain_data) to a temporary file to avoid XCom
    size limits. Returns file path via XCom instead of the actual data.

    Args:
        chunks_result: Result from extract_domain_terms task (contains data_file path)
        batch_size: Number of chunks to process in each batch (max 100 for Azure OpenAI)

    Returns:
        Dict with file path to embedded data and metadata
    """
    from src.knowledge.embeddings import EmbeddingService
    import asyncio
    import logging
    import pickle
    import tempfile
    import os
    from uuid import uuid4

    log = logging.getLogger(__name__)

    if not chunks_result:
        log.warning("No chunks result provided")
        return {
            'chunk_count': 0,
            'doc_type': 'excel',
            'original_filename': '',
        }

    # Load data from pickle file if data_file path is provided
    # Preserve original metadata before loading pickle
    original_metadata = {k: v for k, v in chunks_result.items() if k != 'data_file'}
    data_file = chunks_result.get('data_file')
    if data_file and os.path.exists(data_file):
        log.info(f"Loading chunk data from {data_file}")
        with open(data_file, 'rb') as f:
            chunks_result = pickle.load(f)
        # Restore important metadata from XCom (like 'adapted' flag)
        for key in ['adapted', 'document_id', 'filename']:
            if key in original_metadata and key not in chunks_result:
                chunks_result[key] = original_metadata[key]
        # Clean up the temp file after loading
        try:
            os.remove(data_file)
            log.info(f"Cleaned up chunk data file: {data_file}")
        except Exception as e:
            log.warning(f"Failed to clean up chunk data file {data_file}: {e}")

    # Handle both NEW pipeline format ('chunks') and OLD pipeline format ('records')
    chunks = chunks_result.get('chunks') or chunks_result.get('records', [])
    domain_data = chunks_result.get('domain_data', [])  # Load domain data
    if not chunks:
        log.info("No chunks to embed")
        return {
            'chunk_count': 0,
            'doc_type': chunks_result.get('doc_type', 'excel'),
            'original_filename': chunks_result.get('original_filename', ''),
        }

    log.info(f"Generating embeddings for {len(chunks)} chunks (batch_size={batch_size})")

    # Extract text from chunks for embedding
    texts = []
    valid_chunks = []

    for idx, chunk in enumerate(chunks):
        if isinstance(chunk, dict):
            text = chunk.get('text') or chunk.get('content') or chunk.get('raw_content', '')
            text = str(text) if text else ''
            if text.strip():
                texts.append(text)
                valid_chunks.append(chunk)
            else:
                log.warning(f"Skipping chunk {idx} with no valid text content")
        else:
            text = str(chunk)
            if text.strip():
                texts.append(text)
                valid_chunks.append(chunk)

    chunks = valid_chunks

    if not texts:
        log.warning("No valid texts to embed after filtering empty content")
        return {
            'chunk_count': 0,
            'doc_type': chunks_result.get('doc_type', 'excel'),
            'original_filename': chunks_result.get('original_filename', ''),
        }

    skipped_count = len(chunks_result.get('chunks', [])) - len(texts)
    if skipped_count > 0:
        log.info(f"Filtered out {skipped_count} chunks with empty content, {len(texts)} remaining")

    # Generate unique run ID for checkpointing
    run_id = uuid4().hex[:8]
    checkpoint_file = f'/opt/airflow/logs/embedding_checkpoint_{run_id}.pkl'

    async def embed_with_checkpoints():
        # Initialize embedding service
        # Note: EmbeddingService creates and closes clients per batch internally
        service = EmbeddingService()
        all_embeddings = []
        start_batch = 0

        # Check for existing checkpoint (retry scenario)
        if os.path.exists(checkpoint_file):
            log.info(f"Found checkpoint, resuming from {checkpoint_file}...")
            with open(checkpoint_file, 'rb') as f:
                checkpoint = pickle.load(f)
                all_embeddings = checkpoint['embeddings']
                start_batch = checkpoint['next_batch']
            log.info(f"Resuming from batch {start_batch}, {len(all_embeddings)} embeddings recovered")

        # Generate embeddings in batches (max 100 per Azure OpenAI call)
        effective_batch_size = min(batch_size, 100)
        total_batches = (len(texts) - 1) // effective_batch_size + 1

        for i in range(start_batch * effective_batch_size, len(texts), effective_batch_size):
            batch = texts[i:i + effective_batch_size]
            batch_num = i // effective_batch_size + 1

            log.info(f"Processing batch {batch_num}/{total_batches}")
            batch_embeddings = await service.embed_batch(batch)
            all_embeddings.extend(batch_embeddings)

            # Save checkpoint after each batch
            checkpoint = {
                'embeddings': all_embeddings,
                'next_batch': batch_num,
                'total_batches': total_batches,
            }
            with open(checkpoint_file, 'wb') as f:
                pickle.dump(checkpoint, f)

            log.info(f"Checkpoint saved: {len(all_embeddings)}/{len(texts)} embeddings")

        # Clean up checkpoint on success
        if os.path.exists(checkpoint_file):
            os.remove(checkpoint_file)

        return all_embeddings

    all_embeddings = asyncio.run(embed_with_checkpoints())
    log.info(f"Generated {len(all_embeddings)} embeddings")

    # Store large data to temporary file to avoid XCom size limits
    # Use /opt/airflow/logs as shared storage (mounted volume)
    data_dir = '/opt/airflow/logs/embedding_data'
    os.makedirs(data_dir, exist_ok=True)
    doc_id = chunks_result.get('document_id', uuid4().hex)
    data_file = os.path.join(data_dir, f'embeddings_{doc_id}_{uuid4().hex[:8]}.pkl')

    embedded_data = {
        'chunks': chunks,
        'embeddings': all_embeddings,
        'domain_data': domain_data,  # Pass through domain data
        'chunk_count': len(chunks),
        'doc_type': chunks_result.get('doc_type'),
        'original_filename': chunks_result.get('original_filename', ''),
    }

    with open(data_file, 'wb') as f:
        pickle.dump(embedded_data, f)

    file_size_mb = os.path.getsize(data_file) / (1024 * 1024)
    log.info(f"Saved embedded data to {data_file} ({file_size_mb:.2f} MB)")

    # Return only metadata and file path via XCom (small payload)
    result = {
        'data_file': data_file,
        'chunk_count': len(chunks),
        'doc_type': chunks_result.get('doc_type'),
        'original_filename': chunks_result.get('original_filename', ''),
    }
    # Propagate document_id for downstream tasks and scoped cleanup
    if chunks_result.get('document_id'):
        result['document_id'] = chunks_result['document_id']
    return result


@task.external_python(python=EXTERNAL_PYTHON)
def generate_embeddings_chunks(
    detailed_pipeline_result: Dict[str, Any]
) -> Dict[str, Any]:
    """
    Generate embeddings for DETAILED pipeline chunks.

    Converts detailed extraction chunks (images, shapes, flowcharts, text)
    to ExtractedRecord structure and generates embeddings using LegacyToMilvusAdapter.

    Args:
        detailed_pipeline_result: Result from parse_excel_detailed task (contains data_file path)

    Returns:
        Dict with embedded chunks data_file path and metadata
    """
    # All imports must be inside for external_python
    import sys
    sys.path.insert(0, '/opt/airflow')

    from src.adapters.legacy_to_milvus_adapter import LegacyToMilvusAdapter
    from uuid import UUID
    import logging
    import pickle
    import os
    import asyncio

    log = logging.getLogger(__name__)

    if not detailed_pipeline_result:
        log.warning("No DETAILED pipeline result provided")
        return {'record_count': 0, 'adapted': False}

    # Load chunks from pickle file
    data_file = detailed_pipeline_result.get('data_file')
    if not data_file or not os.path.exists(data_file):
        log.error(f"DETAILED pipeline data file not found: {data_file}")
        return {'record_count': 0, 'adapted': False}

    log.info(f"Loading DETAILED pipeline chunks from {data_file}")
    with open(data_file, 'rb') as f:
        chunks = pickle.load(f)

    # Clean up DETAILED pipeline file after loading
    try:
        os.remove(data_file)
        log.info(f"Cleaned up DETAILED pipeline data file: {data_file}")
    except Exception as e:
        log.warning(f"Failed to clean up DETAILED pipeline file {data_file}: {e}")

    if not chunks:
        log.info("No chunks to adapt")
        return {'record_count': 0, 'adapted': True}

    log.info(f"Adapting {len(chunks)} DETAILED pipeline chunks to schema")

    # Get metadata from result
    document_id = UUID(detailed_pipeline_result.get('document_id'))
    filename = detailed_pipeline_result.get('filename', 'document.xlsx')

    # Convert chunks to dicts if they're dataclass instances
    chunks_as_dicts = []
    for chunk in chunks:
        if hasattr(chunk, '__dict__'):
            # Dataclass instance - convert to dict
            chunk_dict = {
                'element_type': getattr(chunk, 'element_type', 'Text'),
                'text': getattr(chunk, 'text', ''),
                'keywords_text': getattr(chunk, 'keywords_text', ''),
                'created_at': getattr(chunk, 'created_at', ''),
                'metadata': getattr(chunk, 'metadata', {}),
            }
            chunks_as_dicts.append(chunk_dict)
        elif isinstance(chunk, dict):
            chunks_as_dicts.append(chunk)
        else:
            log.warning(f"Skipping unknown chunk type: {type(chunk)}")

    # Initialize adapter
    adapter = LegacyToMilvusAdapter()

    # Adapt chunks (async)
    async def adapt():
        return await adapter.adapt_chunks(
            legacy_chunks=chunks_as_dicts,
            document_id=document_id,
            filename=filename,
        )

    adapted_result = asyncio.run(adapt())

    log.info(f"Adapted {adapted_result['record_count']} records with embeddings")

    # Save adapted result to pickle file
    adapted_file = f"/opt/airflow/logs/adapted_{document_id}.pkl"
    with open(adapted_file, 'wb') as f:
        pickle.dump(adapted_result, f)

    log.info(f"Saved adapted result to {adapted_file}")

    # Return metadata via XCom
    return {
        'data_file': adapted_file,
        'document_id': str(document_id),
        'filename': filename,
        'record_count': adapted_result['record_count'],
        'adapted': True,
    }


@task.external_python(python=EXTERNAL_PYTHON)
def merge_pipeline_results(
    new_pipeline_result: Dict[str, Any],
    adapted_old_result: Dict[str, Any]
) -> Dict[str, Any]:
    """
    Merge NEW pipeline results with DETAILED pipeline results.

    Combines embeddings and records from both pipelines into a single
    dataset ready for Milvus storage.

    Args:
        new_pipeline_result: Result from generate_embeddings_large_tables (NEW pipeline)
        adapted_old_result: Result from generate_embeddings_chunks (DETAILED pipeline)

    Returns:
        Dict with merged data_file path and combined metadata
    """
    import logging
    import pickle
    import os

    log = logging.getLogger(__name__)

    # Load NEW pipeline data
    new_chunks = []
    new_embeddings = []
    new_domain_data = []
    if new_pipeline_result and new_pipeline_result.get('chunk_count', 0) > 0:
        new_data_file = new_pipeline_result.get('data_file', '')
        if new_data_file and os.path.exists(new_data_file):
            log.info(f"Loading NEW pipeline data from {new_data_file}")
            with open(new_data_file, 'rb') as f:
                new_data = pickle.load(f)
            new_chunks = new_data.get('chunks', [])
            new_embeddings = new_data.get('embeddings', [])
            new_domain_data = new_data.get('domain_data', [])
            # Clean up temp file
            try:
                os.remove(new_data_file)
                log.info(f"Cleaned up NEW pipeline file: {new_data_file}")
            except Exception as e:
                log.warning(f"Failed to clean up NEW pipeline file: {e}")
        else:
            log.info("NEW pipeline returned 0 chunks (no large tables found)")
    else:
        log.info("NEW pipeline skipped or returned no results (no large tables found)")

    # Load adapted DETAILED pipeline data
    detailed_records = []
    detailed_embeddings = []
    detailed_domain_data = []
    if adapted_old_result and adapted_old_result.get('adapted'):
        detailed_data_file = adapted_old_result.get('data_file')
        if detailed_data_file and os.path.exists(detailed_data_file):
            log.info(f"Loading adapted DETAILED pipeline data from {detailed_data_file}")
            with open(detailed_data_file, 'rb') as f:
                detailed_data = pickle.load(f)
            detailed_records = detailed_data.get('records', [])
            detailed_embeddings = detailed_data.get('embeddings', [])
            detailed_domain_data = detailed_data.get('domain_data', [])
            log.info(f"DEBUG merge: DETAILED pipeline loaded - records={len(detailed_records)}, embeddings={len(detailed_embeddings)}, domain_data={len(detailed_domain_data)}")
            # Clean up temp file
            try:
                os.remove(detailed_data_file)
                log.info(f"Cleaned up adapted DETAILED pipeline file: {detailed_data_file}")
            except Exception as e:
                log.warning(f"Failed to clean up adapted file: {e}")

    log.info(
        f"Merging pipelines: NEW ({len(new_chunks)} records) + "
        f"DETAILED ({len(detailed_records)} records)"
    )

    # Merge records, embeddings, and domain data
    merged_chunks = new_chunks + detailed_records
    merged_embeddings = new_embeddings + detailed_embeddings
    merged_domain_data = new_domain_data + detailed_domain_data

    if not merged_chunks:
        log.warning("No records to merge from either pipeline")
        return {
            'chunk_count': 0,
            'merged': False,
            'doc_type': 'excel',
            'original_filename': '',
        }

    log.info(f"Merged total: {len(merged_chunks)} records with embeddings and domain data")

    # Get metadata from whichever result is available
    doc_type = (new_pipeline_result.get('doc_type') if new_pipeline_result
                else 'excel')
    original_filename = (new_pipeline_result.get('original_filename') if new_pipeline_result
                         else adapted_old_result.get('filename', ''))

    # Save merged data to pickle file
    log.info(f"DEBUG merge: About to save merged_data with domain_data length = {len(merged_domain_data)}")
    merged_data = {
        'chunks': merged_chunks,
        'embeddings': merged_embeddings,
        'domain_data': merged_domain_data,
        'chunk_count': len(merged_chunks),
        'doc_type': doc_type,
        'original_filename': original_filename,
    }

    # Use document_id from either result
    document_id = (new_pipeline_result.get('document_id') if new_pipeline_result and 'document_id' in new_pipeline_result
                   else adapted_old_result.get('document_id', 'merged'))

    merged_file = f"/opt/airflow/logs/merged_{document_id}.pkl"
    with open(merged_file, 'wb') as f:
        pickle.dump(merged_data, f)

    log.info(f"Saved merged result to {merged_file}")

    # Return metadata via XCom
    return {
        'data_file': merged_file,
        'chunk_count': len(merged_chunks),
        'doc_type': doc_type,
        'original_filename': original_filename,
        'merged': True,
        'new_count': len(new_chunks),
        'detailed_count': len(detailed_records),
    }


@task.external_python(python=EXTERNAL_PYTHON)
def upload_images_to_seaweedfs(
    document_id: str,
    detailed_pipeline_result: Dict[str, Any],
    tenant_id: str,
    node_id: str,
) -> Dict[str, Any]:
    """
    Upload images from DETAILED pipeline to SeaweedFS.

    Processes chunks with image_base64 data, uploads to SeaweedFS,
    and replaces base64 with image_url and image_s3_key.

    Args:
        document_id: Document UUID string
        detailed_pipeline_result: Result from parse_excel_detailed task (contains data_file path)
        tenant_id: Tenant UUID string for S3 key scoping
        node_id: Node UUID string for S3 key scoping

    Returns:
        Updated result dict with images uploaded and metadata updated
    """
    from uuid import UUID
    from src.services.seaweedfs_service import SeaweedFSService
    import logging
    import pickle
    import os

    log = logging.getLogger(__name__)

    doc_uuid = UUID(document_id)

    if not detailed_pipeline_result:
        log.info("No detailed pipeline result provided")
        return {
            'uploaded_count': 0,
            'success': True,
            'data_file': None,
            'document_id': document_id,
            'filename': '',
            'chunk_count': 0,
        }

    # Load data from pickle file if data_file path is provided
    data_file = detailed_pipeline_result.get('data_file')
    if not data_file or not os.path.exists(data_file):
        log.warning("No data_file found in detailed_pipeline_result")
        return {
            'uploaded_count': 0,
            'success': True,
            'data_file': None,
            'document_id': detailed_pipeline_result.get('document_id'),
            'filename': detailed_pipeline_result.get('filename', ''),
            'chunk_count': 0,
        }

    log.info(f"Loading detailed parse result from {data_file}")
    with open(data_file, 'rb') as f:
        parse_result = pickle.load(f)

    # parse_to_markdown returns {sheets, image_chunks, flowchart_chunks, color_map, cell_values_map}
    image_chunks = parse_result.get('image_chunks', [])

    if not image_chunks:
        log.info("No image chunks to upload")
        return {
            'uploaded_count': 0,
            'success': True,
            'data_file': data_file,
            'document_id': detailed_pipeline_result.get('document_id'),
            'filename': detailed_pipeline_result.get('filename', ''),
        }

    # Access metadata via attribute (ChunkedElementDTO) or dict key
    def get_meta(chunk):
        if hasattr(chunk, 'metadata'):
            return chunk.metadata
        elif isinstance(chunk, dict):
            return chunk.get('metadata', {})
        return {}

    # Count images with base64 data
    image_count = sum(
        1 for chunk in image_chunks
        if 'image_base64' in get_meta(chunk)
    )

    log.info(f"Found {image_count} images to upload")

    if image_count == 0:
        log.info("No images with base64 data found")
        return {
            'uploaded_count': 0,
            'success': True,
            'data_file': data_file,
            'document_id': detailed_pipeline_result.get('document_id'),
            'filename': detailed_pipeline_result.get('filename', ''),
        }

    # Upload images one at a time, freeing base64 data immediately after upload
    log.info("Uploading images to SeaweedFS...")
    service = SeaweedFSService()
    uploaded_count = 0

    for chunk in image_chunks:
        metadata = get_meta(chunk)
        if 'image_base64' not in metadata:
            continue

        image_number = metadata.get('image_number', 1)
        sheet_name = metadata.get('sheet_name', 'Sheet1')
        mime_type = metadata.get('mime_type', 'image/png')

        result = service.upload_image(
            image_base64=metadata['image_base64'],
            document_id=doc_uuid,
            image_number=image_number,
            sheet_name=sheet_name,
            mime_type=mime_type,
            tenant_id=tenant_id,
            node_id=node_id,
        )

        if result['success']:
            metadata['image_s3_key'] = result['image_s3_key']
            metadata['image_url'] = result['image_url']
            # Remove base64 immediately to free memory
            del metadata['image_base64']
            uploaded_count += 1
            log.info(f"Uploaded image {image_number}: {result['image_s3_key']}")
        else:
            # Still remove base64 to avoid OOM on large files
            metadata.pop('image_base64', None)
            log.warning(f"Failed to upload image {image_number}: {result.get('error', 'Unknown')}")

    log.info(f"Successfully uploaded {uploaded_count}/{image_count} images to SeaweedFS")

    # Save updated parse result back (image base64 removed, URLs added)
    with open(data_file, 'wb') as f:
        pickle.dump(parse_result, f)

    log.info(f"Updated parse result saved to {data_file}")

    return {
        'uploaded_count': uploaded_count,
        'total_images': image_count,
        'success': True,
        'data_file': data_file,
        'document_id': detailed_pipeline_result.get('document_id'),
        'filename': detailed_pipeline_result.get('filename'),
    }


@task.external_python(python=EXTERNAL_PYTHON)
def store_vectors_task(
    document_id: str,
    embedded_result: Dict[str, Any],
    pg_records_result: Dict[str, Any] = None,
    tenant_id: str = "",
    node_id: str = "",
) -> Dict[str, Any]:
    """
    Store vectors in Milvus.

    Args:
        document_id: Document UUID string
        embedded_result: Result from generate_embeddings task (may contain data_file path)
        pg_records_result: Result from save_extracted_records_task with record_ids_file.
            If provided, uses PostgreSQL record IDs for Milvus linking.
            If None, generates new UUIDs (legacy behavior, no PostgreSQL linking).
        tenant_id: Tenant UUID string for Milvus data isolation
        node_id: Node UUID string for Milvus data isolation

    Returns:
        Storage result dict
    """
    from uuid import UUID, uuid4
    from src.knowledge.vector_store import VectorStore
    from src.resilience.breakers import async_breaker_call, milvus_breaker
    import asyncio
    import logging
    import pickle
    import os

    log = logging.getLogger(__name__)

    doc_uuid = UUID(document_id)

    if not embedded_result:
        log.warning("No embedded result provided")
        return {'stored_count': 0, 'success': False}

    # Load data from file if data_file path is provided (large data case)
    data_file = embedded_result.get('data_file', '')
    if data_file and os.path.exists(data_file):
        log.info(f"Loading embedded data from {data_file}")
        with open(data_file, 'rb') as f:
            embedded_result = pickle.load(f)
        # Clean up the temp file after loading
        try:
            os.remove(data_file)
            log.info(f"Cleaned up temp file: {data_file}")
        except Exception as e:
            log.warning(f"Failed to clean up temp file {data_file}: {e}")

    chunks = embedded_result.get('chunks', [])
    embeddings = embedded_result.get('embeddings', [])
    domain_data = embedded_result.get('domain_data', [])  # Extract domain data

    # DEBUG: Log domain_data status
    log.info(f"DEBUG store_vectors: embedded_result keys = {list(embedded_result.keys())}")
    log.info(f"DEBUG store_vectors: domain_data type = {type(domain_data)}, len = {len(domain_data) if domain_data else 0}")

    if not chunks or not embeddings:
        log.info("No chunks or embeddings to store")
        return {'stored_count': 0, 'success': True}

    if len(chunks) != len(embeddings):
        log.error(f"Mismatch: {len(chunks)} chunks vs {len(embeddings)} embeddings")
        return {'stored_count': 0, 'success': False}

    # Log domain data statistics
    if domain_data:
        category_counts = {}
        for data in domain_data:
            category = data.get('domain_category', 'general')
            category_counts[category] = category_counts.get(category, 0) + 1
        log.info(f"Domain category distribution: {category_counts}")

    log.info(f"Storing {len(chunks)} vectors to Milvus")

    # Load PostgreSQL record IDs if available (for proper PostgreSQL-Milvus linking)
    record_ids = None
    record_ids_file = pg_records_result.get('record_ids_file', '') if pg_records_result else ''
    if record_ids_file and os.path.exists(record_ids_file):
        log.info(f"Loading PostgreSQL record IDs from {record_ids_file}")
        with open(record_ids_file, 'rb') as f:
            record_ids = pickle.load(f)
        # Clean up the record IDs file
        try:
            os.remove(record_ids_file)
            log.info(f"Cleaned up record IDs file: {record_ids_file}")
        except Exception as e:
            log.warning(f"Failed to clean up record IDs file {record_ids_file}: {e}")

        if len(record_ids) != len(chunks):
            log.error(
                f"Mismatch: {len(record_ids)} PostgreSQL record IDs vs {len(chunks)} chunks. "
                "Falling back to generated UUIDs (no PostgreSQL linking)."
            )
            record_ids = None
        else:
            log.info(f"Using {len(record_ids)} PostgreSQL record IDs for Milvus linking")

    # Fall back to generating new UUIDs if no PostgreSQL IDs available
    if record_ids is None:
        log.warning("No PostgreSQL record IDs available - generating new UUIDs (no PostgreSQL linking)")
        record_ids = [uuid4() for _ in chunks]

    # Determine document type from embedded_result (passed through from chunk task)
    doc_type = embedded_result.get('doc_type') or 'text'
    original_filename = embedded_result.get('original_filename', '')

    # Convert Excel dict chunks to ExtractedRecord objects for VectorStore compatibility
    if doc_type == 'excel':
        from src.extraction.models import ExtractedRecord
        records = []
        for chunk in chunks:
            if isinstance(chunk, dict):
                content = chunk.get('content', {})
                source = chunk.get('_source', {})
                # Ensure all required _source fields are present
                complete_source = {
                    'document_id': source.get('document_id', str(doc_uuid)),
                    'filename': source.get('filename', ''),
                    'sheet': source.get('sheet', 'Sheet1'),
                    'table_id': source.get('table_id', 0),
                    'row': source.get('row', 0),
                    'col_range': source.get('col_range', ''),
                }
                # Skip if content is empty
                if not content:
                    log.debug(f"Skipping chunk with empty content")
                    continue
                record = ExtractedRecord(
                    content=content,
                    headers=chunk.get('headers', list(content.keys())),
                    _source=complete_source,
                    _merge_metadata=chunk.get('_merge_metadata'),
                    resolved_content=chunk.get('resolved_content'),
                    _legacy_metadata=chunk.get('_legacy_metadata'),
                    tags=chunk.get('tags', []),
                )
                records.append(record)
            else:
                records.append(chunk)
    else:
        records = chunks

    async def store():
        store = VectorStore()
        count = await async_breaker_call(
            milvus_breaker,
            store.index_with_vectors,
            records=records,
            vectors=embeddings,
            document_id=doc_uuid,
            record_ids=record_ids,
            document_type=doc_type,
            filename=original_filename,
            domain_data=domain_data if domain_data else None,
            tenant_id=tenant_id,
            node_id=node_id,
        )
        return count

    stored_count = asyncio.run(store())
    log.info(f"Successfully stored {stored_count} vectors")

    return {
        'stored_count': stored_count,
        'success': True,
    }

@task.external_python(python=EXTERNAL_PYTHON)
def save_document_metadata_task(
    document_id: str,
    fetch_result: Dict[str, Any],
    validation_result: Dict[str, Any],
    tenant_id: str,
    node_id: str,
    correlation_id: str,
) -> Dict[str, Any]:
    """
    Create Document record in PostgreSQL.

    Args:
        document_id: Document UUID string
        fetch_result: Result from fetch_document task with file info
        validation_result: Result from validate_event task
        tenant_id: Tenant UUID string for data isolation
        node_id: Node UUID string for data isolation
        correlation_id: Airflow dag_run_id, for pipeline-event correlation

    Returns:
        Dict with document_id and success status
    """
    import asyncio
    import logging
    from uuid import UUID
    from datetime import datetime
    from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
    from sqlalchemy.orm import sessionmaker
    from src.db.models import Document
    from src.core.config import settings
    from src.services.pipeline_events import publish_stage_event

    log = logging.getLogger(__name__)

    publish_stage_event(
        "pending", "started",
        document_id=document_id,
        tenant_id=tenant_id,
        node_id=node_id,
        correlation_id=correlation_id,
    )

    # Extract file info
    filename = fetch_result.get('filename', validation_result.get('object_key', 'unknown'))
    file_path = validation_result.get('object_key', '')
    file_size = fetch_result.get('size', 0)

    log.info(f"Creating document record: {document_id}, filename: {filename}")
    
    async def save_document():
        # Create async engine
        engine = create_async_engine(
            settings.database_url,
            echo=False,
            pool_pre_ping=True,
        )
        
        async_session = sessionmaker(
            engine,
            class_=AsyncSession,
            expire_on_commit=False,
        )
        
        async with async_session() as session:
            try:
                # Set RLS tenant context for row-level security policy
                from sqlalchemy import text as sa_text
                if tenant_id:
                    from uuid import UUID as _UUID
                    _tid = str(_UUID(tenant_id))  # validate UUID format
                    await session.execute(sa_text(f"SET LOCAL app.tenant_id = '{_tid}'"))

                # Upsert: if upload endpoint already created the record (status=pending),
                # update it to processing. Otherwise insert a new record (legacy CLI flow).
                from sqlalchemy.dialects.postgresql import insert as pg_insert

                stmt = pg_insert(Document).values(
                    id=UUID(document_id),
                    filename=filename,
                    file_path=file_path,
                    file_size_bytes=file_size,
                    status='processing',
                    created_at=datetime.utcnow(),
                    tenant_id=UUID(tenant_id) if tenant_id else None,
                    node_id=UUID(node_id) if node_id else None,
                ).on_conflict_do_update(
                    index_elements=['id'],
                    set_={
                        'status': 'processing',
                    },
                )
                await session.execute(stmt)
                await session.commit()

                log.info(f"Document record upserted successfully: {document_id}")
                return True

            except Exception as e:
                await session.rollback()
                log.error(f"Failed to upsert document record: {e}")
                raise
            finally:
                await engine.dispose()
    
    success = asyncio.run(save_document())

    if success:
        publish_stage_event(
            "processing", "started",
            document_id=document_id,
            tenant_id=tenant_id,
            node_id=node_id,
            correlation_id=correlation_id,
        )

    return {
        'document_id': document_id,
        'success': success,
    }


@task.external_python(python=EXTERNAL_PYTHON)
def save_extracted_records_task(
    document_id: str,
    chunks_result: Dict[str, Any],
) -> Dict[str, Any]:
    """
    Save ExtractedRecords to PostgreSQL using StructuredStore.

    This task MUST run BEFORE store_vectors_task so that record IDs exist
    in PostgreSQL before Milvus references them.

    Args:
        document_id: Document UUID string
        chunks_result: Result containing chunks with embeddings

    Returns:
        Dict with record_ids_file (path to pickled UUIDs) and record_count.
        The record_ids are the PostgreSQL primary keys that must be used
        by store_vectors_task to link Milvus vectors to PostgreSQL records.
    """
    import asyncio
    import logging
    import pickle
    import os
    from uuid import UUID
    from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
    from sqlalchemy.orm import sessionmaker
    from src.knowledge.structured_store import StructuredStore
    from src.extraction.models import ExtractedRecord
    from src.core.config import settings

    log = logging.getLogger(__name__)

    # Load chunks from pickle if data_file is provided
    data_file = chunks_result.get('data_file')
    if data_file and os.path.exists(data_file):
        log.info(f"Loading chunks from {data_file}")
        with open(data_file, 'rb') as f:
            chunks_result = pickle.load(f)

    chunks = chunks_result.get('chunks', [])

    if not chunks:
        log.warning(f"No chunks to save for document {document_id}")
        return {
            'record_ids_file': '',  # Empty string for Airflow 3.x XCom compatibility
            'record_count': 0,
            'success': True,
            'document_id': document_id,
        }

    log.info(f"Saving {len(chunks)} records to PostgreSQL for document {document_id}")

    async def save_records():
        # Create async engine
        engine = create_async_engine(
            settings.database_url,
            echo=False,
            pool_pre_ping=True,
        )

        async_session = sessionmaker(
            engine,
            class_=AsyncSession,
            expire_on_commit=False,
        )

        async with async_session() as session:
            try:
                # Convert chunks to ExtractedRecord objects
                records = []

                for chunk_idx, chunk in enumerate(chunks):
                    if isinstance(chunk, dict):
                        # Handle text documents (Word/PDF/PowerPoint) which have 'text' field
                        # vs Excel documents which have 'content' dict
                        if 'text' in chunk and 'content' not in chunk:
                            # Text document chunk - convert to ExtractedRecord format
                            text_content = chunk.get('text', '')
                            if not text_content.strip():
                                log.debug(f"Skipping empty text chunk at index {chunk_idx}")
                                continue
                            content = {'text': text_content}
                            headers = ['text']
                            # Use metadata from text chunk
                            metadata = chunk.get('metadata', {})
                            source = {
                                'document_id': document_id,
                                'filename': metadata.get('source_filename', chunks_result.get('original_filename', 'unknown')),
                                'sheet': 'Page',  # Use 'Page' for text documents
                                'table_id': metadata.get('page_number', 1) or 1,
                                'row': chunk_idx,
                                'col_range': '',
                            }
                        else:
                            # Excel chunk - use existing format
                            content = chunk.get('content', {})
                            headers = chunk.get('headers', list(content.keys()) if content else [])
                            source = chunk.get('_source', {})

                            # Ensure source has required fields
                            if 'document_id' not in source:
                                source['document_id'] = document_id
                            if 'filename' not in source:
                                source['filename'] = chunks_result.get('original_filename', 'unknown')
                            if 'sheet' not in source:
                                source['sheet'] = 'Sheet1'
                            if 'table_id' not in source:
                                source['table_id'] = 1
                            if 'row' not in source:
                                source['row'] = chunk_idx
                            if 'col_range' not in source:
                                source['col_range'] = ''

                        # Skip if content is empty
                        if not content:
                            log.debug(f"Skipping chunk {chunk_idx} with empty content")
                            continue

                        record = ExtractedRecord(
                            content=content,
                            headers=headers,
                            _source=source,
                            resolved_content=chunk.get('resolved_content'),
                            tags=chunk.get('tags', []),
                        )
                        records.append(record)
                    elif isinstance(chunk, ExtractedRecord):
                        records.append(chunk)

                # Use StructuredStore to save records - now returns StoreResult with IDs
                store = StructuredStore(session)
                result = await store.store_records(
                    records=records,
                    document_id=UUID(document_id),
                )

                log.info(f"Saved {result.count} records to PostgreSQL with {len(result.record_ids)} IDs")

                return result.record_ids, result.count

            except Exception as e:
                await session.rollback()
                log.error(f"Failed to save records: {e}")
                raise
            finally:
                await engine.dispose()

    record_ids, saved_count = asyncio.run(save_records())

    # Save record_ids to pickle file (can be large for many records)
    # These IDs link PostgreSQL records to Milvus vectors
    record_ids_file = ''
    if record_ids:
        record_ids_dir = '/opt/airflow/logs/record_ids'
        os.makedirs(record_ids_dir, exist_ok=True)
        record_ids_file = os.path.join(record_ids_dir, f'record_ids_{document_id}.pkl')
        with open(record_ids_file, 'wb') as f:
            pickle.dump(record_ids, f)
        log.info(f"Saved {len(record_ids)} record IDs to {record_ids_file}")

    return {
        'record_ids_file': record_ids_file,
        'record_count': saved_count,
        'success': True,
        'document_id': document_id,
    }


@task.external_python(python=EXTERNAL_PYTHON)
def update_document_status_task(
    document_id: str,
    outcome: Dict[str, Any],
    record_count: int = 0,
    tenant_id: str = '',
    node_id: str = '',
    correlation_id: str = '',
) -> Dict[str, Any]:
    """
    Write terminal document status to PostgreSQL and emit the terminal pipeline event.

    Status is resolved upstream by ``resolve_terminal_outcome`` (regular @task in
    the DAG file, which has access to Airflow's task-instance state via context).

    Args:
        document_id: Document UUID string
        outcome: Dict from resolve_terminal_outcome with keys ``status`` ('completed'|'failed')
                 and ``error`` (str | None).
        record_count: Number of records extracted
        tenant_id: Tenant UUID string for RLS context + event routing
        node_id: Node UUID string for event routing
        correlation_id: Airflow dag_run_id, for event correlation

    Returns:
        Dict with document_id, status, success
    """
    import asyncio
    import logging
    from uuid import UUID
    from datetime import datetime
    from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
    from sqlalchemy.orm import sessionmaker
    from sqlalchemy import update, text as sa_text
    from src.db.models import Document
    from src.core.config import settings
    from src.services.pipeline_events import publish_stage_event

    log = logging.getLogger(__name__)

    outcome = outcome or {}
    status = outcome.get('status', 'failed')
    error_msg = outcome.get('error') if status != 'completed' else None
    if status != 'completed' and not error_msg:
        error_msg = 'terminal outcome unavailable'
    log.info(f"Updating document {document_id} status to {status} (error={error_msg})")

    async def update_document():
        engine = create_async_engine(
            settings.database_url,
            echo=False,
            pool_pre_ping=True,
        )
        async_session = sessionmaker(
            engine,
            class_=AsyncSession,
            expire_on_commit=False,
        )
        async with async_session() as session:
            try:
                if tenant_id:
                    _tid = str(UUID(tenant_id))  # validate UUID format
                    await session.execute(sa_text(f"SET LOCAL app.tenant_id = '{_tid}'"))

                stmt = (
                    update(Document)
                    .where(Document.id == UUID(document_id))
                    .values(
                        status=status,
                        record_count=record_count,
                        processed_at=datetime.utcnow() if status == 'completed' else None,
                    )
                )
                await session.execute(stmt)
                await session.commit()
                log.info(f"Document {document_id} updated successfully with status={status}")
                return True
            except Exception as e:
                await session.rollback()
                log.error(f"Failed to update document: {e}")
                raise
            finally:
                await engine.dispose()

    success = asyncio.run(update_document())

    publish_stage_event(
        'completed' if status == 'completed' else 'failed',
        'succeeded' if status == 'completed' else 'failed',
        document_id=document_id,
        tenant_id=tenant_id,
        node_id=node_id,
        correlation_id=correlation_id,
        **({'error': error_msg} if error_msg else {}),
    )

    return {
        'document_id': document_id,
        'status': status,
        'success': success,
    }


@task.external_python(python=EXTERNAL_PYTHON)
def hierarchical_chunk_excel_detailed_task(
    document_id: str,
    parse_result: Dict[str, Any],
    chunk_sizes: List[int] = None,
) -> Dict[str, Any]:
    """
    Hierarchically chunk Excel Detailed pipeline output using LlamaIndex.

    Accepts per-sheet markdown from parse_to_markdown. Creates one
    LlamaIndex Document per sheet (preserving sheet_name in metadata),
    then applies a 2-level hierarchy via the mermaid-aware Excel parser map:

        MarkdownNodeParser           -> sections (parents, by ATX headers)
        MermaidAwareLineSplitter(512)-> leaves (line-based, mermaid atomic)

    Mermaid fences stay in-context (no orphan leaves); pure-prose leaves
    stay <= 512 tokens, mermaid-bearing leaves may grow up to the embedding
    model context. Header-only nodes are pruned post-parse.

    Image and flowchart content remain flat leaf-level TextNodes.

    Args:
        document_id: Document UUID string
        parse_result: Result from upload_images_to_seaweedfs (contains data_file path
            to pickle with {sheets, image_chunks, flowchart_chunks, color_map, cell_values_map})
        chunk_sizes: DEPRECATED -- ignored. Excel branch now uses a fixed
            2-level mermaid-aware hierarchy. Kept for back-compat with callers
            that pass it.

    Returns:
        Dict with nodes_file path and metadata
    """
    import logging
    import pickle
    import os
    from uuid import UUID

    from llama_index.core import Document
    from llama_index.core.node_parser import HierarchicalNodeParser, get_leaf_nodes

    from src.core.config import Settings
    from src.extraction_v2.mermaid_chunking import (
        EXCEL_PARSER_IDS,
        build_excel_parser_map,
        prune_header_only_nodes,
    )
    from src.llm.factory import get_token_counter, get_tokenizer

    log = logging.getLogger(__name__)

    if chunk_sizes is not None:
        log.warning(
            f"chunk_sizes={chunk_sizes} is deprecated for the Excel branch; "
            "ignoring and using the fixed mermaid-aware 2-level hierarchy."
        )

    doc_uuid = UUID(document_id)

    if not parse_result:
        log.warning("No parse result provided")
        return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0}

    # Load data from pickle file
    data_file = parse_result.get('data_file')
    if not data_file or not os.path.exists(data_file):
        log.warning(f"Data file not found: {data_file}")
        return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0}

    log.info(f"Loading Excel Detailed parse result from {data_file}")
    with open(data_file, 'rb') as f:
        parsed_data = pickle.load(f)

    sheets = parsed_data.get('sheets', [])
    image_chunks = parsed_data.get('image_chunks', [])
    flowchart_chunks = parsed_data.get('flowchart_chunks', [])
    color_map = parsed_data.get('color_map', {})
    cell_values_map = parsed_data.get('cell_values_map', {})
    original_filename = parse_result.get('filename', 'document.xlsx')

    if not sheets and not image_chunks and not flowchart_chunks:
        log.warning("No content to chunk")
        return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0}

    log.info(f"Hierarchical chunking Excel Detailed: {original_filename}")
    log.info(f"Input: {len(sheets)} sheets, {len(image_chunks)} images, {len(flowchart_chunks)} flowcharts")

    tokenizer = get_tokenizer()
    token_count = get_token_counter()
    embed_max_tokens = Settings().embedding_model_context_tokens
    # Headroom against the 8192-token embed ceiling: covers retained
    # metadata (~100 tok template + filename/sheet_name/document_id)
    # plus Azure-tokenizer variance vs local cl100k_base counts.
    EMBED_SAFETY_BUFFER = 400
    chunker_max = embed_max_tokens - EMBED_SAFETY_BUFFER
    log.info(
        f"Excel hierarchy: 2-level [markdown, chunk_size_512], "
        f"embed_max={embed_max_tokens} chunker_max={chunker_max} buffer={EMBED_SAFETY_BUFFER}"
    )

    node_parser_map = build_excel_parser_map(
        tokenize_fn=token_count,
        embed_max_tokens=chunker_max,
    )
    node_parser = HierarchicalNodeParser(
        node_parser_ids=EXCEL_PARSER_IDS,
        node_parser_map=node_parser_map,
    )

    # Create one LlamaIndex Document per sheet from markdown
    documents = []
    for sheet_info in sheets:
        sheet_name = sheet_info['sheet_name']
        markdown = sheet_info['markdown']

        if not markdown or not markdown.strip():
            continue

        doc_metadata = {
            "document_id": str(doc_uuid),
            "filename": original_filename,
            "document_type": "excel",
            "sheet_name": sheet_name,
        }

        doc = Document(
            doc_id=str(doc_uuid),
            text=markdown,
            metadata=doc_metadata,
        )
        documents.append(doc)
        log.info(f"Sheet '{sheet_name}': {token_count(markdown)} tokens")

    # Add image descriptions as separate Documents (leaf-level, not hierarchically chunked)
    def _dto_to_dict(chunk):
        if hasattr(chunk, 'model_dump'):
            return chunk.model_dump()
        elif hasattr(chunk, 'dict'):
            return chunk.dict()
        elif hasattr(chunk, '__dict__'):
            return {'text': getattr(chunk, 'text', ''), 'metadata': getattr(chunk, 'metadata', {}),
                    'element_type': getattr(chunk, 'element_type', 'Image')}
        return chunk

    image_documents = []
    for img_chunk in image_chunks:
        img = _dto_to_dict(img_chunk)
        text = img.get('text', '')
        if not text or not text.strip():
            continue
        meta = img.get('metadata', {})
        doc_metadata = {
            "document_id": str(doc_uuid),
            "filename": original_filename,
            "document_type": "excel",
            "sheet_name": meta.get('sheet_name', 'Unknown'),
            "element_type": "Image",
        }
        for key in ['image_url', 'image_s3_key', 'image_number', 'position_key']:
            if key in meta:
                doc_metadata[key] = meta[key]
        image_documents.append(Document(doc_id=str(doc_uuid), text=text, metadata=doc_metadata))

    flowchart_documents = []
    for fc_chunk in flowchart_chunks:
        fc = _dto_to_dict(fc_chunk)
        text = fc.get('text', '')
        if not text or not text.strip():
            continue
        meta = fc.get('metadata', {})
        doc_metadata = {
            "document_id": str(doc_uuid),
            "filename": original_filename,
            "document_type": "excel",
            "sheet_name": meta.get('sheet_name', 'Unknown'),
            "element_type": "Flowchart",
        }
        flowchart_documents.append(Document(doc_id=str(doc_uuid), text=text, metadata=doc_metadata))

    if not documents and not image_documents and not flowchart_documents:
        log.warning("No content to chunk after processing")
        return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0}

    log.info(f"Created {len(documents)} sheet Documents for hierarchical chunking")

    # Apply hierarchical parsing to sheet documents
    nodes = []
    if documents:
        nodes = node_parser.get_nodes_from_documents(documents)
        log.info(f"Created {len(nodes)} hierarchical nodes from {len(documents)} sheets")
        nodes = list(prune_header_only_nodes(nodes))
        log.info(f"After header-only prune: {len(nodes)} nodes")

    # Add image/flowchart as flat TextNodes (no hierarchy needed for these)
    from llama_index.core.schema import TextNode, NodeRelationship, RelatedNodeInfo
    doc_id_str = str(doc_uuid)
    for doc in image_documents + flowchart_documents:
        leaf_node = TextNode(
            text=doc.text,
            metadata=doc.metadata,
        )
        leaf_node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=doc_id_str)
        nodes.append(leaf_node)

    # Count leaf nodes
    leaf_nodes = get_leaf_nodes(nodes)
    mermaid_bearing_leaves = sum(
        1 for n in leaf_nodes if n.metadata.get("has_mermaid")
    )
    log.info(
        f"Total nodes: {len(nodes)}, leaf nodes: {len(leaf_nodes)}, "
        f"mermaid-bearing leaves: {mermaid_bearing_leaves}"
    )

    leaf_ids = {l.node_id for l in leaf_nodes}

    # Ensure ALL nodes have ref_doc_id set (used as doc_id in Milvus)
    for node in nodes:
        if not node.ref_doc_id:
            node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=doc_id_str)

        node.metadata["token_count"] = token_count(node.text)
        node.metadata["char_count"] = len(node.text)

        # cell_colors is only attached to non-leaf (parent / sheet-level)
        # nodes. Leaves go to Milvus, where the dynamic field is capped at
        # 64KB; a single highlighted sheet can hold 10k+ colored cells
        # (~200KB JSON) and would blow that cap on every leaf. Parents go
        # to PostgresDocumentStore with no such limit, and the
        # AutoMergingRetriever pulls them at query time so highlight-aware
        # consumers still see cell_colors via the parent chain.
        # cell_values is dropped entirely: it's a much larger duplicate
        # of the markdown content and nothing in this repo reads it back.
        if node.node_id not in leaf_ids:
            sheet_name = node.metadata.get("sheet_name")
            if sheet_name and sheet_name in color_map:
                node.metadata["cell_colors"] = color_map[sheet_name]

    # Per-leaf token audit -- including the metadata that actually ends up
    # in the embed text -- so any leaf that exceeds the embed ceiling is
    # visible in the chunking log instead of only at the storage stage.
    from llama_index.core.schema import MetadataMode
    embed_max = Settings().embedding_model_context_tokens
    leaf_sizes = []
    for idx, leaf in enumerate(leaf_nodes):
        embed_text = leaf.get_content(metadata_mode=MetadataMode.EMBED)
        tot = token_count(embed_text)
        text_tok = token_count(leaf.text)
        leaf_sizes.append((idx, tot, text_tok))
        if tot > embed_max:
            log.warning(
                f"leaf[{idx}] OVER BUDGET total={tot} text_tok={text_tok} "
                f"excluded_keys={leaf.excluded_embed_metadata_keys} "
                f"sheet={leaf.metadata.get('sheet_name')!r}"
            )
    leaf_sizes.sort(key=lambda x: x[1], reverse=True)
    log.info(
        f"Leaf embed-text token sizes (top 5): {leaf_sizes[:5]} "
        f"(embed_max={embed_max})"
    )

    # Save nodes to pickle file
    nodes_dir = '/opt/airflow/logs/hierarchical_nodes'
    os.makedirs(nodes_dir, exist_ok=True)
    nodes_file = os.path.join(nodes_dir, f'nodes_excel_detailed_{document_id}.pkl')

    nodes_data = {
        'nodes': nodes,
        'node_count': len(nodes),
        'leaf_count': len(leaf_nodes),
        'doc_type': 'excel',
        'original_filename': original_filename,
        'document_id': document_id,
        'parser_ids': EXCEL_PARSER_IDS,
        'mermaid_bearing_leaves': mermaid_bearing_leaves,
    }

    with open(nodes_file, 'wb') as f:
        pickle.dump(nodes_data, f)

    log.info(f"Saved {len(nodes)} nodes to {nodes_file}")

    return {
        'nodes_file': nodes_file,
        'node_count': len(nodes),
        'leaf_count': len(leaf_nodes),
        'doc_type': 'excel',
        'original_filename': original_filename,
        'document_id': document_id,
    }


@task.external_python(python=EXTERNAL_PYTHON)
def store_to_llamaindex_task(
    document_id: str,
    nodes_result: Dict[str, Any],
    tenant_id: str,
    node_id: str,
    correlation_id: str,
) -> Dict[str, Any]:
    """
    Store hierarchical nodes to LlamaIndex MilvusVectorStore + PostgresDocumentStore.

    Pattern: ALL nodes → PostgresDocumentStore, LEAF nodes only → Milvus.
    This enables AutoMergingRetriever to expand context at query time.

    Args:
        document_id: Document UUID string
        nodes_result: Result from hierarchical_chunk task (contains nodes_file path)
        tenant_id: Tenant UUID for pipeline-event routing
        node_id: Node UUID for pipeline-event routing
        correlation_id: Airflow dag_run_id, for event correlation

    Returns:
        Dict with stored counts and success status
    """
    import asyncio
    import logging
    import pickle
    import os
    import sys
    from uuid import UUID

    from llama_index.core import VectorStoreIndex, StorageContext, Settings
    from llama_index.core.node_parser import get_leaf_nodes
    from llama_index.vector_stores.milvus import MilvusVectorStore
    from llama_index.vector_stores.milvus.utils import BM25BuiltInFunction
    from llama_index.storage.docstore.postgres import PostgresDocumentStore
    from src.core.config import settings
    from src.llm import get_embed_model
    from src.resilience.breakers import breaker_call, milvus_breaker
    from src.services.pipeline_events import publish_stage_event

    if '/opt/airflow/dags' not in sys.path:
        sys.path.insert(0, '/opt/airflow/dags')
    from tasks.chunk_schema import CHUNK_SCHEMA_VERSION

    log = logging.getLogger(__name__)

    publish_stage_event(
        "indexing", "started",
        document_id=document_id,
        tenant_id=tenant_id,
        node_id=node_id,
        correlation_id=correlation_id,
    )

    doc_uuid = UUID(document_id)

    if not nodes_result:
        log.warning("No nodes result provided")
        return {'stored_count': 0, 'leaf_count': 0, 'success': False}

    # Load nodes from pickle file
    nodes_file = nodes_result.get('nodes_file')
    if not nodes_file or not os.path.exists(nodes_file):
        log.warning(f"Nodes file not found: {nodes_file}")
        return {'stored_count': 0, 'leaf_count': 0, 'success': False}

    log.info(f"Loading nodes from {nodes_file}")
    with open(nodes_file, 'rb') as f:
        nodes_data = pickle.load(f)

    nodes = nodes_data.get('nodes', [])
    if not nodes:
        log.warning("No nodes to store")
        return {'stored_count': 0, 'leaf_count': 0, 'success': True}

    # Inject tenant/node context + chunk schema version into all node metadata
    for node in nodes:
        node.metadata['tenant_id'] = tenant_id
        node.metadata['node_id'] = node_id
        node.metadata['schema_version'] = CHUNK_SCHEMA_VERSION
        # TEXTIQ-569: ensure key is always present so the Milvus dynamic
        # envelope carries it on every new chunk (tag-cascade filter on
        # missing key would silently no-match). tag_records_task already
        # populates it; this is the safety default for paths that skip
        # tagging (e.g., tagging failure, empty text).
        node.metadata.setdefault('verified_tag_names', [])

    # Constrain the per-node embed text to semantic-only metadata. Excel
    # leaves can be ~8k tokens already; serialising bulk metadata such as
    # cell_colors, cell_values, the tag list, or routing IDs blows past
    # the 8192-token embedding ceiling. Whitelist the few keys that
    # genuinely help retrieval, then exclude everything else.
    EMBED_INCLUDE_KEYS = {
        "filename", "sheet_name", "document_type", "element_type",
    }
    for node in nodes:
        excl = [k for k in node.metadata.keys() if k not in EMBED_INCLUDE_KEYS]
        node.excluded_embed_metadata_keys = excl
        node.excluded_llm_metadata_keys = excl

    log.info(f"Storing {len(nodes)} nodes to LlamaIndex stores")

    async def store_nodes():
        """Async wrapper for LlamaIndex operations (MilvusVectorStore uses async internally)."""
        # Configure embedding model via provider factory
        Settings.embed_model = get_embed_model()

        # Initialize MilvusVectorStore with hybrid search (dense + BM25 sparse)
        milvus_uri = f"http://{settings.milvus_host}:{settings.milvus_port}"
        log.info(f"Connecting to Milvus: {milvus_uri}")

        vector_store = MilvusVectorStore(
            uri=milvus_uri,
            collection_name=settings.milvus_collection,
            dim=settings.embedding_dimensions,
            overwrite=False,  # Don't drop existing data
            enable_sparse=True,
            sparse_embedding_function=BM25BuiltInFunction(),
            hybrid_ranker="RRFRanker",
            hybrid_ranker_params={"k": 60},
        )

        # Initialize PostgresDocumentStore
        log.info("Connecting to PostgreSQL docstore")

        # Create engines with RLS tenant context set on every connection
        from uuid import UUID as _UUID
        from sqlalchemy import create_engine as _create_engine, event as sa_event
        from sqlalchemy.ext.asyncio import create_async_engine as _create_async_engine

        _tid = str(_UUID(tenant_id))  # validate UUID
        _nid = str(_UUID(node_id))   # validate UUID
        db_url = settings.database_url
        db_url_sync = db_url.replace("+asyncpg", "")

        # asyncpg: use server_settings to set tenant context at connection time
        async_engine = _create_async_engine(
            db_url,
            connect_args={"server_settings": {"app.tenant_id": _tid, "app.node_id": _nid}},
        )
        # psycopg2: use connect event to set tenant context
        sync_engine = _create_engine(db_url_sync)

        @sa_event.listens_for(sync_engine, "connect")
        def _set_tenant(dbapi_conn, connection_record):
            cursor = dbapi_conn.cursor()
            cursor.execute(f"SET app.tenant_id = '{_tid}'; SET app.node_id = '{_nid}'")
            cursor.close()

        from llama_index.storage.kvstore.postgres import PostgresKVStore
        kvstore = PostgresKVStore(
            table_name=settings.docstore_table,
            engine=sync_engine,
            async_engine=async_engine,
            perform_setup=True,
            use_jsonb=True,
        )
        docstore = PostgresDocumentStore(postgres_kvstore=kvstore, namespace=settings.docstore_namespace)

        # Create storage context
        storage_context = StorageContext.from_defaults(
            vector_store=vector_store,
            docstore=docstore,
        )

        # Add ALL nodes to docstore (for parent lookup during retrieval)
        log.info(f"Adding {len(nodes)} nodes to PostgresDocumentStore")
        storage_context.docstore.add_documents(nodes)

        # Get leaf nodes for vector indexing
        leaf_nodes = get_leaf_nodes(nodes)
        log.info(f"Indexing {len(leaf_nodes)} leaf nodes to Milvus")

        # Create index with only leaf nodes (vectors stored in Milvus).
        # Wrapped with milvus_breaker — VectorStoreIndex(...) executes the
        # underlying Milvus inserts during construction.
        index = breaker_call(
            milvus_breaker,
            VectorStoreIndex,
            leaf_nodes,
            storage_context=storage_context,
            show_progress=True,
        )

        log.info(f"Successfully stored {len(nodes)} nodes ({len(leaf_nodes)} leaf nodes)")
        return len(nodes), len(leaf_nodes)

    # Run async operations
    stored_count, leaf_count = asyncio.run(store_nodes())

    # Clean up nodes file
    try:
        os.remove(nodes_file)
        log.info(f"Cleaned up nodes file: {nodes_file}")
    except Exception as e:
        log.warning(f"Failed to clean up nodes file: {e}")

    return {
        'stored_count': stored_count,
        'leaf_count': leaf_count,
        'document_id': document_id,
        'success': True,
    }


@task.external_python(python=EXTERNAL_PYTHON)
def store_large_table_to_llamaindex_task(
    document_id: str,
    nl_result: Dict[str, Any],
    tenant_id: str,
    node_id: str,
    correlation_id: str,
) -> Dict[str, Any]:
    """
    Store Large Table records to LlamaIndex MilvusVectorStore.

    For Large Table pipeline: 1 record = 1 Document (no hierarchical parsing).
    Text = NL description, metadata.record_json = original JSON.

    Args:
        document_id: Document UUID string
        nl_result: Result from convert_large_tables_to_natural_language task
        tenant_id: Tenant UUID for pipeline-event routing
        node_id: Node UUID for pipeline-event routing
        correlation_id: Airflow dag_run_id, for event correlation

    Returns:
        Dict with stored count and success status
    """
    import asyncio
    import logging
    import pickle
    import os
    import sys
    from uuid import UUID

    from llama_index.core import Document, VectorStoreIndex, StorageContext, Settings
    from llama_index.vector_stores.milvus import MilvusVectorStore
    from llama_index.vector_stores.milvus.utils import BM25BuiltInFunction
    from src.core.config import settings
    from src.llm import get_embed_model
    from src.resilience.breakers import breaker_call, milvus_breaker
    from src.services.pipeline_events import publish_stage_event

    if '/opt/airflow/dags' not in sys.path:
        sys.path.insert(0, '/opt/airflow/dags')
    from tasks.chunk_schema import CHUNK_SCHEMA_VERSION

    log = logging.getLogger(__name__)

    publish_stage_event(
        "indexing", "started",
        document_id=document_id,
        tenant_id=tenant_id,
        node_id=node_id,
        correlation_id=correlation_id,
    )

    doc_uuid = UUID(document_id)

    if not nl_result:
        log.warning("No NL result provided")
        return {'stored_count': 0, 'success': False}

    # Load data from pickle file
    data_file = nl_result.get('data_file')
    if not data_file or not os.path.exists(data_file):
        log.warning(f"Data file not found: {data_file}")
        return {'stored_count': 0, 'success': False}

    log.info(f"Loading NL records from {data_file}")
    with open(data_file, 'rb') as f:
        nl_data = pickle.load(f)

    chunks = nl_data.get('chunks', [])
    if not chunks:
        log.warning("No records to store")
        return {'stored_count': 0, 'success': True}

    original_filename = nl_data.get('original_filename', '')
    log.info(f"Storing {len(chunks)} Large Table records to LlamaIndex")

    # Convert records to LlamaIndex Documents (before async block)
    documents = []
    for record in chunks:
        nl_text = record.get('text', '')
        if not nl_text or not nl_text.strip():
            continue

        source = record.get('_source', {})

        doc = Document(
            doc_id=str(doc_uuid),  # Set doc_id to our document_id
            text=nl_text,
            metadata={
                "document_id": str(doc_uuid),
                "tenant_id": tenant_id,
                "node_id": node_id,
                "schema_version": CHUNK_SCHEMA_VERSION,
                "filename": source.get('filename', original_filename),
                "document_type": "excel",
                "sheet_name": source.get('sheet', ''),
                "row_number": source.get('row', 0),
                "record_json": record.get('content', {}),
                "domain_category": record.get('domain_category', 'general'),
                "domain_terms": record.get('domain_terms', []),
                "tags": record.get('tags', []),
                # TEXTIQ-569: emitted by SemanticTagger.tag_records_async
                # post-inheritance. Stored under the Milvus dynamic envelope
                # via LlamaIndex's node_to_metadata_dict; indexed via the
                # JSON-path INVERTED index for tag-cascade filtering.
                "verified_tag_names": record.get('verified_tag_names', []),
            },
        )
        documents.append(doc)

    if not documents:
        log.warning("No valid documents to store")
        return {'stored_count': 0, 'success': True}

    log.info(f"Created {len(documents)} LlamaIndex Documents")

    async def store_documents():
        """Async wrapper for LlamaIndex operations (MilvusVectorStore uses async internally)."""
        # Configure embedding model via provider factory
        Settings.embed_model = get_embed_model()

        # Initialize MilvusVectorStore with hybrid search (dense + BM25 sparse)
        milvus_uri = f"http://{settings.milvus_host}:{settings.milvus_port}"
        log.info(f"Connecting to Milvus: {milvus_uri}")

        vector_store = MilvusVectorStore(
            uri=milvus_uri,
            collection_name=settings.milvus_collection,
            dim=settings.embedding_dimensions,
            overwrite=False,
            enable_sparse=True,
            sparse_embedding_function=BM25BuiltInFunction(),
            hybrid_ranker="RRFRanker",
            hybrid_ranker_params={"k": 60},
        )

        storage_context = StorageContext.from_defaults(vector_store=vector_store)

        # Store directly via VectorStoreIndex (no hierarchy).
        # Wrapped with milvus_breaker — from_documents executes Milvus inserts.
        index = breaker_call(
            milvus_breaker,
            VectorStoreIndex.from_documents,
            documents,
            storage_context=storage_context,
            show_progress=True,
        )

        return len(documents)

    stored_count = asyncio.run(store_documents())
    log.info(f"Successfully stored {stored_count} Large Table records")

    # Clean up data file
    try:
        os.remove(data_file)
        log.info(f"Cleaned up data file: {data_file}")
    except Exception as e:
        log.warning(f"Failed to clean up data file: {e}")

    return {
        'stored_count': stored_count,
        'document_id': document_id,
        'success': True,
    }
