"""
Document Extraction DAG - Multi-format document processing pipeline.

This DAG processes documents through extraction, chunking, embedding,
and storage stages with file-type-based branching.

Uses @task.external_python to run heavy processing tasks in the uv
virtual environment at /opt/airflow/.venv which contains all required
dependencies (openpyxl, docling, etc.)

**HIERARCHICAL CHUNKING (LlamaIndex)**:
- Text documents use HierarchicalNodeParser with 3 levels [4096/1024/256 tokens]
- Storage uses LlamaIndex MilvusVectorStore + PostgresDocumentStore
- Parent-child node relationships enable AutoMergingRetriever at query time
"""

import os
from datetime import datetime, timedelta, timezone
from typing import Dict, Any
from airflow import DAG
from airflow.models.param import Param
from airflow.decorators import task
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.utils.task_group import TaskGroup

# Import configuration
from dag_config import CONFIG, default_args

# Import regular operators (no src.* imports)
from operators import (
    ValidateEventOperator,
    FetchDocumentOperator,
)

# Import external_python task functions
from tasks.parse_tasks import (
    parse_excel_large_tables,
    parse_word_document,
    parse_pdf_document,
    parse_powerpoint_document,
    parse_excel_detailed,
)
from tasks.processing_tasks import (
    # Excel Large Table pipeline (flat records, not hierarchical)
    detect_headers_task,
    chunk_large_table_records,
    convert_large_tables_to_natural_language,
    extract_domain_terms_task,
    tag_records_task,
    # Hierarchical chunking tasks (LlamaIndex)
    hierarchical_chunk_text_task,
    hierarchical_chunk_excel_detailed_task,
    extract_domain_terms_nodes_task,
    tag_records_nodes_task,
    store_to_llamaindex_task,
    store_large_table_to_llamaindex_task,
    # Excel Detailed pipeline
    upload_images_to_seaweedfs,
    # Document metadata
    save_document_metadata_task,
    update_document_status_task,
)

# Note: default_args and timeouts now imported from config.py

# Failure callback for debugging
def on_failure_callback(context):
    """Log failure details for debugging."""
    import logging
    log = logging.getLogger(__name__)

    task_instance = context.get('task_instance')
    dag = context.get('dag')
    log.error(
        f"Task failed: {getattr(task_instance, 'task_id', 'unknown')}",
        extra={
            'dag_id': getattr(dag, 'dag_id', 'unknown'),
            'run_id': context.get('run_id'),
            'logical_date': str(context.get('logical_date') or context.get('execution_date') or ''),
            'exception': str(context.get('exception', 'Unknown')),
        }
    )

# SLA miss callback
def sla_miss_callback(dag, task_list, blocking_task_list, slas, blocking_tis):
    """Handle SLA misses."""
    import logging
    log = logging.getLogger(__name__)
    log.warning(f"SLA missed for tasks: {[t.task_id for t in task_list]}")

# Update default_args with callbacks and SLA
dag_default_args = {
    **default_args,
    'on_failure_callback': on_failure_callback,
    'sla': timedelta(hours=2),  # 2-hour SLA for all tasks
}


@task.branch
def detect_format(event: dict) -> list:
    """
    Branch based on file extension.

    Returns the task_id(s) of the appropriate parse task(s) based on file type.
    For Excel files, returns BOTH parse_excel_large_tables and parse_excel_detailed to run in parallel.
    """
    if not event:
        return ['unsupported_format']
    file_path = event.get('object_key', event.get('file_path', ''))
    ext = file_path.lower().rsplit('.', 1)[-1] if '.' in file_path else ''

    if ext in ('xlsx', 'xls', 'xlsm'):
        # Excel: Run BOTH NEW and OLD pipelines in parallel
        return ['parse_excel_large_tables', 'parse_excel_detailed']
    elif ext in ('docx', 'doc', 'docm'):
        return ['parse_word']
    elif ext == 'pdf':
        return ['parse_pdf']
    elif ext in ('pptx', 'ppt', 'pptm'):
        return ['parse_powerpoint']
    elif ext == 'md':
        # Markdown needs no parse step — chunk task reads the .md file directly.
        return ['hierarchical_chunk_text']
    else:
        return ['unsupported_format']


