"""Celery tasks for V2 document extraction pipeline (Docling-based)."""

import asyncio
from uuid import UUID

import redis as sync_redis
import structlog

from src.core.config import settings
from src.db.session import get_sync_session
from src.extraction_v2.pipeline import ExtractionPipelineV2
from src.extraction_v2.large_table_extractor import LargeTableExtractor, LargeTableConfig
from src.extraction.pipeline import ExtractionPipeline  # V1 fallback
from src.workers.celery_app import celery_app

logger = structlog.get_logger(__name__)


@celery_app.task(name="extract_document_v2", bind=True, max_retries=3)
def extract_document_v2(self, document_id: str, file_path: str):
    """
    Extract document using V2 ExtractionPipeline (Docling-based).

    This task processes an uploaded Excel document through the V2 workflow:
    1. Excel → HTML conversion (Docling + LibreOffice + border detection)
    2. HTML table extraction (BeautifulSoup)
    3. LLM-based header detection (GPT-4o)
    4. JSON record building

    Falls back to V1 pipeline (openpyxl-based) if V2 fails.

    Args:
        document_id: UUID string of the document to process
        file_path: Path to the Excel file

    Returns:
        Dictionary with extraction summary

    Raises:
        Exception: If both V2 and V1 extraction fail after retries
    """
    try:
        logger.info(
            "extraction_v2_task_started",
            document_id=document_id,
            file_path=file_path,
            task_id=self.request.id,
        )

        # Create Redis client for progress tracking
        redis_client = None
        try:
            redis_client = sync_redis.from_url(
                settings.redis_url, decode_responses=False
            )
        except Exception as e:
            logger.warning(
                "redis_connection_failed",
                error=str(e),
                note="Proceeding without progress tracking",
            )

        # Update progress: Starting V2 pipeline
        if redis_client:
            _update_progress(redis_client, document_id, "processing_v2", 10)

        # Try V2 pipeline first
        doc_uuid = UUID(document_id)

        logger.info("attempting_v2_extraction", document_id=document_id)

        try:
            v2_pipeline = ExtractionPipelineV2()

            # Detect file type and route to appropriate processing method
            from pathlib import Path
            file_ext = Path(file_path).suffix.lower()
            is_word_doc = file_ext in ['.docx', '.doc', '.docm']
            is_pdf = file_ext == '.pdf'
            is_pptx = file_ext in ['.pptx', '.ppt', '.pptm']

            # Update progress: Converting document
            if redis_client:
                _update_progress(redis_client, document_id, "converting_html", 25)

            # Route based on file type
            if is_word_doc:
                logger.info("routing_to_word_pipeline", document_id=document_id, extension=file_ext)
                result = v2_pipeline.process_word_document_sync(doc_uuid, file_path)
            elif is_pdf:
                logger.info("routing_to_pdf_pipeline", document_id=document_id, extension=file_ext)
                result = v2_pipeline.process_pdf_document_sync(doc_uuid, file_path)
            elif is_pptx:
                logger.info("routing_to_pptx_pipeline", document_id=document_id, extension=file_ext)
                result = v2_pipeline.process_pptx_document_sync(doc_uuid, file_path)
            else:
                # Use large table detection for Excel files
                # This automatically routes large tables (>1000 rows) to chunked processing
                logger.info("routing_to_excel_pipeline_with_large_table_detection", document_id=document_id, extension=file_ext)
                result = v2_pipeline.process_document_with_large_table_detection_sync(doc_uuid, file_path)

            if result.success:
                logger.info(
                    "v2_extraction_succeeded",
                    document_id=document_id,
                    record_count=result.record_count,
                    table_count=result.table_count,
                )

                # Update progress: Saving records
                if redis_client:
                    _update_progress(redis_client, document_id, "saving_records", 75)

                # Save records to database and mark document as completed
                # Determine if this is a text document (Word, PDF, PowerPoint) for semantic chunking metadata
                is_text_doc = is_word_doc or is_pdf or is_pptx
                with get_sync_session() as db:
                    _save_v2_records(db, result, is_word_doc=is_text_doc)
                    _mark_document_completed(db, document_id, result)

                # Update progress: Indexing vectors
                if redis_client:
                    _update_progress(redis_client, document_id, "indexing_vectors", 85)

                # Index records in Milvus (pass document type and filename)
                try:
                    from pathlib import Path
                    filename = Path(file_path).name
                    # Determine document type for metadata
                    if is_word_doc:
                        doc_type = "word"
                    elif is_pdf:
                        doc_type = "pdf"
                    elif is_pptx:
                        doc_type = "powerpoint"
                    else:
                        doc_type = "excel"

                    vector_count = _index_vectors_sync(
                        result.records,
                        doc_uuid,
                        file_path,
                        doc_type,
                        filename
                    )
                    logger.info(
                        "v2_indexing_succeeded",
                        document_id=document_id,
                        vector_count=vector_count,
                        document_type=doc_type,
                    )
                except Exception as e:
                    logger.error(
                        "v2_indexing_failed",
                        document_id=document_id,
                        error=str(e),
                    )
                    # Don't fail the whole task, records are already saved
                    # Mark document as completed with warning
                    with get_sync_session() as db:
                        from src.db.models import Document
                        doc = db.query(Document).filter(Document.id == doc_uuid).first()
                        if doc:
                            doc.error_message = f"Vector indexing failed: {str(e)[:500]}"
                            db.commit()

                # Update progress: Complete
                if redis_client:
                    _update_progress(redis_client, document_id, "completed", 100)

                return {
                    "status": "success",
                    "pipeline": "v2",
                    "document_id": document_id,
                    "record_count": result.record_count,
                    "table_count": result.table_count,
                }
            else:
                logger.warning(
                    "v2_extraction_failed_fallback_to_v1",
                    document_id=document_id,
                    error=result.error,
                )

                # V2 failed, fallback to V1
                return _fallback_to_v1(
                    self, document_id, file_path, redis_client,
                    v2_error=result.error
                )

        except Exception as e:
            logger.error(
                "v2_extraction_exception_fallback_to_v1",
                document_id=document_id,
                error=str(e),
                exc_info=True,
            )

            # V2 exception, fallback to V1
            return _fallback_to_v1(
                self, document_id, file_path, redis_client,
                v2_error=str(e)
            )

    except Exception as e:
        logger.error(
            "extraction_task_failed",
            document_id=document_id,
            error=str(e),
            exc_info=True,
        )

        # Update progress: Failed
        if redis_client:
            _update_progress(redis_client, document_id, "failed", 0)

        # Mark document as failed in database
        with get_sync_session() as db:
            _mark_document_failed(db, document_id, str(e))

        raise


