from __future__ import annotations

import base64
from io import BytesIO
from typing import Any

import dataiku
from PIL import Image

from backend.config import get_conversation_vision_llm, get_default_llm_id, get_text_extraction_type
from backend.constants import ExtractionMode
from backend.services.document_extractor_utils import (
    decode_document_extractor_image_payload,
    resize_image_with_ratio,
    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_pdf_llm_readable_content(*, part: dict[str, Any], extraction_mode: str) -> dict[str, Any]:
    """Build mode-aligned LLM-readable content for one inline PDF part using DSS document services."""

    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("PDF part does not contain inline data")

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

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

    if extraction_mode == ExtractionMode.PAGES_TEXT.value:
        text = _extract_pdf_text(doc_extractor=doc_extractor, pdf_bytes=pdf_bytes, filename=filename)
        if not text:
            raise RuntimeError("No readable PDF text could be extracted")
        return {"text": text}

    if extraction_mode == ExtractionMode.PAGES_SCREENSHOTS.value:
        images = _extract_pdf_images(doc_extractor=doc_extractor, pdf_bytes=pdf_bytes, filename=filename)
        if not images:
            raise RuntimeError("No readable PDF page screenshots could be extracted")
        return {"images": images}

    raise RuntimeError(f"Unsupported extraction mode: {extraction_mode}")


def _extract_pdf_text(*, doc_extractor, pdf_bytes: bytes, filename: str) -> str:
    image_handling_mode = get_text_extraction_type()
    extract_kwargs: dict[str, Any] = {
        "document": _build_local_pdf_document_ref(pdf_bytes=pdf_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()
    return structured_document_chunks_to_text(structured_response.text_chunks or [])


def _extract_pdf_images(*, doc_extractor, pdf_bytes: bytes, filename: str) -> list[dict[str, Any]]:
    screenshots_response = doc_extractor.generate_pages_screenshots(
        document=_build_local_pdf_document_ref(pdf_bytes=pdf_bytes, filename=filename),
        output_managed_folder=None,
        fetch_size=50,
    )
    images: list[dict[str, Any]] = []
    for page_index in range(screenshots_response.total_count):
        screenshot = screenshots_response.fetch_screenshot(page_index)
        image_bytes = decode_document_extractor_image_payload(screenshot.image)
        image = Image.open(BytesIO(image_bytes))
        image = resize_image_with_ratio(image)
        out_buf = BytesIO()
        image.save(out_buf, format="PNG")
        images.append(
            {
                "pageNumber": page_index + 1,
                "mimeType": "image/png",
                "dataBase64": base64.b64encode(out_buf.getvalue()).decode("utf-8"),
            }
        )
    return images


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

    pdf_stream = BytesIO(pdf_bytes)
    pdf_stream.name = filename
    return LocalFileDocumentRef(pdf_stream)


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