from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Any, Optional

import dataiku

from backend.config import (
    get_conversation_vision_llm,
    get_extraction_mode,
    get_text_extraction_type,
    get_uploads_managedfolder_id,
)
from backend.constants import ExtractionMode
from backend.services.derived_documents_service import (
    DocumentExtractor,
    ManagedFolderDocumentRef,
    _extract_document_text,
)
from backend.utils.logging_utils import get_logger

logger = get_logger(__name__)


class TextExtractionMode(str, Enum):
    IGNORE = "IGNORE"
    OCR = "OCR"
    VLM_ANNOTATE = "VLM_ANNOTATE"


class TextExtractionResolutionReason(str, Enum):
    ADMIN_PAGES_SCREENSHOTS_REQUIRES_OCR = "admin_pages_screenshots_requires_ocr"
    STANDALONE_IMAGE_STRUCTURED_EXTRACT_REQUIRES_OCR = "standalone_image_structured_extract_requires_ocr"
    ADMIN_TEXT_EXTRACTION_SETTING = "admin_text_extraction_setting"


@dataclass(frozen=True)
class ResolvedTextExtractionStrategy:
    requested_mode: TextExtractionMode
    effective_mode: TextExtractionMode
    reason: TextExtractionResolutionReason
    llm_id: Optional[str] = None


@dataclass(frozen=True)
class ConversationTextFallbackCandidate:
    document_name: str
    document_path: str
    metadata: dict[str, Any]


def generate_missing_conversation_text_fallbacks(*, store, conv_id: str) -> bool:
    """
    Generate missing extracted-text fallbacks for existing conversation documents.

    This service is intentionally scoped to the non-multimodal compatibility flow:
    it only targets documents that already have snapshots but no persisted `text_path`.
    It does not replace the main document derivation pipeline.
    """
    candidates = list_documents_missing_text_fallback(store=store, conv_id=conv_id)
    if not candidates:
        logger.info("No missing text fallback detected for conversation %s", conv_id)
        return False

    logger.debug(
        "Generating conversation text fallback for conv=%s documents=%s",
        conv_id,
        [candidate.document_name for candidate in candidates],
    )

    if DocumentExtractor is None or ManagedFolderDocumentRef is None:
        logger.warning("DocumentExtractor not available; skipping conversation text fallback generation")
        return False

    folder_id = get_uploads_managedfolder_id()
    client = dataiku.api_client()
    project = client.get_default_project()
    project_key = project.project_key
    folder = project.get_managed_folder(folder_id)
    doc_extractor = DocumentExtractor(client, project_key)

    for candidate in candidates:
        generate_text_fallback_for_document(
            store=store,
            conv_id=conv_id,
            folder=folder,
            folder_id=folder_id,
            doc_extractor=doc_extractor,
            candidate=candidate,
        )

    return True


def list_documents_missing_text_fallback(*, store, conv_id: str) -> list[ConversationTextFallbackCandidate]:
    """Return the conversation documents that have snapshots but no extracted-text fallback yet."""
    candidates: list[ConversationTextFallbackCandidate] = []
    for doc in store.get_derived_documents(conv_id) or []:
        metadata = doc.document_metadata or {}
        if metadata.get("text_path"):
            continue
        if not metadata.get("snapshots"):
            continue
        candidates.append(
            ConversationTextFallbackCandidate(
                document_name=doc.document_name,
                document_path=doc.document_path,
                metadata=metadata,
            )
        )
    return candidates