def _fallback_to_v1(
    task,
    document_id: str,
    file_path: str,
    redis_client,
    v2_error: str | None = None
):
    """
    Fallback to V1 extraction pipeline.

    Args:
        task: Celery task instance
        document_id: Document UUID
        file_path: Path to Excel file
        redis_client: Redis client for progress tracking
        v2_error: Error message from V2 attempt

    Returns:
        Dictionary with extraction summary

    Raises:
        ExtractionError: If file is a Word document (V1 only supports Excel)
    """
    from pathlib import Path
    from src.core.exceptions import ExtractionError

    # Check if file is a Word document - V1 pipeline only supports Excel
    file_ext = Path(file_path).suffix.lower()
    is_word_doc = file_ext in ['.docx', '.doc', '.docm']

    if is_word_doc:
        logger.error(
            "v1_fallback_skipped_word_document",
            document_id=document_id,
            file_ext=file_ext,
            v2_error=v2_error,
        )
        # Mark document as failed since V1 can't process Word documents
        with get_sync_session() as db:
            _mark_document_failed(
                db,
                document_id,
                f"Word document processing failed: {v2_error}. V1 fallback not available for Word documents."
            )
        raise ExtractionError(
            message=f"Word document extraction failed: {v2_error}",
            details={"document_id": document_id, "file_ext": file_ext}
        )

    logger.info(
        "starting_v1_fallback",
        document_id=document_id,
        v2_error=v2_error,
    )

    # Update progress: Fallback to V1
    if redis_client:
        _update_progress(redis_client, document_id, "fallback_v1", 30)

    # Use V1 extraction pipeline with db session
    doc_uuid = UUID(document_id)
    with get_sync_session() as db_session:
        v1_pipeline = ExtractionPipeline(db_session=db_session, redis_client=redis_client)

        # V1 pipeline processing (only takes document_id)
        v1_result = v1_pipeline.process_document(doc_uuid)

    logger.info(
        "v1_extraction_completed",
        document_id=document_id,
        record_count=v1_result.records_extracted if hasattr(v1_result, 'records_extracted') else 0,
    )

    # Update progress: Complete
    if redis_client:
        _update_progress(redis_client, document_id, "completed", 100)

    record_count = v1_result.records_extracted if hasattr(v1_result, 'records_extracted') else 0
    return {
        "status": "success",
        "pipeline": "v1",
        "fallback_reason": v2_error,
        "document_id": document_id,
        "record_count": record_count,
    }