@task(trigger_rule='all_done')
def resolve_terminal_outcome() -> Dict[str, Any]:
    """
    Resolve terminal pipeline outcome (completed vs failed) via XCom pulls.

    Airflow 3's task runtime exposes only a slim ``RuntimeDagRun`` (no ORM access),
    so we can't enumerate TI states directly. Instead, pull XCom from each ``store_*``
    task: the pipeline completed successfully only if at least one store both
    returned ``success`` *and* persisted at least one chunk. Skipped and failed
    stores yield ``None``; a store that "succeeded" with zero chunks is treated
    as a failure so we don't mark empty documents as completed.
    """
    from airflow.sdk import get_current_context
    ti = get_current_context()['ti']
    store_results = {
        'store_large_tables_llamaindex': ti.xcom_pull(task_ids='store_large_tables_llamaindex'),
        'store_excel_detailed_llamaindex': ti.xcom_pull(task_ids='store_excel_detailed_llamaindex'),
        'store_text_llamaindex': ti.xcom_pull(task_ids='store_text_llamaindex'),
    }
    succeeded = [
        name for name, result in store_results.items()
        if isinstance(result, dict)
        and result.get('success')
        and (result.get('stored_count') or 0) > 0
    ]
    if succeeded:
        return {'status': 'completed'}
    any_success_flag = any(
        isinstance(r, dict) and r.get('success') for r in store_results.values()
    )
    error = (
        'all store tasks stored 0 chunks'
        if any_success_flag
        else 'no store task succeeded (see task logs for details)'
    )
    return {'status': 'failed', 'error': error}


@task
def get_document_id(event: dict) -> str:
    """Extract document_id from validate event result."""
    return event.get('document_id', '')


@task
def get_tenant_id(event: dict) -> str:
    """Extract tenant_id from validate event result."""
    return event.get('tenant_id', '')


@task
def get_node_id(event: dict) -> str:
    """Extract node_id from validate event result."""
    return event.get('node_id', '')


@task(trigger_rule='all_done')
def cleanup_temp_files(document_id: str) -> Dict[str, Any]:
    """
    Clean up all temporary pickle files for this document.

    Runs regardless of success or failure (trigger_rule='all_done') to
    prevent accumulation of orphaned files on task failures.
    """
    import os
    import glob

    patterns = [
        f'/opt/airflow/logs/parse_*_{document_id}.pkl',
        f'/opt/airflow/logs/chunks_*_{document_id}.pkl',
        f'/opt/airflow/logs/hierarchical_nodes/nodes_*_{document_id}.pkl',
        f'/opt/airflow/logs/parse_detailed_{document_id}.pkl',
        f'/opt/airflow/logs/adapted_{document_id}.pkl',
        f'/opt/airflow/logs/merged_{document_id}.pkl',
        f'/opt/airflow/logs/nl_chunks_{document_id}_*.pkl',
        f'/opt/airflow/logs/embedding_data/embeddings_{document_id}_*.pkl',
        f'/opt/airflow/logs/domain_data/domain_terms_{document_id}_*.pkl',
        f'/opt/airflow/logs/tag_data/tagged_{document_id}_*.pkl',
        f'/opt/airflow/logs/record_ids/record_ids_{document_id}.pkl',
    ]

    cleaned = 0
    for pattern in patterns:
        for f in glob.glob(pattern):
            try:
                os.remove(f)
                cleaned += 1
            except Exception:
                pass  # Ignore errors, file might be already deleted

    return {'cleaned_files': cleaned, 'document_id': document_id}


@task(trigger_rule='none_failed_min_one_success')
def select_parse_result(
    word_result: dict = None,
    pdf_result: dict = None,
    ppt_result: dict = None
) -> dict:
    """Select the non-None parse result from text document branches."""
    return word_result or pdf_result or ppt_result or {}