def generate_text_fallback_for_document(
    *,
    store,
    conv_id: str,
    folder,
    folder_id: str,
    doc_extractor,
    candidate: ConversationTextFallbackCandidate,
) -> bool:
    """
    Generate and persist the text fallback for one derived conversation document.

    This logic intentionally duplicates a narrow subset of document-processing behavior.
    We keep it isolated from `process_derived_documents()` for now to avoid introducing
    regressions in the main document derivation pipeline.
    A future refactor may extract shared helpers once both flows are better covered by tests.
    """
    strategy = resolve_text_extraction_strategy_for_document(document_name=candidate.document_name)
    logger.debug(
        "Resolved conversation text fallback strategy for doc=%s requested=%s effective=%s reason=%s",
        candidate.document_name,
        strategy.requested_mode.value,
        strategy.effective_mode.value,
        strategy.reason.value,
    )

    output_base = build_output_base_path(candidate.document_path)
    document_ref = build_document_ref(document_path=candidate.document_path, folder_id=folder_id)
    text_metadata = _extract_document_text(
        doc_extractor,
        document_ref,
        folder,
        output_base,
        image_handling_mode=strategy.effective_mode.value,
        llm_id=strategy.llm_id,
    )
    if not text_metadata:
        logger.warning(
            "Text fallback generation failed for document %s in conv %s; keeping existing snapshots only",
            candidate.document_name,
            conv_id,
        )
        return False

    persist_text_fallback_metadata(
        store=store,
        conv_id=conv_id,
        candidate=candidate,
        text_metadata=text_metadata,
    )
    return True


def resolve_text_extraction_strategy_for_document(*, document_name: str) -> ResolvedTextExtractionStrategy:
    """
    Resolve the effective text extraction strategy for one conversation document.

    The resolution follows the current admin settings as closely as possible while
    applying DSS backend constraints for the `structured_extract` flow.
    """
    if get_extraction_mode() == ExtractionMode.PAGES_SCREENSHOTS.value:
        return ResolvedTextExtractionStrategy(
            requested_mode=TextExtractionMode.OCR,
            effective_mode=TextExtractionMode.OCR,
            reason=TextExtractionResolutionReason.ADMIN_PAGES_SCREENSHOTS_REQUIRES_OCR,
        )

    requested_mode = TextExtractionMode(get_text_extraction_type())
    if is_standalone_image_file(document_name) and requested_mode == TextExtractionMode.VLM_ANNOTATE:
        return ResolvedTextExtractionStrategy(
            requested_mode=requested_mode,
            effective_mode=TextExtractionMode.OCR,
            reason=TextExtractionResolutionReason.STANDALONE_IMAGE_STRUCTURED_EXTRACT_REQUIRES_OCR,
        )

    return ResolvedTextExtractionStrategy(
        requested_mode=requested_mode,
        effective_mode=requested_mode,
        reason=TextExtractionResolutionReason.ADMIN_TEXT_EXTRACTION_SETTING,
        llm_id=get_conversation_vision_llm() if requested_mode == TextExtractionMode.VLM_ANNOTATE else None,
    )


def is_standalone_image_file(document_name: str) -> bool:
    """Return True when the document is a standalone image file handled through structured_extract."""
    lowered = document_name.lower()
    return lowered.endswith((".png", ".jpg", ".jpeg"))


def build_output_base_path(document_path: str) -> str:
    """Return the managed-folder output prefix associated with a document input path."""
    if document_path.startswith("inputs/"):
        return document_path.replace("inputs/", "outputs/", 1)
    return f"outputs/{document_path}"


def build_document_ref(*, document_path: str, folder_id: str):
    """Create the DSS document reference used by the structured extraction endpoint."""
    return ManagedFolderDocumentRef(document_path, folder_id)


def persist_text_fallback_metadata(
    *,
    store,
    conv_id: str,
    candidate: ConversationTextFallbackCandidate,
    text_metadata: dict[str, Any],
) -> None:
    """Persist the generated text fallback while preserving the existing derived-document metadata."""
    metadata_payload = {
        **candidate.metadata,
        "status": candidate.metadata.get("status") or "ready",
        **text_metadata,
    }
    store.upsert_derived_document(
        conv_id,
        candidate.document_name,
        candidate.document_path,
        metadata_payload,
    )