def _update_progress(redis_client, document_id: str, status: str, progress: int):
    """Update extraction progress in Redis."""
    try:
        key = f"extraction_progress:{document_id}"
        redis_client.hset(
            key,
            mapping={
                "status": status,
                "progress": progress,
            }
        )
        redis_client.expire(key, 3600)  # 1 hour TTL
    except Exception as e:
        logger.warning(
            "progress_update_failed",
            document_id=document_id,
            error=str(e),
        )


def _save_v2_records(db, result, is_word_doc: bool = False):
    """
    Save V2 extraction records to database.

    Handles both standard V2 format and large table extraction format.

    Args:
        db: Database session
        result: ExtractionResult from V2 pipeline
        is_word_doc: True if this is a Word document (different record format)
    """
    from src.db.models import ExtractedRecord

    logger.info(
        "saving_v2_records",
        document_id=str(result.document_id),
        record_count=result.record_count,
        is_word_doc=is_word_doc,
    )

    if not result.records:
        logger.warning("no_records_to_save", document_id=str(result.document_id))
        return

    # Prepare batch data for bulk insert
    batch_data = []
    for record in result.records:
        # Check if this is a large table extraction record (has 'extraction_method' key)
        if record.get("extraction_method") == "large_table_chunked":
            # Large table format: {document_id, sheet_name, table_range, row_index, data, extraction_method, header_depth}
            content = record.get("data", {})
            headers = list(content.keys())

            batch_data.append({
                "document_id": result.document_id,
                "sheet_name": record.get("sheet_name", "Unknown"),
                "row_number": record.get("row_index", 0),
                "col_range": record.get("table_range"),
                "content": content,
                "resolved_content": None,
                "headers": headers,
            })
        else:
            # Standard V2 format: {_source: {...}, content: {...}}
            source = record.get("_source", {})
            content = record.get("content", {})


    if is_word_doc:
        # Word document chunks have different format: {file_id, text, metadata, element_type}
        for idx, record in enumerate(result.records):
            text = record.get("text", "")
            metadata = record.get("metadata", {})

            # Store text in content as a structured object
            content = {
                "text": text,
                "element_type": record.get("element_type", "Text"),
            }

            batch_data.append({
                "document_id": result.document_id,
                "sheet_name": metadata.get("source_filename", "Document"),
                "row_number": idx,  # Use index as row number for ordering
                "col_range": None,
                "content": content,
                "resolved_content": None,
                "headers": ["text", "element_type"],
            })
    else:
        # Excel records have format: {_source, content}
        for record in result.records:
            source = record.get("_source", {})
            content = record.get("content", {})

            # Extract flattened headers from content keys
            headers = list(content.keys())

            batch_data.append({
                "document_id": result.document_id,
                "sheet_name": source.get("sheet", "Unknown"),
                "row_number": source.get("row", 0),
                "col_range": None,  # V2 doesn't track col_range
                "content": content,
                "resolved_content": None,  # Will be set by symbol resolution pass
                "headers": headers,
            })

    # Perform bulk insert (efficient for large batches)
    try:
        db.bulk_insert_mappings(ExtractedRecord, batch_data)
        db.commit()

        logger.info(
            "v2_records_saved",
            document_id=str(result.document_id),
            records_inserted=len(batch_data),
        )

    except Exception as e:
        db.rollback()
        logger.error(
            "v2_records_save_failed",
            document_id=str(result.document_id),
            error=str(e),
        )
        raise


