import base64
import io
import json
import os
from typing import Any, Dict, List, Optional, Tuple

import dataiku
from PIL import Image

from backend.config import (
    get_conversation_vision_llm,
    get_default_llm_id,
    get_guardrails_enabled,
    get_image_validation,
    get_text_extraction_type,
    get_uploads_managedfolder_id,
)
from backend.schemas import schemas
from backend.services.document_extractor_utils import (
    decode_document_extractor_image_payload,
    resize_image_with_ratio,
)
from backend.utils.logging_utils import get_logger

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

logger = get_logger(__name__)


def _ensure_folder_prefix(folder, base_path: str) -> None:
    keep_path = f"{base_path}/.init"
    try:
        folder.put_file(keep_path, b"")
    except Exception:
        pass


def _extract_document_text(
    doc_extractor,
    document_ref,
    folder,
    output_base: str,
    image_handling_mode: str = "IGNORE",
    llm_id: Optional[str] = None,
    output_managed_folder: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Run structured extraction to obtain text-only content for a document.

    Args:
        doc_extractor: DocumentExtractor instance
        document_ref: ManagedFolderDocumentRef instance
        folder: Dataiku folder instance
        output_base: Base path for output files
        image_handling_mode: How to handle images ("IGNORE", "OCR", or "VLM_ANNOTATE")
        llm_id: LLM ID for VLM_ANNOTATE mode (optional)
        output_managed_folder: Managed folder ID for storing images in VLM_ANNOTATE mode (optional)
    """
    try:
        extract_kwargs = {
            "document": document_ref,
            "image_handling_mode": image_handling_mode,
            "image_validation": get_image_validation(),
        }

        if image_handling_mode == "OCR":
            extract_kwargs["ocr_engine"] = "AUTO"
        elif image_handling_mode == "VLM_ANNOTATE":
            if llm_id:
                extract_kwargs["llm_id"] = llm_id
            if output_managed_folder:
                extract_kwargs["output_managed_folder"] = output_managed_folder

        logger.debug(f"Text extraction parameters: {extract_kwargs}")
        structured_response = doc_extractor.structured_extract(**extract_kwargs)
        structured_response._fail_unless_success()

        chunks = structured_response.text_chunks or []
    except Exception as err:  # pragma: no cover - rely on logging
        logger.exception(
            "Failed to extract text for document %s",
            document_ref.path if hasattr(document_ref, "path") else document_ref,
        )
        return {}

    # Save chunks as JSON
    text_path = f"{output_base}/extracted.json"
    try:
        chunks_json = json.dumps(chunks, ensure_ascii=False, indent=2)
        folder.put_file(text_path, io.BytesIO(chunks_json.encode("utf-8")))
    except Exception as err:  # pragma: no cover - rely on logging
        logger.exception(
            "Failed to persist extracted chunks for %s: %s",
            document_ref.path if hasattr(document_ref, "path") else document_ref,
            err,
        )
        text_path = None
    metadata: Dict[str, Any] = {}
    if text_path:
        metadata["text_path"] = text_path
    return metadata


def _extract_text_metadata_for_document(
    *,
    doc_extractor,
    document_ref,
    folder,
    output_base: str,
    text_extraction_type_override: Optional[str] = None,
) -> Optional[Dict[str, Any]]:
    """Extract text metadata for one document using configured strategy or explicit override."""
    image_handling_mode = text_extraction_type_override or get_text_extraction_type()
    llm_id = get_conversation_vision_llm() if image_handling_mode == "VLM_ANNOTATE" else get_default_llm_id()
    return _extract_document_text(
        doc_extractor,
        document_ref,
        folder,
        output_base,
        image_handling_mode=image_handling_mode,
        llm_id=llm_id,
    )


def _handle_text_mode_image_file(
    *,
    store,
    conv_id: str,
    doc_name: str,
    document_path: str,
    output_base: str,
    file_ext: str,
    document_ref,
    existing_snapshots: list[dict],
    build_image_snapshots,
    extract_text_metadata,
) -> None:
    """
    Handle image attachments in `pagesText` mode:
    keep/build snapshots and add text fallback when extraction succeeds.
    """
    snapshots = existing_snapshots
    if not snapshots:
        try:
            snapshots = build_image_snapshots(document_path, output_base, file_ext)
        except Exception as err:
            logger.exception(
                "Failed to process image attachment %s for conv %s",
                doc_name,
                conv_id,
            )
            store.upsert_derived_document(
                conv_id,
                doc_name,
                document_path,
                {"status": "failed", "error": str(err), "snapshots": []},
            )
            return

    text_metadata = extract_text_metadata(document_ref, output_base)
    if not text_metadata:
        logger.warning(
            "Text extraction failed for image attachment %s in conv %s; keeping snapshots only",
            doc_name,
            conv_id,
        )
        store.upsert_derived_document(
            conv_id,
            doc_name,
            document_path,
            {"snapshots": snapshots, "status": "ready"},
        )
        return

    metadata_payload = {"snapshots": snapshots, "status": "ready", **text_metadata}
    store.upsert_derived_document(conv_id, doc_name, document_path, metadata_payload)


def _build_image_snapshots(
    *,
    folder,
    document_path: str,
    output_base: str,
    file_ext: str,
    ensure_rgb_for_jpeg,
) -> list[dict]:
    with folder.get_file(document_path) as response:
        image_bytes = response.raw.read()

    target_path = f"{output_base}/page_1.{file_ext}"
    img = Image.open(io.BytesIO(image_bytes))
    img = resize_image_with_ratio(img)
    out_buf = io.BytesIO()
    save_format = "JPEG" if file_ext in {"jpg", "jpeg"} else "PNG"
    if save_format == "JPEG":
        img = ensure_rgb_for_jpeg(img)
    img.save(out_buf, format=save_format)
    out_buf.seek(0)
    folder.put_file(target_path, out_buf)
    return [{"page_number": 1, "screenshot_path": target_path}]


def process_derived_documents(
    store,
    conv_id: str,
    user_id: str,
    attachments: List[Dict[str, Any]],
    *,
    extraction_mode: str = "pagesScreenshots",
    text_extraction_type_override: Optional[str] = None,
) -> None:
    """
    Populate derived document artifacts for each attachment.

    Args:
        extraction_mode: "pagesScreenshots" (default) to generate screenshots, or "pagesText" for text-only extraction
        text_extraction_type_override: Optional override for text extraction strategy used only when
            extraction_mode == "pagesText". When unset, the admin-configured text extraction type is used.
    """
    if not attachments:
        return
    if DocumentExtractor is None or ManagedFolderDocumentRef is None:
        logger.warning("DocumentExtractor not available; skipping derived document processing")
        return

    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)
    docs: List[schemas.DerivedDocument] | None = store.get_derived_documents(conv_id) or []
    existing_docs = {doc.document_path: doc for doc in docs}

    def _ensure_rgb_for_jpeg(img: Image.Image) -> Image.Image:
        if img.mode == "RGB":
            return img
        if img.mode in {"RGBA", "LA"} or (img.mode == "P" and "transparency" in img.info):
            rgba_image = img.convert("RGBA")
            background = Image.new("RGBA", rgba_image.size, (255, 255, 255, 255))
            composited = Image.alpha_composite(background, rgba_image)
            return composited.convert("RGB")
        return img.convert("RGB")

    def _handle_simple_text_file(doc_name: str, document_path: str, output_base: str) -> None:
        try:
            with folder.get_file(document_path) as response:
                content_bytes = response.raw.read()
            text_path = f"{output_base}/content.txt"
            with io.BytesIO(content_bytes) as bstream:
                folder.put_file(text_path, bstream)
            metadata_payload = {
                "status": "ready",
                "text_path": text_path,
                "snapshots": [],
            }
            store.upsert_derived_document(conv_id, doc_name, document_path, metadata_payload)
        except Exception as err:
            logger.exception(
                "Failed to process text attachment %s for conv %s",
                doc_name,
                conv_id,
            )
            store.upsert_derived_document(
                conv_id,
                doc_name,
                document_path,
                {"status": "failed", "error": str(err), "snapshots": []},
            )

    def _handle_image_file(doc_name: str, document_path: str, output_base: str, file_ext: str) -> None:
        try:
            snapshots = _build_image_snapshots(
                folder=folder,
                document_path=document_path,
                output_base=output_base,
                file_ext=file_ext,
                ensure_rgb_for_jpeg=_ensure_rgb_for_jpeg,
            )
            metadata_payload = {
                "status": "ready",
                "snapshots": snapshots,
            }
            store.upsert_derived_document(conv_id, doc_name, document_path, metadata_payload)
        except Exception as err:
            logger.exception(
                "Failed to process image attachment %s for conv %s",
                doc_name,
                conv_id,
            )
            store.upsert_derived_document(
                conv_id,
                doc_name,
                document_path,
                {"status": "failed", "error": str(err), "snapshots": []},
            )

    # Build a mapping from document_path to attachment info
    attachment_map = {att.get("document_path"): att for att in attachments if att.get("document_path")}

    # Get all unique source paths from both existing docs and new attachments
    all_source_paths = set(existing_docs.keys()) | set(attachment_map.keys())

    for document_path in all_source_paths:
        attachment = attachment_map.get(document_path)
        doc_record = existing_docs.get(document_path)
        doc_name = None
        if attachment:
            doc_name = attachment.get("document_name")
        elif doc_record:
            doc_name = doc_record.document_name

        if not doc_name or not document_path:
            continue

        file_ext = os.path.splitext(doc_name.lower())[1].lstrip(".")
        doc_meta = doc_record.document_metadata if doc_record else {}
        status = doc_meta.get("status")
        existing_snapshots = doc_meta.get("snapshots", []) if doc_meta else []
        has_snapshots = bool(existing_snapshots)
        has_text = bool(doc_meta.get("text_path")) if doc_meta else False

        if document_path.startswith("inputs/"):
            output_base = document_path.replace("inputs/", "outputs/", 1)
        else:
            output_base = f"outputs/{document_path}"

        _ensure_folder_prefix(folder, output_base)
        document_ref = ManagedFolderDocumentRef(document_path, folder_id)

        is_text_file = file_ext in {"txt", "md", "html"}
        is_image_file = file_ext in {"png", "jpg", "jpeg"}

        if extraction_mode == "pagesText":
            if is_text_file:
                _handle_simple_text_file(doc_name, document_path, output_base)
                continue
            if is_image_file:
                _handle_text_mode_image_file(
                    store=store,
                    conv_id=conv_id,
                    doc_name=doc_name,
                    document_path=document_path,
                    output_base=output_base,
                    file_ext=file_ext,
                    document_ref=document_ref,
                    existing_snapshots=existing_snapshots,
                    build_image_snapshots=lambda doc_path, out_base, ext: _build_image_snapshots(
                        folder=folder,
                        document_path=doc_path,
                        output_base=out_base,
                        file_ext=ext,
                        ensure_rgb_for_jpeg=_ensure_rgb_for_jpeg,
                    ),
                    extract_text_metadata=lambda doc_ref, out_base: _extract_text_metadata_for_document(
                        doc_extractor=doc_extractor,
                        document_ref=doc_ref,
                        folder=folder,
                        output_base=out_base,
                        text_extraction_type_override=text_extraction_type_override,
                    ),
                )
                continue
            if has_text:
                continue

            text_metadata = _extract_text_metadata_for_document(
                doc_extractor=doc_extractor,
                document_ref=document_ref,
                folder=folder,
                output_base=output_base,
                text_extraction_type_override=text_extraction_type_override,
            )
            if not text_metadata:
                store.upsert_derived_document(
                    conv_id,
                    doc_name,
                    document_path,
                    {"snapshots": existing_snapshots, "status": "failed", "error": "text_extraction_failed"},
                )
                continue

            metadata_payload = {"snapshots": existing_snapshots, "status": "ready", **text_metadata}
            store.upsert_derived_document(conv_id, doc_name, document_path, metadata_payload)
            continue

        if is_text_file:
            _handle_simple_text_file(doc_name, document_path, output_base)
            continue

        if is_image_file:
            _handle_image_file(doc_name, document_path, output_base, file_ext)
            continue

        if status == "ready" and has_snapshots:
            continue

        metadata_entries = []
        try:
            screenshots_response = doc_extractor.generate_pages_screenshots(
                document=document_ref,
                output_managed_folder=None,
                fetch_size=50,
            )
            total_pages = screenshots_response.total_count

            for page_index in range(total_pages):
                screenshot = screenshots_response.fetch_screenshot(page_index)
                image_bytes = decode_document_extractor_image_payload(screenshot.image)
                # Resize screenshot and persist
                img = Image.open(io.BytesIO(image_bytes))
                img = resize_image_with_ratio(img)
                out_buf = io.BytesIO()
                img.save(out_buf, format="PNG")
                out_buf.seek(0)

                target_path = f"{output_base}/page_{page_index + 1}.png"
                folder.put_file(target_path, out_buf)
                metadata_entries.append({"page_number": page_index + 1, "screenshot_path": target_path})

            logger.info(
                "Derived document ready for conv=%s doc=%s (%d pages)",
                conv_id,
                doc_name,
                len(metadata_entries),
            )
        except Exception as err:
            logger.exception("Failed to process derived document %s for conv %s", doc_name, conv_id)
            store.upsert_derived_document(
                conv_id,
                doc_name,
                document_path,
                {"snapshots": [], "status": "failed", "error": str(err)},
            )
            continue

        metadata_payload = {"snapshots": metadata_entries, "status": "ready"}
        store.upsert_derived_document(conv_id, doc_name, document_path, metadata_payload)


def get_structured_documents(store, conv_id: str, user_id: str) -> Dict[str, Dict[str, Any]]:
    """Return derived document metadata keyed by source path.

    Documents with guardrails_status 'content_violation' or 'extraction_failed' are filtered out
    and not passed to the LLM, but only when guardrails are enabled. If guardrails are disabled,
    previously-failed documents are still included.
    """
    derived_docs: List[schemas.DerivedDocument] | None = store.get_derived_documents(conv_id)
    documents_out: Dict[str, Dict[str, Any]] = {}
    if not derived_docs:
        return documents_out

    global_snapshot_id = 1
    guardrails_enabled = get_guardrails_enabled()

    for doc in derived_docs:
        metadata = doc.document_metadata or {}
        status = metadata.get("status")
        guardrails_status = metadata.get("guardrails_status")

        # Skip documents that failed guardrails check (either content violation or extraction failure),
        # but only if guardrails are enabled
        if guardrails_enabled and guardrails_status in ("content_violation", "extraction_failed"):
            logger.info(
                f"Filtering out document '{doc.document_name}' from LLM context: guardrails_status={guardrails_status}"
            )
            continue

        snapshots_meta = metadata.get("snapshots") or []
        doc_name = doc.document_name
        prepared_snapshots: list[dict] = []

        # Read text content from the folder where it's stored
        text_content = ""
        text_path = metadata.get("text_path")
        if text_path:
            try:
                folder_id = get_uploads_managedfolder_id()
                import dataiku

                folder = dataiku.Folder(folder_id)
                with folder.get_download_stream(text_path) as stream:
                    text_bytes = stream.read()
                    try:
                        text_content = text_bytes.decode("utf-8", errors="replace")
                    except Exception:
                        text_content = ""
            except Exception:
                # best-effort fallback; keep empty on failure
                text_content = ""

        # If extraction failed and there is no text content, synthesize a short message
        error_msg = metadata.get("error")
        if (not text_content) and status == "failed":
            human_doc_name = doc_name or doc.document_path or "Document"
            reason = f": {error_msg}" if error_msg else ""
            text_content = (
                f'Extraction failed for "{human_doc_name}"{reason}. No readable content is available for this file.'
            )

        for snap in snapshots_meta:
            path = snap.get("screenshot_path")
            if not path:
                continue

            prepared_snapshots.append(
                {
                    "id": global_snapshot_id,
                    "file_name": path.rsplit("/", 1)[-1],
                    "page": snap.get("page") or snap.get("page_number"),
                    "screenshot_path": path,
                }
            )
            global_snapshot_id += 1

        documents_out[doc.document_path] = {
            "name": doc_name,
            "snapshots": prepared_snapshots,
            "status": status,
            "error": error_msg,
            "text": text_content,
            "has_text_fallback": bool(text_path and text_content.strip()),
            "document_path": doc.document_path,
        }

    return documents_out


def count_conversation_images(store, conv_id: str) -> int:
    """
    Count the total number of images (snapshots) generated for all documents in a conversation.

    Args:
        store: The database store instance
        conv_id: The conversation ID

    Returns:
        Total count of image snapshots across all derived documents
    """
    derived_docs: List[schemas.DerivedDocument] | None = store.get_derived_documents(conv_id)
    total_images = 0

    for doc in derived_docs:
        metadata = doc.document_metadata or {}
        snapshots = metadata.get("snapshots") or []
        total_images += len(snapshots)

    logger.debug(f"Conversation {conv_id} has {total_images} total images across {len(derived_docs)} documents")
    return total_images


def predict_new_image_count(attachments: List[Dict[str, Any]], folder_id: str) -> int:
    """
    Predict the number of images that would be generated from new attachments.

    This function extracts page counts from documents without generating screenshots,
    allowing quota decisions to be made before processing.

    Args:
        attachments: List of attachment metadata dicts with 'path' keys
        folder_id: The managed folder ID where documents are stored

    Returns:
        Predicted total number of images (pages) across all attachments
    """
    if not attachments:
        return 0

    try:
        from dataikuapi.dss.document_extractor import DocumentExtractor, ManagedFolderDocumentRef
    except ImportError:
        logger.warning("DocumentExtractor not available; cannot predict image count")
        return 0

    import dataiku

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

    total_pages = 0

    for attachment in attachments:
        path = attachment.get("document_path")
        if not path:
            continue

        try:
            document_ref = ManagedFolderDocumentRef(path, folder_id)
            # Use generate_pages_screenshots with fetch_size=0 to get metadata only
            # This is much faster than actually generating screenshots
            response = doc_extractor.generate_pages_screenshots(
                document=document_ref,
                output_managed_folder=None,
                fetch_size=1,
            )
            page_count = response.total_count
            total_pages += page_count
            logger.debug(f"Document {attachment.get('document_name', path)} has {page_count} pages")
        except Exception as err:
            logger.warning(f"Failed to extract page count for {path}: {err}")
            # Continue with other documents even if one fails
            continue

    logger.info(f"Predicted {total_pages} total images from {len(attachments)} new attachments")
    return total_pages