with DAG(
    dag_id='document_extraction_dag',
    default_args=dag_default_args,
    description='Multi-format document extraction pipeline',
    schedule=None,  # Event-triggered only (Airflow 3.x)
    start_date=datetime(2024, 1, 1, tzinfo=timezone.utc),
    catchup=False,
    max_active_runs=1,
    sla_miss_callback=sla_miss_callback,
    tags=['extraction', 'v2', 'docling', 'hierarchical'],
    doc_md="""
## Document Extraction DAG

Processes documents through the extraction pipeline with **hierarchical chunking**:
1. Validate incoming event
2. Fetch document from storage
3. Branch based on file type
4. Parse document (Excel/Word/PDF/PowerPoint)
5. Hierarchical chunking (LlamaIndex HierarchicalNodeParser)
6. Domain terms extraction + semantic tagging
7. Store to LlamaIndex (MilvusVectorStore + PostgresDocumentStore)

### Hierarchical Chunking

Text documents (Word/PDF/PPTX) use **LlamaIndex HierarchicalNodeParser**:
- 2-level hierarchy: MarkdownNodeParser -> SentenceSplitter(512)
- Parent-child relationships preserved in PostgresDocumentStore
- Leaf nodes indexed in Milvus for vector search
- Enables AutoMergingRetriever for context expansion at query time

### Excel Hybrid Pipeline (Large Table + Detailed)

For Excel documents, this DAG runs TWO pipelines in PARALLEL:

**Large Table Pipeline** (Docling-based):
- parse_excel_large_tables → detect_headers → chunk_records → convert_to_nl → domain_terms → tags → store_large_table_to_llamaindex
- 1 row = 1 Document (flat, no hierarchy)

**Detailed Pipeline** (openpyxl-based with image/shape extraction):
- parse_excel_detailed → upload_images → hierarchical_chunk → domain_terms_nodes → tag_nodes → store_to_llamaindex
- Hierarchical chunking for text content

### Storage Architecture

```
┌─────────────────────────────────────────────────────────────┐
│              LlamaIndex Storage Layer                       │
├─────────────────────────────────────────────────────────────┤
│  MilvusVectorStore          │  PostgresDocumentStore        │
│  ────────────────           │  ─────────────────────        │
│  Leaf nodes only            │  ALL nodes (parent + leaf)    │
│  EMBEDDING_DIMENSIONS-dim   │  Full node hierarchy          │
│  configured embedding model │  Parent-child relationships   │
└─────────────────────────────────────────────────────────────┘
```

### Triggering

Trigger via API with configuration:
```json
{
    "document_id": "uuid",
    "object_key": "path/to/file.xlsx"
}
```
    """,
    params={
        "document_id": Param(type="string", description="Unique document identifier (UUID)"),
        "object_key": Param(type="string", description="Path to the document in storage"),
        "tenant_id": Param(type="string", description="Tenant UUID for data isolation"),
        "node_id": Param(type="string", description="Node UUID for data isolation"),
    },
) as dag:

    # Task 1: Validate incoming event (regular operator)
    validate = ValidateEventOperator(
        task_id='validate_event',
    )

    # Task 2: Fetch document from storage (regular operator)
    fetch = FetchDocumentOperator(
        task_id='fetch_document',
    )

    # Get document_id and tenant context for downstream tasks
    doc_id = get_document_id(validate.output)
    t_id = get_tenant_id(validate.output)
    n_id = get_node_id(validate.output)

    # Task 2.5: Save document metadata to PostgreSQL (external_python task)
    save_doc_metadata = save_document_metadata_task.override(
        task_id='save_document_metadata',
    )(
        document_id=doc_id,
        fetch_result=fetch.output,
        validation_result=validate.output,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # Task 3: Detect format and branch (TaskFlow branch)
    branch = detect_format(validate.output)

    # Unsupported format handler (dead end)
    unsupported = EmptyOperator(
        task_id='unsupported_format',
    )

    # =========================================================================
    # EXCEL LARGE TABLE PIPELINE (flat records, no hierarchy)
    # =========================================================================

    # Task 4a: Parse Excel for LARGE TABLES (external_python task)
    parse_excel_large = parse_excel_large_tables.override(
        task_id='parse_excel_large_tables',
        execution_timeout=CONFIG.PARSE_TIMEOUT,
    )(
        fetch_result=fetch.output,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # Task 5a: Detect headers (Excel only)
    detect_headers = detect_headers_task.override(
        task_id='detect_headers',
    )(parse_result=parse_excel_large)

    # Task 6a: Chunk large table records
    chunk_large_tables = chunk_large_table_records.override(
        task_id='chunk_large_table_records',
    )(
        document_id=doc_id,
        parse_result=parse_excel_large,
        headers_result=detect_headers,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # Task 6.4a: Convert large table records to natural language
    nl_chunks = convert_large_tables_to_natural_language.override(
        task_id='convert_large_tables_to_nl',
        execution_timeout=timedelta(minutes=30),
    )(chunks_result=chunk_large_tables)

    # Task 6.5a: Extract domain terms from Excel chunks (flat records)
    domain_terms_excel = extract_domain_terms_task.override(
        task_id='extract_domain_terms_excel',
        execution_timeout=CONFIG.LLM_TASK_TIMEOUT,
    )(chunks_result=nl_chunks, tenant_id=t_id, node_id=n_id)

    # Task 6.6a: Apply semantic tags to Excel chunks
    tags_excel = tag_records_task.override(
        task_id='tag_records_excel',
        execution_timeout=CONFIG.LLM_TASK_TIMEOUT,
    )(chunks_result=domain_terms_excel)

    # Task 8a: Store Large Table records to LlamaIndex MilvusVectorStore
    # (1 record = 1 Document, no hierarchy, LlamaIndex handles embeddings)
    store_large_tables = store_large_table_to_llamaindex_task.override(
        task_id='store_large_tables_llamaindex',
        execution_timeout=CONFIG.STORE_TIMEOUT,
    )(
        document_id=doc_id,
        nl_result=tags_excel,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # =========================================================================
    # EXCEL DETAILED PIPELINE (hierarchical chunking)
    # =========================================================================

    # Task 4b: Parse Excel with DETAILED extraction (images, shapes, flowcharts)
    parse_excel_det = parse_excel_detailed.override(
        task_id='parse_excel_detailed',
        execution_timeout=CONFIG.PARSE_TIMEOUT,
    )(
        fetch_result=fetch.output,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # Task 5b: Upload images to SeaweedFS
    upload_images = upload_images_to_seaweedfs.override(
        task_id='upload_images_to_seaweedfs',
    )(
        document_id=doc_id,
        detailed_pipeline_result=parse_excel_det,
        tenant_id=t_id,
        node_id=n_id,
    )

    # Task 6b: Hierarchical chunking for Excel Detailed (LlamaIndex)
    hierarchical_chunk_excel = hierarchical_chunk_excel_detailed_task.override(
        task_id='hierarchical_chunk_excel_detailed',
        execution_timeout=CONFIG.EMBEDDING_TIMEOUT,
    )(
        document_id=doc_id,
        parse_result=upload_images,
    )

    # Task 6.5b: Extract domain terms from Excel Detailed nodes
    domain_terms_excel_det = extract_domain_terms_nodes_task.override(
        task_id='extract_domain_terms_excel_detailed',
        execution_timeout=CONFIG.LLM_TASK_TIMEOUT,
    )(nodes_result=hierarchical_chunk_excel, tenant_id=t_id, node_id=n_id)

    # Task 6.6b: Apply semantic tags to Excel Detailed nodes
    tags_excel_det = tag_records_nodes_task.override(
        task_id='tag_records_excel_detailed',
        execution_timeout=CONFIG.LLM_TASK_TIMEOUT,
    )(nodes_result=domain_terms_excel_det)

    # Task 8b: Store Excel Detailed nodes to LlamaIndex
    store_excel_detailed = store_to_llamaindex_task.override(
        task_id='store_excel_detailed_llamaindex',
        execution_timeout=CONFIG.STORE_TIMEOUT,
    )(
        document_id=doc_id,
        nodes_result=tags_excel_det,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # =========================================================================
    # TEXT PIPELINE (Word/PDF/PPTX - hierarchical chunking)
    # =========================================================================

    # Task 4c: Parse Word
    parse_word = parse_word_document.override(
        task_id='parse_word',
        execution_timeout=CONFIG.PARSE_TIMEOUT,
    )(
        fetch_result=fetch.output,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # Task 4d: Parse PDF
    parse_pdf = parse_pdf_document.override(
        task_id='parse_pdf',
        execution_timeout=CONFIG.PARSE_TIMEOUT,
    )(
        fetch_result=fetch.output,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # Task 4e: Parse PowerPoint
    parse_powerpoint = parse_powerpoint_document.override(
        task_id='parse_powerpoint',
        execution_timeout=CONFIG.PARSE_TIMEOUT,
    )(
        fetch_result=fetch.output,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # Select text parse result from branches
    text_parse_result = select_parse_result(
        word_result=parse_word,
        pdf_result=parse_pdf,
        ppt_result=parse_powerpoint,
    )

    # Task 6c: Hierarchical chunking for text documents (LlamaIndex).
    # Two upstream paths feed this task:
    #   - parse_result via select_parse_result (Word/PDF/PowerPoint)
    #   - fetch_result_md directly from fetch (Markdown — no parse step)
    # trigger_rule='none_failed_min_one_success' so it fires when either path produces a result.
    # The Excel branches still cascade a skip via select_parse_result (skipped) + branch
    # (which does not list this task for Excel).
    hierarchical_chunk_text = hierarchical_chunk_text_task.override(
        task_id='hierarchical_chunk_text',
        trigger_rule='none_failed_min_one_success',
    )(
        document_id=doc_id,
        parse_result=text_parse_result,
        fetch_result_md=fetch.output,
    )

    # Task 6.5c: Extract domain terms from text nodes
    domain_terms_text = extract_domain_terms_nodes_task.override(
        task_id='extract_domain_terms_text',
        execution_timeout=CONFIG.LLM_TASK_TIMEOUT,
    )(nodes_result=hierarchical_chunk_text, tenant_id=t_id, node_id=n_id)

    # Task 6.6c: Apply semantic tags to text nodes
    tags_text = tag_records_nodes_task.override(
        task_id='tag_records_text',
        execution_timeout=CONFIG.LLM_TASK_TIMEOUT,
    )(nodes_result=domain_terms_text)

    # Task 8c: Store text nodes to LlamaIndex (MilvusVectorStore + PostgresDocumentStore)
    store_text = store_to_llamaindex_task.override(
        task_id='store_text_llamaindex',
        execution_timeout=CONFIG.STORE_TIMEOUT,
    )(
        document_id=doc_id,
        nodes_result=tags_text,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # =========================================================================
    # FINALIZATION
    # =========================================================================

    # Collect record count from whichever store task ran
    @task(task_id='get_record_count', trigger_rule='all_done')
    def get_record_count(large=None, excel=None, text=None) -> int:
        for result in [large, excel, text]:
            if result and isinstance(result, dict) and result.get('success'):
                return result.get('stored_count', 0)
        return 0

    record_count = get_record_count(
        large=store_large_tables,
        excel=store_excel_detailed,
        text=store_text,
    )

    # Task 9a: Resolve terminal outcome from upstream TI states (runs in airflow-core)
    terminal_outcome = resolve_terminal_outcome.override(
        task_id='resolve_terminal_outcome',
    )()

    # Task 9b: Write terminal status to DB + emit terminal pipeline event
    update_doc_status = update_document_status_task.override(
        task_id='update_document_status',
        trigger_rule='all_done',
    )(
        document_id=doc_id,
        outcome=terminal_outcome,
        record_count=record_count,
        tenant_id=t_id,
        node_id=n_id,
        correlation_id='{{ run_id }}',
    )

    # Define dependencies
    # Initial flow: validate → fetch → save_doc_metadata → branch
    validate >> fetch >> save_doc_metadata >> branch

    # Unsupported format (dead end)
    branch >> unsupported

    # Excel Large Table path:
    # branch → parse_excel_large → detect_headers → chunk → nl → domain_terms → tags → store
    branch >> parse_excel_large >> detect_headers >> chunk_large_tables >> nl_chunks >> domain_terms_excel >> tags_excel >> store_large_tables

    # Excel Detailed path:
    # branch → parse_excel_detailed → upload_images → hierarchical_chunk → domain_terms → tags → store
    branch >> parse_excel_det >> upload_images >> hierarchical_chunk_excel >> domain_terms_excel_det >> tags_excel_det >> store_excel_detailed

    # Word path: branch → parse_word → text_parse_result
    branch >> parse_word >> text_parse_result

    # PDF path: branch → parse_pdf → text_parse_result
    branch >> parse_pdf >> text_parse_result

    # PowerPoint path: branch → parse_powerpoint → text_parse_result
    branch >> parse_powerpoint >> text_parse_result

    # Markdown path: branch → hierarchical_chunk_text (no parse, no select_parse_result).
    # hierarchical_chunk_text reads the .md file directly via fetch_result_md.
    branch >> hierarchical_chunk_text

    # Text pipeline: parse_result → hierarchical_chunk → domain_terms → tags → store
    text_parse_result >> hierarchical_chunk_text >> domain_terms_text >> tags_text >> store_text

    # Terminal outcome resolution must run after all data-pipeline tasks have reached
    # a final state so it can inspect TI states correctly (trigger_rule='all_done').
    [store_large_tables, store_excel_detailed, store_text, record_count] >> terminal_outcome >> update_doc_status

    # Cleanup: Run after all paths complete (success or failure)
    cleanup = cleanup_temp_files(doc_id)
    [update_doc_status, unsupported] >> cleanup