def _mark_document_completed(db, document_id: str, result):
    """
    Mark document as completed in database with extraction metadata.

    Args:
        db: Database session
        document_id: Document UUID
        result: ExtractionResult from V2 pipeline
    """
    from datetime import UTC, datetime

    from src.db.models import Document

    logger.info(
        "marking_document_completed",
        document_id=document_id,
        record_count=result.record_count,
        table_count=result.table_count,
    )

    try:
        # Find document by ID
        doc_uuid = UUID(document_id)
        document = db.query(Document).filter(Document.id == doc_uuid).first()

        if document:
            document.status = "completed"
            document.error_message = None
            document.processed_at = datetime.now(UTC)
            document.record_count = result.record_count
            document.sheet_count = result.table_count  # V2 tracks tables, not sheets
            db.commit()

            logger.info(
                "document_marked_completed",
                document_id=document_id,
                record_count=result.record_count,
            )
        else:
            logger.warning(
                "document_not_found_for_completion_marking",
                document_id=document_id,
            )

    except Exception as e:
        db.rollback()
        logger.error(
            "failed_to_mark_document_completed",
            document_id=document_id,
            error=str(e),
        )
        raise


def _mark_document_failed(db, document_id: str, error: str):
    """
    Mark document as failed in database.

    Args:
        db: Database session
        document_id: Document UUID
        error: Error message
    """
    from src.db.models import Document

    logger.info(
        "marking_document_failed",
        document_id=document_id,
        error=error,
    )

    try:
        # Find document by ID
        doc_uuid = UUID(document_id)
        document = db.query(Document).filter(Document.id == doc_uuid).first()

        if document:
            document.status = "failed"
            document.error_message = error[:1000]  # Truncate to avoid overflow
            document.processed_at = None  # Clear processed timestamp
            db.commit()

            logger.info(
                "document_marked_failed",
                document_id=document_id,
            )
        else:
            logger.warning(
                "document_not_found_for_failure_marking",
                document_id=document_id,
            )

    except Exception as e:
        db.rollback()
        logger.error(
            "failed_to_mark_document_failed",
            document_id=document_id,
            error=str(e),
        )


