from __future__ import annotations

import base64
from io import BytesIO
from typing import Any

import dataiku

from backend.config import get_conversation_vision_llm, get_default_llm_id, get_text_extraction_type
from backend.services.document_extractor_utils import structured_document_chunks_to_text

try:
    from dataikuapi.dss.document_extractor import DocumentExtractor, LocalFileDocumentRef
except ImportError:  # pragma: no cover
    DocumentExtractor = None  # type: ignore
    LocalFileDocumentRef = None  # type: ignore


def build_docx_llm_readable_content(*, part: dict[str, Any]) -> dict[str, Any]:
    """Build text LLM-readable content for one inline DOCX part using DSS structured extraction."""

    if DocumentExtractor is None or LocalFileDocumentRef is None:
        raise RuntimeError("DocumentExtractor not available")

    data_base64 = part.get("dataBase64")
    if not isinstance(data_base64, str) or not data_base64.strip():
        raise ValueError("DOCX part does not contain inline data")

    normalized_data = "".join(data_base64.strip().split())
    docx_bytes = base64.b64decode(normalized_data, validate=True)
    filename = _docx_source_filename(part)

    client = dataiku.api_client()
    project = client.get_default_project()
    doc_extractor = DocumentExtractor(client, project.project_key)

    image_handling_mode = get_text_extraction_type()
    extract_kwargs: dict[str, Any] = {
        "document": _build_local_docx_document_ref(docx_bytes=docx_bytes, filename=filename),
        "image_handling_mode": image_handling_mode,
    }
    if image_handling_mode == "OCR":
        extract_kwargs["ocr_engine"] = "AUTO"
    elif image_handling_mode == "VLM_ANNOTATE":
        extract_kwargs["llm_id"] = get_conversation_vision_llm() or get_default_llm_id()

    structured_response = doc_extractor.structured_extract(**extract_kwargs)
    structured_response._fail_unless_success()
    text = structured_document_chunks_to_text(structured_response.text_chunks or [])
    if not text:
        raise RuntimeError("No readable DOCX text could be extracted")
    return {"text": text}


def _build_local_docx_document_ref(*, docx_bytes: bytes, filename: str):
    """Wrap inline DOCX bytes into the DSS local-file document reference expected by DocumentExtractor."""

    docx_stream = BytesIO(docx_bytes)
    docx_stream.name = filename
    return LocalFileDocumentRef(docx_stream)


def _docx_source_filename(part: dict[str, Any]) -> str:
    filename = part.get("filename")
    if isinstance(filename, str) and filename.strip():
        return filename.strip()
    return "document.docx"
