app.py

from dotenv import load_dotenv

load_dotenv()

import gradio as gr
from chonkie import RecursiveChunker
from typing import Dict, Any, List, Optional
from dataclasses import dataclass, field
import logging
import re
import base64
import hashlib
import mimetypes
import json
from collections import Counter
from datasets import Dataset, Features, Value, Sequence, load_dataset
from datasets.features import Image as HFImage
from huggingface_hub import HfApi, get_token
import huggingface_hub
import os
from mistralai import Mistral
import fitz  # pymupdf
from PIL import Image
import io
import tempfile

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


# --- Exceptions ---


class OCRError(Exception):
    """Raised when OCR processing fails."""

    pass


# --- Mistral Client (lazy init) ---

_client: Mistral | None = None


def get_mistral_client() -> Mistral:
    """Get or initialize the Mistral client."""
    global _client
    if _client is not None:
        return _client

    api_key = os.environ.get("MISTRAL_API_KEY")
    if not api_key:
        logger.warning("MISTRAL_API_KEY not set. Attempting to use Hugging Face token.")
        api_key = get_token()
        if api_key:
            logger.info("Using Hugging Face token as MISTRAL_API_KEY.")

    if not api_key:
        raise OCRError(
            "No API key found. Set MISTRAL_API_KEY or run `huggingface-cli login`."
        )

    _client = Mistral(api_key=api_key)
    logger.info("Mistral client initialized successfully.")
    return _client


# --- Helper Functions ---


def encode_image_bytes(image_bytes: bytes) -> str:
    """Encodes image bytes to a base64 string."""
    return base64.b64encode(image_bytes).decode("utf-8")


def decode_base64_data_uri(data_uri: str) -> Optional[dict]:
    """Decode a base64 data URI to a HF-compatible image bytes dict.

    Args:
        data_uri: A string like "data:image/jpeg;base64,/9j/4AAQ..." or raw base64.

    Returns:
        Dict with {"bytes": <raw bytes>, "path": None} for datasets.Image feature,
        or None if decoding fails.
    """
    try:
        if data_uri.startswith("data:"):
            # Strip the "data:image/...;base64," prefix
            _, encoded = data_uri.split(",", 1)
        else:
            encoded = data_uri
        raw_bytes = base64.b64decode(encoded)
        # Validate it's a real image by opening it
        img = Image.open(io.BytesIO(raw_bytes))
        # Re-encode as PNG for consistency
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return {"bytes": buf.getvalue(), "path": None}
    except Exception as e:
        logger.warning(f"Failed to decode base64 image ({len(data_uri)} chars): {e}")
        return None


def extract_images_from_markdown(markdown_text: str) -> Dict[str, str]:
    """
    Extracts base64 image data URIs from markdown and maps them to reference IDs.
    Returns a dictionary mapping reference IDs to base64 data URIs.
    """
    image_map = {}
    img_refs = re.findall(
        r"!\[.*?\]\((data:image/[a-zA-Z+]+;base64,[A-Za-z0-9+/=]+)\)", markdown_text
    )
    for idx, img_uri in enumerate(img_refs):
        ref_id = f"img_ref_{idx + 1}"
        image_map[ref_id] = img_uri
    return image_map


def replace_image_references(markdown_text: str, image_map: Dict[str, str]) -> str:
    """
    Replaces base64 image data URIs in markdown with reference IDs (e.g., img_ref_1).
    """
    updated_markdown = markdown_text
    for ref_id, img_uri in image_map.items():
        escaped_uri = re.escape(img_uri)
        pattern = r"(!\[.*?\]\()" + escaped_uri + r"(\))"
        updated_markdown = re.sub(pattern, f"\\1{ref_id}\\2", updated_markdown)
    return updated_markdown


def get_combined_markdown(ocr_response: Any) -> tuple[str, str, Dict[str, str]]:
    """Combines markdown from OCR pages, replacing image IDs with base64 data URIs."""
    processed_markdowns = []
    raw_markdowns = []
    image_data_map = {}

    if not hasattr(ocr_response, "pages") or not ocr_response.pages:
        logger.warning("OCR response has no pages.")
        return "", "", {}

    for page_idx, page in enumerate(ocr_response.pages):
        if hasattr(page, "images") and page.images:
            logger.info(f"Page {page_idx}: Found {len(page.images)} images.")
            for img in page.images:
                if (
                    hasattr(img, "id")
                    and hasattr(img, "image_base64")
                    and img.image_base64
                ):
                    image_data_map[img.id] = img.image_base64
                else:
                    logger.warning(
                        f"Page {page_idx}: Image object lacks 'id' or valid 'image_base64'."
                    )
        else:
            logger.info(f"Page {page_idx}: No images found.")

        if not hasattr(page, "markdown"):
            logger.warning(f"Page {page_idx} lacks 'markdown' attribute. Skipping.")
            continue

        current_raw_markdown = page.markdown or ""
        raw_markdowns.append(current_raw_markdown)
        current_processed_markdown = current_raw_markdown

        img_refs = re.findall(r"!\[.*?\]\((.*?)\)", current_processed_markdown)
        for img_id in img_refs:
            if img_id in image_data_map:
                base64_data_uri = image_data_map[img_id]
                escaped_img_id = re.escape(img_id)
                pattern = r"(!\[.*?\]\()" + escaped_img_id + r"(\))"
                current_processed_markdown = re.sub(
                    pattern,
                    r"\1" + base64_data_uri + r"\2",
                    current_processed_markdown,
                )
            elif not img_id.startswith(("http:", "https:", "data:")):
                logger.warning(
                    f"Page {page_idx}: Image ID '{img_id}' not in image data."
                )

        processed_markdowns.append(current_processed_markdown)

    logger.info(
        f"Processed {len(processed_markdowns)} pages with {len(image_data_map)} images."
    )
    return "\n\n".join(processed_markdowns), "\n\n".join(raw_markdowns), image_data_map