async def _index_vectors_async(
    records_data: list[dict],
    document_id: UUID,
    file_path: str,
    document_type: str = "excel",
    filename: str = "",
):
    """
    Index records in Milvus vector store (async).

    Supports both Excel and Word documents.

    Args:
        records_data: List of record dicts (Excel) or chunk dicts (Word)
        document_id: Document UUID
        file_path: Path to the original file
        document_type: "excel" or "word"
        filename: Document filename

    Raises:
        Exception: If indexing fails
    """
    import os
    from src.db.models import ExtractedRecord as ExtractedRecordModel
    from src.db.session import get_async_session
    from src.extraction.models import ExtractedRecord
    from src.knowledge.vector_store import VectorStore

    logger.info(
        "indexing_vectors_started",
        document_id=str(document_id),
        record_count=len(records_data),
        document_type=document_type,
    )

    # Extract filename if not provided
    if not filename:
        filename = os.path.basename(file_path)

    if document_type == "excel":
        # Convert record dicts to ExtractedRecord objects
        # Handle both standard V2 format and large table extraction format
        extracted_records = []
        for rec_data in records_data:
            # Check if this is a large table extraction record
            if rec_data.get("extraction_method") == "large_table_chunked":
                # Large table format
                content = rec_data.get("data", {})
                headers = list(content.keys())

                enriched_source = {
                    "document_id": str(document_id),
                    "filename": filename,
                    "sheet": rec_data.get("sheet_name", "Unknown"),
                    "table_id": 1,
                    "row": rec_data.get("row_index", 0),
                    "col_range": rec_data.get("table_range"),
                }
            else:
                # Standard V2 format
                source = rec_data.get("_source", {})
                content = rec_data.get("content", {})
                headers = list(content.keys())

                # Enrich source with missing fields required by ExtractedRecord
                enriched_source = {
                    "document_id": source.get("document_id", str(document_id)),
                    "filename": filename,
                    "sheet": source.get("sheet", "Unknown"),
                    "table_id": source.get("table_id", 1),
                    "row": source.get("row", 0),
                    "col_range": source.get("col_range"),  # Optional, can be None
                }

            # Skip if content is empty
            if not content:
                logger.debug(
                    "skipping_empty_record",
                    document_id=str(document_id),
                    record_index=idx,
                )
                continue

            # Create ExtractedRecord object
            extracted_record = ExtractedRecord(
                content=content,
                headers=headers,
                _source=enriched_source,
                _merge_metadata=rec_data.get("_merge_metadata"),
            )
            extracted_records.append(extracted_record)
    else:
        # Word chunks - already in the right format (dict)
        extracted_records = records_data

    # Get record IDs from database (they were just inserted)
    # Use sync session to avoid event loop issues in Celery workers
    from src.db.session import get_sync_session
    from sqlalchemy import select

    with get_sync_session() as session:
        # Query record IDs for this document
        result = session.execute(
            select(ExtractedRecordModel.id)
            .where(ExtractedRecordModel.document_id == document_id)
            .order_by(ExtractedRecordModel.row_number)
        )
        record_ids = [row[0] for row in result.fetchall()]

    if len(record_ids) != len(extracted_records):
        logger.warning(
            "record_count_mismatch",
            document_id=str(document_id),
            extracted_count=len(extracted_records),
            db_count=len(record_ids),
        )

    # Index in Milvus
    vector_store = VectorStore()
    await vector_store.create_collection()  # Idempotent
    vector_count = await vector_store.index_records(
        extracted_records,
        document_id,
        record_ids,
        document_type,
        filename,
    )

    logger.info(
        "indexing_vectors_completed",
        document_id=str(document_id),
        vector_count=vector_count,
        document_type=document_type,
    )

    return vector_count


def _index_vectors_sync(
    records_data: list[dict],
    document_id: UUID,
    file_path: str,
    document_type: str = "excel",
    filename: str = ""
) -> int:
    """
    Index records in Milvus vector store (synchronous wrapper).

    Args:
        records_data: List of record dictionaries with '_source' and 'content'
        document_id: Document UUID
        file_path: Path to the original file (for filename extraction)
        document_type: Type of document ("excel" or "word")
        filename: Document filename

    Returns:
        Number of vectors indexed

    Raises:
        Exception: If indexing fails
    """
    try:
        # Use asyncio.run() which properly handles event loop creation and cleanup
        # This avoids "Future attached to a different loop" errors
        return asyncio.run(_index_vectors_async(
            records_data, document_id, file_path, document_type, filename
        ))
    except Exception as e:
        logger.error(
            "vector_indexing_failed",
            document_id=str(document_id),
            error=str(e),
        )
        raise