def perform_ocr(file_path: str) -> tuple[str, str, Dict[str, str]]:
    """Performs OCR on a file using Mistral API.

    Args:
        file_path: Path to the file on disk.

    Returns:
        Tuple of (processed_markdown, raw_markdown, image_data_map).

    Raises:
        OCRError: If OCR processing fails.
    """
    client = get_mistral_client()
    file_name = os.path.basename(file_path)
    file_ext = os.path.splitext(file_name)[1].lower()
    logger.info(f"Performing OCR on file: {file_name}")

    ocr_response = None
    supported_images = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}

    if file_ext == ".pdf":
        uploaded_file_id = None
        try:
            with open(file_path, "rb") as f:
                file_content = f.read()

            logger.info(f"Uploading PDF {file_name} to Mistral...")
            uploaded_pdf = client.files.upload(
                file={"file_name": file_name, "content": file_content},
                purpose="ocr",
            )
            uploaded_file_id = uploaded_pdf.id
            logger.info(f"PDF uploaded. File ID: {uploaded_file_id}")

            signed_url_response = client.files.get_signed_url(file_id=uploaded_file_id)
            ocr_response = client.ocr.process(
                model="mistral-ocr-latest",
                document={
                    "type": "document_url",
                    "document_url": signed_url_response.url,
                },
                include_image_base64=True,
            )
        finally:
            if uploaded_file_id:
                try:
                    client.files.delete(file_id=uploaded_file_id)
                except Exception as delete_err:
                    logger.warning(
                        f"Failed to delete temporary file {uploaded_file_id}: {delete_err}"
                    )

    elif file_ext in supported_images:
        with open(file_path, "rb") as f:
            image_bytes = f.read()
        if not image_bytes:
            raise OCRError(f"Uploaded image file '{file_name}' is empty.")

        base64_encoded = encode_image_bytes(image_bytes)
        mime_type, _ = mimetypes.guess_type(file_path)
        mime_type = mime_type or "image/jpeg"
        data_uri = f"data:{mime_type};base64,{base64_encoded}"
        ocr_response = client.ocr.process(
            model="mistral-ocr-latest",
            document={"type": "image_url", "image_url": data_uri},
            include_image_base64=True,
        )
    else:
        raise OCRError(f"Unsupported file type: '{file_ext}'")

    if not ocr_response:
        raise OCRError(f"OCR returned no response for '{file_name}'.")

    processed_md, raw_md, img_map = get_combined_markdown(ocr_response)
    logger.info(f"Processed markdown length: {len(processed_md)}")
    return processed_md, raw_md, img_map


def _build_header_index(markdown_text: str) -> list[tuple[int, int, str]]:
    """Build a sorted index of (position, level, title) for all markdown headers."""
    headers = []
    for match in re.finditer(r"^(#{1,6})\s+(.+)$", markdown_text, re.MULTILINE):
        level = len(match.group(1))
        title = match.group(2).strip()
        headers.append((match.start(), level, title))
    return headers


def _get_headers_for_position(
    headers: list[tuple[int, int, str]], position: int
) -> dict[str, str]:
    """Given a character position, find the active chapter/section/subsection.

    Maps header levels: H1 -> chapter, H2 -> section, H3+ -> subsection.
    """
    active: dict[int, str] = {}
    for hdr_pos, level, title in headers:
        if hdr_pos > position:
            break
        active[level] = title
        # Clear deeper levels when a higher-level header appears
        for deeper in list(active.keys()):
            if deeper > level:
                del active[deeper]

    return {
        "chapter": active.get(1, ""),
        "section": active.get(2, ""),
        "subsection": active.get(3, active.get(4, active.get(5, active.get(6, "")))),
    }


def _clean_text(text: str) -> str:
    """Remove markdown formatting, image refs, and extra whitespace."""
    cleaned = re.sub(r"!\[.*?\]\(.*?\)", "", text)
    cleaned = re.sub(r"#{1,6}\s+", "", cleaned)
    cleaned = re.sub(r"\*\*(.+?)\*\*", r"\1", cleaned)
    cleaned = re.sub(r"\*(.+?)\*", r"\1", cleaned)
    cleaned = re.sub(r"`(.+?)`", r"\1", cleaned)
    cleaned = re.sub(r"\[(.+?)\]\(.*?\)", r"\1", cleaned)
    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)
    return cleaned.strip()


def chunk_markdown(
    markdown_text_with_images: str,
    chunk_size: int = 512,
) -> list[dict]:
    """Chunks markdown text using chonkie's RecursiveChunker with markdown recipe.

    Args:
        markdown_text_with_images: Markdown text possibly containing base64 image references.
        chunk_size: Maximum character count per chunk.

    Returns:
        List of chunk dicts with the full dataset schema fields.
    """
    if not markdown_text_with_images or not markdown_text_with_images.strip():
        logger.warning("chunk_markdown received empty input.")
        return []

    # Extract images and replace with reference IDs
    image_map = extract_images_from_markdown(markdown_text_with_images)
    updated_markdown = replace_image_references(markdown_text_with_images, image_map)
    logger.info(f"Extracted {len(image_map)} images from markdown.")

    # Build header index for chapter/section/subsection lookup
    header_index = _build_header_index(updated_markdown)

    # Use chonkie's RecursiveChunker with markdown recipe
    chunker = RecursiveChunker.from_recipe(
        "markdown",
        lang="en",
        chunk_size=chunk_size,
    )
    chunks = chunker.chunk(updated_markdown)

    if not chunks:
        logger.warning("No chunks created. Treating entire text as one chunk.")
        all_refs = list(image_map.keys())
        all_images = [
            decoded
            for uri in image_map.values()
            if (decoded := decode_base64_data_uri(uri)) is not None
        ]
        headers = _get_headers_for_position(header_index, 0)
        return [
            {
                "text": updated_markdown,
                "text_clean": _clean_text(updated_markdown),
                "chapter": headers["chapter"],
                "section": headers["section"],
                "subsection": headers["subsection"],
                "images": all_images,
                "image_refs": all_refs,
                "num_images": len(all_images),
                "has_images": len(all_images) > 0,
                "start_index": 0,
                "char_count": len(updated_markdown),
            }
        ]

    result = []
    for chunk in chunks:
        chunk_img_refs = re.findall(r"!\[.*?\]\((img_ref_\d+)\)", chunk.text)
        chunk_images = [
            decoded
            for ref_id in chunk_img_refs
            if ref_id in image_map
            and (decoded := decode_base64_data_uri(image_map[ref_id])) is not None
        ]
        headers = _get_headers_for_position(header_index, chunk.start_index)
        text_clean = _clean_text(chunk.text)

        result.append(
            {
                "text": chunk.text,
                "text_clean": text_clean,
                "chapter": headers["chapter"],
                "section": headers["section"],
                "subsection": headers["subsection"],
                "images": chunk_images,
                "image_refs": chunk_img_refs,
                "num_images": len(chunk_images),
                "has_images": len(chunk_images) > 0,
                "start_index": chunk.start_index,
                "char_count": len(chunk.text),
            }
        )

    logger.info(f"Created {len(result)} chunks.")
    return result


### --- Dataset Builder Pipeline ---


DATASET_SCHEMA = {
    "chunk_id": str,
    "text": str,
    "text_clean": str,
    "chapter": str,
    "section": str,
    "subsection": str,
    "images": list,
    "image_refs": list,
    "num_images": int,
    "has_images": bool,
    "source_filename": str,
    "start_index": int,
    "char_count": int,
}


@dataclass
class QualityConfig:
    """Configuration for quality filtering thresholds."""

    min_char_count: int = 20
    min_clean_char_count: int = 10
    max_char_count: int = 50_000
    min_word_count: int = 3
    max_image_refs_without_text: int = 0
    remove_empty_text: bool = True
    remove_whitespace_only: bool = True


@dataclass
class PipelineStats:
    """Statistics collected during pipeline execution."""

    total_input_chunks: int = 0
    chunks_after_validation: int = 0
    chunks_after_dedup: int = 0
    chunks_after_quality: int = 0
    duplicates_removed: int = 0
    quality_filtered: int = 0
    validation_errors: List[str] = field(default_factory=list)
    quality_reasons: Counter = field(default_factory=Counter)
    source_file_counts: Counter = field(default_factory=Counter)
    avg_char_count: float = 0.0
    avg_images_per_chunk: float = 0.0
    chapters_found: List[str] = field(default_factory=list)

    def summary(self) -> str:
        """Generate a human-readable pipeline summary."""
        lines = [
            "--- Dataset Pipeline Report ---",
            f"Input chunks:            {self.total_input_chunks}",
            f"After validation:        {self.chunks_after_validation}",
            f"Duplicates removed:      {self.duplicates_removed}",
            f"After deduplication:     {self.chunks_after_dedup}",
            f"Quality filtered out:    {self.quality_filtered}",
            f"Final dataset size:      {self.chunks_after_quality}",
            "",
            f"Avg chars/chunk:         {self.avg_char_count:.0f}",
            f"Avg images/chunk:        {self.avg_images_per_chunk:.2f}",
        ]

        if self.source_file_counts:
            lines.append("")
            lines.append("Chunks per source file:")
            for fname, count in sorted(self.source_file_counts.items()):
                lines.append(f"  {fname}: {count}")

        if self.chapters_found:
            unique_chapters = sorted(set(c for c in self.chapters_found if c))
            if unique_chapters:
                lines.append("")
                lines.append(f"Chapters found ({len(unique_chapters)}):")
                for ch in unique_chapters[:20]:
                    lines.append(f"  - {ch}")
                if len(unique_chapters) > 20:
                    lines.append(f"  ... and {len(unique_chapters) - 20} more")

        if self.quality_reasons:
            lines.append("")
            lines.append("Quality filter reasons:")
            for reason, count in self.quality_reasons.most_common():
                lines.append(f"  {reason}: {count}")

        if self.validation_errors:
            lines.append("")
            lines.append(f"Validation errors ({len(self.validation_errors)}):")
            for err in self.validation_errors[:10]:
                lines.append(f"  - {err}")
            if len(self.validation_errors) > 10:
                lines.append(f"  ... and {len(self.validation_errors) - 10} more")

        lines.append("-------------------------------")
        return "\n".join(lines)


class DatasetBuilder:
    """Pipeline for building high-quality datasets before pushing to HF Hub.

    Stages:
        1. Validate  -- ensure every chunk matches the expected schema
        2. Deduplicate -- remove chunks with identical content hashes
        3. Quality filter -- remove empty, too-short, or malformed chunks
        4. Statistics -- compute summary stats for review
        5. Push -- incremental append or full overwrite to HF Hub
    """

    def __init__(
        self,
        quality_config: Optional[QualityConfig] = None,
    ):
        self.quality_config = quality_config or QualityConfig()
        self.stats = PipelineStats()
        self._chunks: List[Dict[str, Any]] = []
        self._seen_hashes: set = set()

    def add_chunks(self, chunks: List[Dict[str, Any]], source_filename: str) -> None:
        """Add raw chunks from a processed file into the pipeline.

        Each chunk gets its source_filename attached and is tracked for stats.
        """
        for chunk in chunks:
            chunk_with_source = {**chunk, "source_filename": source_filename}
            self._chunks.append(chunk_with_source)
        self.stats.source_file_counts[source_filename] += len(chunks)

    def _content_hash(self, chunk: Dict[str, Any]) -> str:
        """Compute a stable hash of chunk content for deduplication."""
        text = chunk.get("text_clean", chunk.get("text", ""))
        source = chunk.get("source_filename", "")
        return hashlib.sha256(f"{source}::{text}".encode("utf-8")).hexdigest()

    # --- Stage 1: Validation ---

    def _validate(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Validate every chunk conforms to the expected schema.

        Drops chunks with missing required fields and logs errors.
        """
        valid = []
        required_keys = set(DATASET_SCHEMA.keys())

        for i, chunk in enumerate(chunks):
            missing = required_keys - set(chunk.keys())
            if missing:
                self.stats.validation_errors.append(
                    f"Chunk {i} ({chunk.get('chunk_id', '?')}): missing fields {missing}"
                )
                continue

            type_ok = True
            for key, expected_type in DATASET_SCHEMA.items():
                val = chunk[key]
                if not isinstance(val, expected_type):
                    self.stats.validation_errors.append(
                        f"Chunk {i} ({chunk.get('chunk_id', '?')}): "
                        f"field '{key}' expected {expected_type.__name__}, "
                        f"got {type(val).__name__}"
                    )
                    type_ok = False
                    break

            if type_ok:
                valid.append(chunk)

        return valid

    # --- Stage 2: Deduplication ---

    def _deduplicate(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Remove chunks with identical content hashes."""
        unique = []
        for chunk in chunks:
            h = self._content_hash(chunk)
            if h not in self._seen_hashes:
                self._seen_hashes.add(h)
                unique.append(chunk)
        return unique

    # --- Stage 3: Quality Filter ---

    def _quality_filter(self, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Filter out low-quality chunks based on configurable thresholds."""
        cfg = self.quality_config
        passed = []

        for chunk in chunks:
            text = chunk.get("text", "")
            text_clean = chunk.get("text_clean", "")
            char_count = chunk.get("char_count", len(text))

            # Empty text
            if cfg.remove_empty_text and not text.strip():
                self.stats.quality_reasons["empty_text"] += 1
                continue

            # Whitespace only
            if cfg.remove_whitespace_only and not text_clean.strip():
                self.stats.quality_reasons["whitespace_only"] += 1
                continue

            # Too short
            if char_count < cfg.min_char_count:
                self.stats.quality_reasons[
                    f"below_min_chars({cfg.min_char_count})"
                ] += 1
                continue

            # Clean text too short
            if len(text_clean.strip()) < cfg.min_clean_char_count:
                self.stats.quality_reasons[
                    f"clean_text_too_short({cfg.min_clean_char_count})"
                ] += 1
                continue

            # Too long (likely malformed)
            if char_count > cfg.max_char_count:
                self.stats.quality_reasons[
                    f"above_max_chars({cfg.max_char_count})"
                ] += 1
                continue

            # Too few words
            word_count = len(text_clean.split())
            if word_count < cfg.min_word_count:
                self.stats.quality_reasons[
                    f"below_min_words({cfg.min_word_count})"
                ] += 1
                continue

            # Image-only chunk with no text
            if (
                cfg.max_image_refs_without_text == 0
                and chunk.get("num_images", 0) > 0
                and word_count == 0
            ):
                self.stats.quality_reasons["image_only_no_text"] += 1
                continue

            passed.append(chunk)

        return passed

    # --- Stage 4: Compute Stats ---

    def _compute_stats(self, chunks: List[Dict[str, Any]]) -> None:
        """Compute summary statistics on the final dataset."""
        if not chunks:
            return

        total_chars = sum(c.get("char_count", 0) for c in chunks)
        total_images = sum(c.get("num_images", 0) for c in chunks)
        self.stats.avg_char_count = total_chars / len(chunks)
        self.stats.avg_images_per_chunk = total_images / len(chunks)
        self.stats.chapters_found = [c.get("chapter", "") for c in chunks]

        # Update source file counts to reflect final dataset
        final_counts: Counter = Counter()
        for c in chunks:
            final_counts[c.get("source_filename", "unknown")] += 1
        self.stats.source_file_counts = final_counts

    # --- Run Full Pipeline ---

    def build(self) -> tuple[Dict[str, list], PipelineStats]:
        """Run the full pipeline and return columnar data + stats.

        Returns:
            Tuple of (columnar_data_dict, pipeline_stats).
        """
        chunks = list(self._chunks)
        self.stats.total_input_chunks = len(chunks)
        logger.info(f"Pipeline: {len(chunks)} input chunks")

        # Stage 1: Validate
        chunks = self._validate(chunks)
        self.stats.chunks_after_validation = len(chunks)
        logger.info(f"Pipeline: {len(chunks)} after validation")

        # Stage 2: Deduplicate
        before_dedup = len(chunks)
        chunks = self._deduplicate(chunks)
        self.stats.duplicates_removed = before_dedup - len(chunks)
        self.stats.chunks_after_dedup = len(chunks)
        logger.info(
            f"Pipeline: {len(chunks)} after dedup ({self.stats.duplicates_removed} removed)"
        )

        # Stage 3: Quality filter
        before_quality = len(chunks)
        chunks = self._quality_filter(chunks)
        self.stats.quality_filtered = before_quality - len(chunks)
        self.stats.chunks_after_quality = len(chunks)
        logger.info(
            f"Pipeline: {len(chunks)} after quality filter ({self.stats.quality_filtered} removed)"
        )

        # Stage 4: Stats
        self._compute_stats(chunks)

        # Convert to columnar format
        all_data: Dict[str, list] = {key: [] for key in DATASET_SCHEMA.keys()}
        for chunk in chunks:
            for key in DATASET_SCHEMA.keys():
                all_data[key].append(chunk[key])

        return all_data, self.stats

    # --- Push to Hub ---

    @staticmethod
    def push(
        all_data: Dict[str, list],
        repo_name: str,
        hf_token: str,
        stats: PipelineStats,
        append: bool = False,
    ) -> str:
        """Push the built dataset to Hugging Face Hub.

        Args:
            all_data: Columnar data dict from build().
            repo_name: HF repo in 'username/dataset-name' format.
            hf_token: Hugging Face API token.
            stats: Pipeline stats for the dataset card.
            append: If True, append to existing dataset instead of overwriting.

        Returns:
            Status message string.
        """
        if not all_data or not all_data.get("chunk_id"):
            return "Error: No data to push after pipeline."

        api = HfApi(token=hf_token)

        try:
            user_info = api.whoami()
            logger.info(f"Authenticated as: {user_info['name']}")
        except Exception as auth_err:
            return f"Error: Invalid HF token - authentication failed: {auth_err}"

        # Create repo if needed
        try:
            api.repo_info(repo_id=repo_name, repo_type="dataset")
            logger.info(f"Repository '{repo_name}' exists.")
        except huggingface_hub.utils.RepositoryNotFoundError:
            api.create_repo(repo_id=repo_name, repo_type="dataset", private=False)
            logger.info(f"Created repository '{repo_name}'.")

        if append:
            # Incremental append: load existing, concatenate, push
            try:
                existing_ds = load_dataset(repo_name, token=hf_token, split="train")
                existing_data = existing_ds.to_dict()
                for key in all_data:
                    if key in existing_data:
                        existing_data[key].extend(all_data[key])
                    else:
                        existing_data[key] = all_data[key]
                # Deduplicate by chunk_id across old + new
                seen_ids = set()
                deduped: Dict[str, list] = {key: [] for key in existing_data}
                for i, cid in enumerate(existing_data["chunk_id"]):
                    if cid not in seen_ids:
                        seen_ids.add(cid)
                        for key in existing_data:
                            deduped[key].append(existing_data[key][i])
                merged_dataset = Dataset.from_dict(deduped)
                total_chunks = len(deduped["chunk_id"])
                new_chunks = total_chunks - len(existing_ds)
                commit_msg = f"Append {new_chunks} new chunks (total: {total_chunks})"
            except Exception as e:
                logger.warning(
                    f"Could not load existing dataset for append, doing full push: {e}"
                )
                merged_dataset = Dataset.from_dict(all_data)
                total_chunks = len(all_data["chunk_id"])
                commit_msg = f"Add {total_chunks} chunks"
        else:
            merged_dataset = Dataset.from_dict(all_data)
            total_chunks = len(all_data["chunk_id"])
            commit_msg = f"Add OCR data: {total_chunks} chunks"

        # Cast the images column so the HF Dataset Viewer renders actual images
        # instead of showing raw base64 strings
        try:
            merged_dataset = merged_dataset.cast_column("images", Sequence(HFImage()))
            logger.info(
                "Cast 'images' column to Sequence(Image()) for viewer rendering."
            )
        except Exception as e:
            logger.warning(f"Could not cast images column to Image feature: {e}")

        merged_dataset.push_to_hub(
            repo_name,
            token=hf_token,
            commit_message=commit_msg,
        )

        # Generate and upload dataset card
        card_content = DatasetBuilder._generate_dataset_card(repo_name, stats, all_data)
        try:
            api.upload_file(
                path_or_fileobj=card_content.encode("utf-8"),
                path_in_repo="README.md",
                repo_id=repo_name,
                repo_type="dataset",
                commit_message="Update dataset card with pipeline stats",
            )
        except Exception as e:
            logger.warning(f"Failed to update dataset card: {e}")

        repo_url = f"https://huggingface.co/datasets/{repo_name}"
        return f"Success! {total_chunks} chunks pushed to: {repo_url}"

    @staticmethod
    def _generate_dataset_card(
        repo_name: str,
        stats: PipelineStats,
        all_data: Dict[str, list],
    ) -> str:
        """Generate a dataset card (README.md) with schema and stats."""
        total = stats.chunks_after_quality
        sources = sorted(stats.source_file_counts.items())
        unique_chapters = sorted(set(c for c in stats.chapters_found if c))

        card = f"""---
license: mit
task_categories:
  - text-generation
  - question-answering
language:
  - en
tags:
  - pdf2dataset
  - ocr
  - chunked
size_categories:
  - {"1K<n<10K" if total >= 1000 else "n<1K"}
---

# {repo_name.split("/")[-1]}

Dataset created with [PDF2Dataset](https://github.com/svngoku/PDF2Dataset) -- OCR + structure-aware chunking pipeline.

## Dataset Summary

| Metric | Value |
|---|---|
| Total chunks | {total} |
| Avg chars/chunk | {stats.avg_char_count:.0f} |
| Avg images/chunk | {stats.avg_images_per_chunk:.2f} |
| Source files | {len(sources)} |
| Duplicates removed | {stats.duplicates_removed} |
| Quality filtered | {stats.quality_filtered} |

## Schema

| Column | Type | Description |
|---|---|---|
| `chunk_id` | `string` | Unique identifier: `filename_chunk_N` |
| `text` | `string` | Raw markdown chunk with image refs |
| `text_clean` | `string` | Cleaned text without markdown formatting |
| `chapter` | `string` | H1 header active at chunk position |
| `section` | `string` | H2 header active at chunk position |
| `subsection` | `string` | H3+ header active at chunk position |
| `images` | `list[Image]` | Rendered images extracted from chunk (viewable in Dataset Viewer) |
| `image_refs` | `list[string]` | Image reference IDs in chunk text |
| `num_images` | `int` | Number of images in chunk |
| `has_images` | `bool` | Whether chunk contains images |
| `source_filename` | `string` | Original source file name |
| `start_index` | `int` | Character offset in source document |
| `char_count` | `int` | Character count of chunk text |

## Source Files

| File | Chunks |
|---|---|
"""
        for fname, count in sources:
            card += f"| `{fname}` | {count} |\n"

        if unique_chapters:
            card += "\n## Document Structure\n\n"
            card += "Chapters found in the source documents:\n\n"
            for ch in unique_chapters[:30]:
                card += f"- {ch}\n"
            if len(unique_chapters) > 30:
                card += f"- ... and {len(unique_chapters) - 30} more\n"

        card += """
## Pipeline

This dataset was processed through the PDF2Dataset pipeline:

1. **OCR** -- Mistral OCR extracts text and images from PDF/image files
2. **Chunking** -- Structure-aware recursive splitting preserves document hierarchy
3. **Validation** -- Schema validation ensures every chunk has required fields
4. **Deduplication** -- Content-hash based dedup removes identical chunks
5. **Quality Filtering** -- Removes empty, too-short, or malformed chunks
"""
        return card


def get_hf_token(explicit_token: str | None = None) -> str | None:
    """Retrieve Hugging Face token with fallback mechanisms."""
    if explicit_token and explicit_token.strip() and explicit_token.startswith("hf_"):
        return explicit_token.strip()

    env_token = os.environ.get("HF_TOKEN")
    if env_token and env_token.startswith("hf_"):
        return env_token

    try:
        stored_token = huggingface_hub.get_token()
        if stored_token:
            return stored_token
    except Exception as e:
        logger.warning(f"Could not retrieve token from Hugging Face config: {e}")

    return None


def process_files(
    file_paths: list[str],
    chunk_size: int,
    hf_token: str | None,
    repo_name: str,
    append_mode: bool = False,
    min_chunk_chars: int = 20,
    min_words: int = 3,
) -> str:
    """Orchestrates OCR, chunking, pipeline processing, and push to HF Hub.

    Pipeline stages:
        1. OCR each file with Mistral
        2. Chunk markdown with structure-aware splitting
        3. Validate schema on every chunk
        4. Deduplicate by content hash
        5. Quality-filter (min chars, min words, empty, etc.)
        6. Compute statistics and generate report
        7. Push to HF Hub (overwrite or append)

    Args:
        file_paths: List of file paths to process.
        chunk_size: Maximum character count per chunk.
        hf_token: Explicit HF token (optional).
        repo_name: HF dataset repository in 'username/dataset-name' format.
        append_mode: If True, append to existing dataset instead of replacing.
        min_chunk_chars: Minimum characters per chunk for quality filter.
        min_words: Minimum words per chunk for quality filter.

    Returns:
        Status message string with pipeline report.
    """
    if not file_paths:
        return "Error: No files uploaded."

    if not repo_name or "/" not in repo_name:
        return "Error: Invalid repository name (use 'username/dataset-name')."

    chunk_size = max(0, chunk_size)

    effective_hf_token = get_hf_token(hf_token)
    if not effective_hf_token:
        return (
            "Error: No valid Hugging Face token found.\n"
            "Please either:\n"
            "1. Sign in with Hugging Face when using the Space\n"
            "2. Set HF_TOKEN environment variable\n"
            "3. Run `huggingface-cli login` in your terminal"
        )

    try:
        # Initialize pipeline with quality config
        quality_cfg = QualityConfig(
            min_char_count=min_chunk_chars,
            min_word_count=min_words,
        )
        builder = DatasetBuilder(quality_config=quality_cfg)

        files_processed = 0
        error_messages = []

        for file_idx, file_path in enumerate(file_paths, 1):
            source_filename = os.path.basename(file_path)
            logger.info(
                f"--- Processing file {file_idx}/{len(file_paths)}: {source_filename} ---"
            )

            try:
                processed_markdown, raw_markdown, img_map = perform_ocr(file_path)
            except OCRError as e:
                error_messages.append(f"File '{source_filename}': {e}")
                logger.error(f"Failed to process file {source_filename}: {e}")
                continue

            chunks = chunk_markdown(processed_markdown, chunk_size)
            if not chunks:
                error_messages.append(
                    f"File '{source_filename}': Failed to chunk the document."
                )
                logger.error(f"Failed to chunk file {source_filename}")
                continue

            # Assign chunk_id before adding to pipeline
            for i, chunk in enumerate(chunks):
                chunk["chunk_id"] = f"{source_filename}_chunk_{i}"

            builder.add_chunks(chunks, source_filename)
            files_processed += 1
            logger.info(
                f"File {source_filename}: queued {len(chunks)} chunks for pipeline"
            )

        if files_processed == 0:
            return "Error: No files were processed successfully.\n" + "\n".join(
                error_messages
            )

        # Run the pipeline
        all_data, stats = builder.build()

        if not all_data or not all_data.get("chunk_id"):
            return (
                "Error: All chunks were filtered out by the pipeline.\n"
                + stats.summary()
                + (
                    "\n\nOCR Errors:\n" + "\n".join(error_messages)
                    if error_messages
                    else ""
                )
            )

        # Push to Hub
        push_result = DatasetBuilder.push(
            all_data=all_data,
            repo_name=repo_name,
            hf_token=effective_hf_token,
            stats=stats,
            append=append_mode,
        )

        # Build final report
        report_parts = [push_result, "", stats.summary()]
        if error_messages:
            report_parts.append(f"\nOCR Errors ({len(error_messages)}):")
            report_parts.extend(f"  - {e}" for e in error_messages)

        return "\n".join(report_parts)

    except huggingface_hub.utils.HfHubHTTPError as hf_http_err:
        status = getattr(hf_http_err.response, "status_code", "Unknown")
        if status == 401:
            return "Error: Invalid or unauthorized Hugging Face token."
        elif status == 403:
            return "Error: Token lacks write permission."
        return f"Error: Hugging Face Hub Error (Status {status}): {hf_http_err}"
    except Exception as e:
        logger.error(f"Unexpected error: {e}", exc_info=True)
        return f"Unexpected error: {e}"


# --- Preview ---


def render_preview(file_objs) -> list[Image.Image]:
    """Render uploaded files as preview images.

    PDFs are rendered page-by-page using PyMuPDF. Images are returned directly.
    """
    if not file_objs:
        return []
    if not isinstance(file_objs, list):
        file_objs = [file_objs]

    images = []
    for file_obj in file_objs:
        file_path = file_obj.name if hasattr(file_obj, "name") else str(file_obj)
        ext = os.path.splitext(file_path)[1].lower()

        if ext == ".pdf":
            try:
                doc = fitz.open(file_path)
                for page in doc:
                    pix = page.get_pixmap(dpi=150)
                    img = Image.open(io.BytesIO(pix.tobytes("png")))
                    images.append(img)
                doc.close()
            except Exception as e:
                logger.error(f"Failed to render PDF preview: {e}")
        elif ext in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}:
            try:
                images.append(Image.open(file_path))
            except Exception as e:
                logger.error(f"Failed to open image preview: {e}")

    return images


# --- Gradio Interface ---


MISTRAL_CSS = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;500;600&family=Inter:wght@400;500;600;700;800&display=swap');

:root {
    --mistral-bg: #FFFAEB;
    --mistral-bg-grid: #E9E2CB;
    --mistral-panel: #FFFAEB;
    --mistral-panel-warm: #FFF0C3;
    --mistral-border: #E9E2CB;
    --mistral-text: #1E1E1E;
    --mistral-muted: #444444;
    --mistral-soft-muted: #766B54;
    --mistral-accent: #FF8205;
    --mistral-accent-hover: #E67200;
    --mistral-shadow: rgba(30, 30, 30, 0.08);
    --mistral-grid-opacity: 0.05;
}

html,
body,
gradio-app,
.gradio-container,
.app,
main {
    background-color: var(--mistral-bg) !important;
    background-image:
        linear-gradient(var(--mistral-bg-grid) 1px, transparent 1px),
        linear-gradient(90deg, var(--mistral-bg-grid) 1px, transparent 1px) !important;
    background-size: 40px 40px !important;
    color: var(--mistral-text) !important;
    font-family: 'Inter', sans-serif !important;
}

html,
body,
gradio-app {
    min-height: 100% !important;
    width: 100% !important;
}

.gradio-container {
    box-sizing: border-box !important;
    margin: 0 auto !important;
    max-width: 1440px !important;
    padding: clamp(1rem, 2.5vw, 2rem) !important;
    width: 100% !important;
}

.app,
main {
    min-height: 100vh !important;
    max-width: 100% !important;
    width: 100% !important;
}

footer {
    display: none !important;
}

#app-shell {
    gap: 1rem !important;
}

#brand-hero {
    background: linear-gradient(135deg, var(--mistral-panel) 0%, var(--mistral-panel-warm) 100%) !important;
    border: 2px solid var(--mistral-border) !important;
    border-top: 5px solid var(--mistral-accent) !important;
    box-shadow: 0 8px 32px var(--mistral-shadow) !important;
    padding: clamp(1.25rem, 3vw, 2rem) !important;
}

#brand-hero h1 {
    color: var(--mistral-text) !important;
    font-size: clamp(2rem, 4vw, 4.5rem) !important;
    font-weight: 800 !important;
    line-height: 0.95 !important;
    letter-spacing: 0 !important;
    margin: 0 0 0.75rem !important;
}

#brand-hero p {
    color: var(--mistral-muted) !important;
    font-size: clamp(1rem, 1.6vw, 1.25rem) !important;
    line-height: 1.55 !important;
    margin: 0 !important;
    max-width: 62rem !important;
}

#brand-hero strong {
    color: var(--mistral-text) !important;
    font-weight: 700 !important;
}

.mistral-panel {
    background: var(--mistral-panel) !important;
    border: 2px solid var(--mistral-border) !important;
    box-shadow: 0 8px 32px var(--mistral-shadow) !important;
    padding: clamp(1rem, 2vw, 1.5rem) !important;
}

.mistral-panel .markdown h3,
.mistral-section h3 {
    color: var(--mistral-text) !important;
    font-size: 0.82rem !important;
    font-weight: 800 !important;
    letter-spacing: 0.08em !important;
    margin: 0.35rem 0 0.85rem !important;
    text-transform: uppercase !important;
}

.mistral-section {
    border-top: 1px solid var(--mistral-border) !important;
    margin-top: 1rem !important;
    padding-top: 1rem !important;
}

.mistral-panel label,
.mistral-panel .wrap label,
.mistral-panel span {
    color: var(--mistral-text) !important;
    font-family: 'Inter', sans-serif !important;
}

.mistral-panel input,
.mistral-panel textarea,
.mistral-panel select {
    background: #FFF8E3 !important;
    border-color: var(--mistral-border) !important;
    color: var(--mistral-text) !important;
    font-family: 'Inter', sans-serif !important;
}

.mistral-panel input:focus,
.mistral-panel textarea:focus {
    border-color: var(--mistral-accent) !important;
    box-shadow: 0 0 0 2px rgba(255, 130, 5, 0.18) !important;
}

#process-button {
    background: var(--mistral-accent) !important;
    border: 0 !important;
    border-radius: 0 !important;
    color: #FFFFFF !important;
    font-weight: 800 !important;
    letter-spacing: 0.06em !important;
    min-height: 3rem !important;
    text-transform: uppercase !important;
}

#process-button:hover {
    background: var(--mistral-accent-hover) !important;
}

.mistral-panel .gr-group,
.mistral-panel .styler {
    background: transparent !important;
    border-color: var(--mistral-border) !important;
}

.mistral-panel button:not(#process-button):not(.reset-button):not(.center) {
    background: #FFF8E3 !important;
    border: 1px solid var(--mistral-border) !important;
    color: var(--mistral-text) !important;
    font-weight: 700 !important;
    letter-spacing: 0 !important;
    text-transform: none !important;
}

.mistral-panel button:not(#process-button):not(.reset-button):not(.center):hover {
    background: var(--mistral-panel-warm) !important;
    border-color: var(--mistral-accent) !important;
}

.mistral-panel button.center.boundedheight {
    border: 2px dashed var(--mistral-border) !important;
    color: var(--mistral-muted) !important;
    min-height: 12rem !important;
}

.mistral-panel button.center.boundedheight svg {
    color: var(--mistral-accent) !important;
}

.mistral-panel button.reset-button {
    background: transparent !important;
    border: 0 !important;
    color: var(--mistral-soft-muted) !important;
    min-height: auto !important;
}

#pipeline-report textarea {
    background-color: var(--mistral-panel) !important;
    background-image:
        linear-gradient(rgba(0, 0, 0, var(--mistral-grid-opacity)) 1px, transparent 1px),
        linear-gradient(90deg, rgba(0, 0, 0, var(--mistral-grid-opacity)) 1px, transparent 1px) !important;
    background-size: 20px 20px !important;
    color: var(--mistral-text) !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 0.95rem !important;
    line-height: 1.7 !important;
}

#document-preview {
    background: var(--mistral-panel) !important;
    border: 2px solid var(--mistral-border) !important;
    box-shadow: 0 8px 32px var(--mistral-shadow) !important;
    padding: clamp(1rem, 2vw, 1.5rem) !important;
}

#document-preview h3 {
    color: var(--mistral-text) !important;
    font-size: 0.82rem !important;
    font-weight: 800 !important;
    letter-spacing: 0.08em !important;
    text-transform: uppercase !important;
}

#document-preview .grid-wrap,
#document-preview .thumbnail-item {
    background: #FFF8E3 !important;
}

.mistral-note {
    color: var(--mistral-soft-muted) !important;
    font-size: 0.9rem !important;
    margin-top: 0.5rem !important;
}
"""


MISTRAL_THEME = gr.themes.Soft(primary_hue="orange", secondary_hue="yellow")
GRADIO_BLOCKS_KWARGS = {"title": "PDF2Dataset -- Mistral OCR Pipeline"}
GRADIO_LAUNCH_KWARGS = {}

try:
    _GRADIO_MAJOR_VERSION = int(gr.__version__.split(".", 1)[0])
except (AttributeError, ValueError):
    _GRADIO_MAJOR_VERSION = 5

if _GRADIO_MAJOR_VERSION >= 6:
    GRADIO_LAUNCH_KWARGS.update(theme=MISTRAL_THEME, css=MISTRAL_CSS)
else:
    GRADIO_BLOCKS_KWARGS.update(theme=MISTRAL_THEME, css=MISTRAL_CSS)


def _gradio_process(
    file_objs: list[str] | str | None,
    chunk_size: int,
    repo_name: str,
    append_mode: bool,
    min_chars: int,
    min_words: int,
    oauth_token: gr.OAuthToken | None,
) -> str:
    """Bridge between Gradio inputs, the signed-in user, and core processing."""
    if oauth_token is None:
        return "Error: Sign in with Hugging Face before processing files."
    if not file_objs:
        return "Error: No files uploaded."
    if not isinstance(file_objs, list):
        file_objs = [file_objs]
    file_paths = [f.name if hasattr(f, "name") else str(f) for f in file_objs]
    return process_files(
        file_paths,
        chunk_size,
        oauth_token.token,
        repo_name,
        append_mode=append_mode,
        min_chunk_chars=min_chars,
        min_words=min_words,
    )


with gr.Blocks(**GRADIO_BLOCKS_KWARGS) as demo:
    gr.Markdown(
        """
        # PDF2Dataset

        Convert PDFs and images into clean Hugging Face datasets with **Mistral OCR**,
        structure-aware chunking, validation, deduplication, and quality filtering.
        """,
        elem_id="brand-hero",
    )

    with gr.Column(elem_id="app-shell"):
        with gr.Row():
            with gr.Column(scale=1, elem_classes=["mistral-panel"]):
                file_input = gr.File(
                    label="Source documents",
                    file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp", ".bmp"],
                    type="filepath",
                    file_count="multiple",
                )

                with gr.Group(elem_classes=["mistral-section"]):
                    gr.Markdown("### Chunking")
                    chunk_size = gr.Slider(
                        minimum=0,
                        maximum=4096,
                        value=512,
                        step=64,
                        label="Max chunk size",
                        info="Character budget per chunk. Use 0 to keep each document as one chunk.",
                    )

                with gr.Group(elem_classes=["mistral-section"]):
                    gr.Markdown("### Quality filters")
                    with gr.Row():
                        min_chars = gr.Slider(
                            minimum=0,
                            maximum=500,
                            value=20,
                            step=5,
                            label="Minimum characters",
                        )
                        min_words = gr.Slider(
                            minimum=0,
                            maximum=50,
                            value=3,
                            step=1,
                            label="Minimum words",
                        )

                with gr.Group(elem_classes=["mistral-section"]):
                    gr.Markdown("### Hugging Face output")
                    gr.LoginButton()
                    repo_name = gr.Textbox(
                        label="Dataset repository",
                        placeholder="your-username/your-dataset-name",
                    )
                    append_mode = gr.Checkbox(
                        label="Append to existing dataset",
                        value=False,
                    )

                submit_btn = gr.Button(
                    "Process and push",
                    variant="primary",
                    elem_id="process-button",
                )

            with gr.Column(scale=1, elem_classes=["mistral-panel"]):
                output = gr.Textbox(
                    label="Pipeline report",
                    lines=30,
                    interactive=False,
                    elem_id="pipeline-report",
                )

        with gr.Group(elem_id="document-preview"):
            gr.Markdown("### Document preview")
            preview_gallery = gr.Gallery(
                label="Uploaded documents",
                columns=2,
                height="auto",
                object_fit="contain",
            )

        gr.Markdown(
            "*Requires `MISTRAL_API_KEY`. Sign in with Hugging Face before pushing a dataset.*",
            elem_classes=["mistral-note"],
        )

    file_input.change(
        fn=render_preview,
        inputs=[file_input],
        outputs=[preview_gallery],
    )

    submit_btn.click(
        fn=_gradio_process,
        inputs=[
            file_input,
            chunk_size,
            repo_name,
            append_mode,
            min_chars,
            min_words,
        ],
        outputs=output,
    )

    gr.Examples(
        examples=[
            [None, 512, "hf-username/my-first-ocr-dataset", False, 20, 3],
            [None, 1024, "hf-username/large-chunk-ocr-data", True, 50, 5],
            [None, 0, "hf-username/no-split-ocr-data", False, 0, 0],
        ],
        inputs=[
            file_input,
            chunk_size,
            repo_name,
            append_mode,
            min_chars,
            min_words,
        ],
        outputs=output,
        fn=_gradio_process,
        cache_examples=False,
    )


def main():
    """Entry point for the application."""
    demo.launch(
        share=os.getenv("GRADIO_SHARE", "False").lower() == "true",
        debug=True,
        **GRADIO_LAUNCH_KWARGS,
    )


if __name__ == "__main__":
    main()