#!/usr/bin/env python3
"""
Local GPU OCR Worker for OCR Provenance MCP System

Replaces python/ocr_worker.py (Datalab cloud API) with Marker-pdf v1.10.2
for 100% local GPU processing. NO cloud API calls, NO fallbacks.

FAIL-FAST: Errors propagate immediately with clear categories.
ALL logging to stderr (stdout is JSON-RPC protocol).
"""

import argparse
import base64
import csv
import gc
import hashlib
import io
import json
import logging
import os
import re
import signal
import sys
import time
import uuid
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional, Tuple, Union

# Configure logging FIRST - ALL to stderr, NEVER stdout
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    stream=sys.stderr,
)
logger = logging.getLogger(__name__)


# =============================================================================
# ERROR CLASSES
# =============================================================================


class LocalOCRError(Exception):
    """Base local OCR error with category for structured error handling."""

    def __init__(self, message: str, category: str, details: Optional[dict] = None):
        super().__init__(message)
        self.category = category
        self.details = details or {}


class OCRFileError(LocalOCRError):
    """File access or validation errors."""

    def __init__(self, message: str, file_path: str):
        super().__init__(message, "OCR_FILE_ERROR", {"file_path": file_path})


class OCRTimeoutError(LocalOCRError):
    """Processing timeout exceeded."""

    def __init__(self, message: str, timeout_seconds: int):
        super().__init__(message, "OCR_TIMEOUT", {"timeout_seconds": timeout_seconds})


class OCRModelError(LocalOCRError):
    """Model loading or inference errors."""

    def __init__(self, message: str, details: Optional[dict] = None):
        super().__init__(message, "OCR_MODEL_ERROR", details)


class OCRGPUError(LocalOCRError):
    """GPU memory or CUDA errors."""

    def __init__(self, message: str, details: Optional[dict] = None):
        super().__init__(message, "OCR_GPU_ERROR", details)


# =============================================================================
# DATA STRUCTURES (match src/models/document.ts exactly)
# =============================================================================


@dataclass
class PageOffset:
    """
    Character offset for a single page.
    MUST match src/models/document.ts PageOffset interface.
    """

    page: int  # 1-indexed page number
    char_start: int  # Start offset in full text
    char_end: int  # End offset in full text


@dataclass
class OCRResult:
    """
    Result from local OCR processing.
    MUST match PythonOCRResponse interface in src/services/ocr/marker.ts.
    """

    # Required fields
    id: str
    provenance_id: str
    document_id: str
    extracted_text: str
    text_length: int
    datalab_request_id: str  # "local-marker-{uuid}"
    datalab_mode: Literal["fast", "balanced", "accurate"]
    parse_quality_score: Optional[float]
    page_count: int
    cost_cents: Optional[float]  # Always 0 for local
    content_hash: str  # sha256:...
    processing_started_at: str  # ISO 8601
    processing_completed_at: str  # ISO 8601
    processing_duration_ms: int

    # Page offsets
    page_offsets: list[PageOffset]
    error: Optional[str] = None

    # Images extracted by Marker (filename -> base64 data)
    images: Optional[Dict[str, str]] = None

    # JSON block hierarchy from Marker JSON renderer
    json_blocks: Optional[dict] = None

    # Marker metadata (table_of_contents, page_stats, etc.)
    metadata: Optional[dict] = None

    # Always null for local processing (requires Google API key)
    extraction_json: Union[dict, list, None] = None

    # Always null for local processing (no cloud costs)
    cost_breakdown_full: Optional[dict] = None

    # Not applicable for local processing
    extras_features: Optional[dict] = None

    # Document metadata from PDF
    doc_title: Optional[str] = None
    doc_author: Optional[str] = None
    doc_subject: Optional[str] = None


# =============================================================================
# SUPPORTED FILE TYPES (Marker-pdf supported formats)
# =============================================================================

SUPPORTED_EXTENSIONS = frozenset(
    {
        ".pdf",
        ".png",
        ".jpg",
        ".jpeg",
        ".tiff",
        ".tif",
        ".bmp",
        ".gif",
        ".webp",
        ".docx",
        ".doc",
        ".pptx",
        ".ppt",
        ".xlsx",
        ".xls",
        ".epub",
        ".html",
        ".htm",
        ".txt",
        ".csv",
        ".md",
    }
)

# Marker page separator: 48 dashes
PAGE_SEPARATOR = "-" * 48

# Max rows before spreadsheet uses direct-to-markdown conversion (skip weasyprint)
SPREADSHEET_MAX_ROWS_FOR_MARKER = 500
# Max rows per "page" for direct spreadsheet conversion
SPREADSHEET_ROWS_PER_PAGE = 50


# =============================================================================
# CUDA STARTUP VALIDATION (runs on module load)
# =============================================================================

try:
    from gpu_utils import clear_gpu_memory, startup_gpu
except ImportError:
    from .gpu_utils import clear_gpu_memory, startup_gpu

# Device startup: configure memory pool, validate environment, log device info
startup_gpu("OCR worker")


# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================


def _output_json(data: dict[str, Any]) -> None:
    """Output a single JSON line to stdout and flush. Used by daemon mode."""
    print(json.dumps(data, ensure_ascii=False), file=sys.stdout, flush=True)


class TqdmProgressReporter:
    """Intercepts tqdm stderr output and emits structured PROGRESS_JSON lines.

    tqdm progress bars from Surya models (Recognizing Layout, Recognizing Text, etc.)
    are parsed and re-emitted as JSON for the TypeScript side to capture.
    """

    _TQDM_RE = re.compile(r"^(.*?):\s*(\d+)%\|.*?\|\s*(\d+)/(\d+)")
    _current_doc_id: str = ""

    def __init__(self) -> None:
        self._original_stderr = sys.__stderr__

    def set_document_id(self, doc_id: str) -> None:
        self._current_doc_id = doc_id

    def write(self, s: str) -> None:
        self._original_stderr.write(s)
        if not self._current_doc_id:
            return
        for line in s.split("\n"):
            m = self._TQDM_RE.search(line)
            if m:
                self._original_stderr.write(
                    f"PROGRESS_JSON:{json.dumps({'document_id': self._current_doc_id, 'stage': m.group(1).strip(), 'percent': int(m.group(2)), 'current': int(m.group(3)), 'total': int(m.group(4))})}\n"
                )
                self._original_stderr.flush()

    def flush(self) -> None:
        self._original_stderr.flush()


# Install the progress reporter globally so all tqdm output is intercepted
_progress_reporter = TqdmProgressReporter()
sys.stderr = _progress_reporter


def output_error(category: str, message: str, details: Optional[dict] = None) -> None:
    """Output structured error JSON to stdout and exit."""
    json.dump(
        {"error": message, "category": category, "details": details or {}},
        sys.stdout,
    )
    sys.stdout.write("\n")
    sys.stdout.flush()
    sys.exit(1)


def validate_file(file_path: str) -> Path:
    """
    Validate file exists and is supported type.
    FAIL-FAST: Raises immediately on any issue.
    """
    path = Path(file_path).resolve()

    if not path.exists():
        raise OCRFileError(f"File not found: {file_path}", str(path))

    if not path.is_file():
        raise OCRFileError(f"Not a file: {file_path}", str(path))

    if path.suffix.lower() not in SUPPORTED_EXTENSIONS:
        raise OCRFileError(
            f"Unsupported file type: {path.suffix}. "
            f"Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}",
            str(path),
        )

    return path


def compute_content_hash(content: str) -> str:
    """
    Compute SHA-256 hash matching src/utils/hash.ts format.
    Returns: 'sha256:' + 64 lowercase hex characters
    """
    hash_hex = hashlib.sha256(content.encode("utf-8")).hexdigest()
    return f"sha256:{hash_hex}"


def check_timeout(start_time: float, timeout: int, phase: str) -> None:
    """
    Cooperative timeout check between processing phases.
    The TS caller (python-shell) handles hard kill via SIGTERM/SIGKILL,
    but we check cooperatively between phases for clean error messages.
    """
    elapsed = time.time() - start_time
    if elapsed >= timeout:
        raise OCRTimeoutError(
            f"OCR processing timed out after {elapsed:.0f}s (limit: {timeout}s) "
            f"during {phase}. Try reducing document size with page_range or max_pages.",
            timeout,
        )


def parse_marker_page_offsets(markdown: str) -> list[PageOffset]:
    """
    Parse page delimiters from Marker paginated markdown output.

    Marker with paginate_output=True adds markers like:
    \\n\\n{page_id}------------------------------------------------\\n\\n

    where page_id is 0-indexed. We convert to 1-indexed for consistency
    with the existing system.

    Note: Marker always emits a page separator for page 0, so the text
    typically starts with \\n\\n{0}---...\\n\\n even for single-page docs.
    We handle this by collecting all page segments and ignoring any
    empty prefix before the first separator.
    """
    # Pattern matches: \n\n{N}(48 dashes)\n\n
    page_pattern = r"\n\n\{(\d+)\}" + re.escape(PAGE_SEPARATOR) + r"\n\n"

    parts = re.split(page_pattern, markdown)

    if len(parts) == 1:
        # No page markers = single page document
        return [PageOffset(page=1, char_start=0, char_end=len(markdown))]

    offsets = []
    current_offset = 0

    # First part is content before any page marker (usually empty/whitespace)
    prefix = parts[0]
    if prefix.strip():
        # Non-empty content before first marker = treat as page 1
        offsets.append(PageOffset(page=1, char_start=0, char_end=len(prefix)))
    current_offset = len(prefix)

    # Subsequent parts: alternating page_id, content
    for i in range(1, len(parts), 2):
        if i + 1 < len(parts):
            page_id = int(parts[i])  # 0-indexed from Marker
            content = parts[i + 1]
            # Calculate the actual separator length in the original text
            marker_text = f"\n\n{{{page_id}}}{PAGE_SEPARATOR}\n\n"
            marker_len = len(marker_text)
            char_start = current_offset + marker_len
            char_end = char_start + len(content)
            offsets.append(
                PageOffset(
                    page=page_id + 1,  # Convert to 1-indexed
                    char_start=char_start,
                    char_end=char_end,
                )
            )
            current_offset = char_end

    # If no offsets were created (shouldn't happen), fallback
    if not offsets:
        return [PageOffset(page=1, char_start=0, char_end=len(markdown))]

    return offsets


def compute_quality_score(
    text: str, page_count: int, json_blocks: Optional[dict]
) -> float:
    """
    Compute heuristic quality score (0.0-5.0) for OCR output.

    Factors:
    - Characters per page (low = poor quality)
    - ASCII ratio (low = possible garbled text)
    - Block count from JSON (low = poor structure detection)
    """
    score = 5.0

    # Factor 1: Characters per page
    if page_count > 0:
        chars_per_page = len(text) / page_count
        if chars_per_page < 50:
            score -= 2.0  # Very low content
        elif chars_per_page < 200:
            score -= 1.0  # Low content
        elif chars_per_page < 500:
            score -= 0.5  # Below average

    # Factor 2: ASCII ratio (printable characters)
    if len(text) > 0:
        ascii_chars = sum(1 for c in text if c.isprintable() or c in "\n\r\t")
        ascii_ratio = ascii_chars / len(text)
        if ascii_ratio < 0.7:
            score -= 1.5  # Lots of garbled text
        elif ascii_ratio < 0.85:
            score -= 0.5  # Some garbled text

    # Factor 3: Block count from JSON
    if json_blocks is not None:
        children = json_blocks.get("children", [])
        if isinstance(children, list):
            block_count = len(children)
            if block_count == 0:
                score -= 1.0
            elif block_count < page_count and page_count > 1:
                score -= 0.5  # Fewer blocks than pages

    # Factor 4: Empty text penalty
    if len(text.strip()) == 0:
        score = 0.0

    return max(0.0, min(5.0, round(score, 2)))


def images_to_base64(images: dict) -> dict[str, str]:
    """
    Convert Marker PIL Image dict to base64 string dict.

    Marker returns: {block_id: PIL.Image.Image}
    We need: {filename: base64_string}
    """
    result = {}
    for block_id, pil_image in images.items():
        try:
            buf = io.BytesIO()
            img_format = "PNG"
            if hasattr(pil_image, "format") and pil_image.format:
                img_format = pil_image.format
            pil_image.save(buf, format=img_format)
            b64_data = base64.b64encode(buf.getvalue()).decode("ascii")
            buf.close()  # F-4: free BytesIO buffer immediately
            # Use block_id as filename with appropriate extension
            ext = img_format.lower()
            if ext == "jpeg":
                ext = "jpg"
            filename = f"{block_id}.{ext}"
            result[filename] = b64_data
        except Exception as e:
            logger.warning(f"Failed to encode image {block_id}: {e}")
        finally:
            pil_image.close()  # F-4: free PIL image memory
    return result


def extract_pdf_metadata(file_path: Path) -> dict:
    """
    Extract document metadata (title, author, subject) from PDF.
    Returns dict with keys: title, author, subject (all optional).
    """
    metadata = {}
    if file_path.suffix.lower() != ".pdf":
        return metadata

    try:
        import pypdf

        reader = pypdf.PdfReader(str(file_path))
        info = reader.metadata
        if info:
            if info.title:
                metadata["title"] = str(info.title)
            if info.author:
                metadata["author"] = str(info.author)
            if info.subject:
                metadata["subject"] = str(info.subject)
    except ImportError:
        logger.debug("pypdf not available, skipping PDF metadata extraction")
    except Exception as e:
        logger.warning(f"Failed to extract PDF metadata: {e}")

    return metadata


# =============================================================================
# MODE CONFIGURATION
# =============================================================================


def configure_mode(mode: str) -> None:
    """
    Set Surya environment variables for the OCR mode.
    MUST be called BEFORE importing Marker models (create_model_dict).

    torch.compile is only enabled for "accurate" mode because Surya models
    process variable-size images. torch.compile assumes static shapes by default,
    causing dimension mismatches on different input sizes.
    """
    if mode == "fast":
        os.environ["DETECTOR_BATCH_SIZE"] = "8"
        os.environ["RECOGNITION_BATCH_SIZE"] = "64"
        os.environ["TABLE_REC_BATCH_SIZE"] = "8"
        logger.info("Mode: fast (reduced batch sizes for speed)")
    elif mode == "accurate":
        os.environ["COMPILE_ALL"] = "true"
        logger.info("Mode: accurate (torch.compile enabled, slower warmup)")
    else:
        # balanced: use Surya defaults, no torch.compile
        logger.info("Mode: balanced (Surya default batch sizes)")


# =============================================================================
# MAIN PROCESSING
# =============================================================================


def _convert_text_direct(
    validated_path: Path,
    timeout: int,
    start_time: float,
    max_pages: Optional[int] = None,
) -> Tuple[str, int, Optional[dict], Optional[dict], Optional[Dict[str, str]]]:
    """
    Convert plain text files (.txt, .csv, .md) directly to paginated markdown.

    These formats don't need OCR — they are already machine-readable text.
    - .txt and .md: read content directly as UTF-8
    - .csv: parse with csv module, convert to markdown table format

    Returns:
        (markdown_text, page_count, json_blocks, metadata, images_b64=None)
    """
    ext = validated_path.suffix.lower()
    file_size = validated_path.stat().st_size
    logger.info(
        f"Text file detected ({ext}) — reading directly, no OCR needed "
        f"(size: {file_size} bytes)"
    )

    check_timeout(start_time, timeout, f"reading {ext} file")

    # Read file content
    try:
        raw_content = validated_path.read_text(encoding="utf-8")
    except UnicodeDecodeError:
        # Try with latin-1 as fallback for non-UTF8 text files
        try:
            raw_content = validated_path.read_text(encoding="latin-1")
            logger.warning(
                f"File {validated_path.name} is not UTF-8, decoded as latin-1"
            )
        except Exception as e:
            raise OCRFileError(
                f"Cannot read text file {validated_path.name}: {e}",
                str(validated_path),
            )
    except Exception as e:
        raise OCRFileError(
            f"Failed to read file {validated_path.name}: {e}",
            str(validated_path),
        )

    if not raw_content.strip():
        # Empty file — return minimal result
        markdown_text = f"\n\n{{0}}{PAGE_SEPARATOR}\n\n(empty file)"
        return markdown_text, 1, {"type": "text_direct", "lines": 0}, {
            "conversion_method": "text_direct",
            "file_type": ext,
            "file_size_bytes": file_size,
        }, None

    # Convert CSV to markdown table
    if ext == ".csv":
        check_timeout(start_time, timeout, "CSV parsing")
        content = _csv_to_markdown(raw_content, validated_path.name)
    else:
        # .txt and .md — content is already readable text
        content = raw_content

    # Paginate: split into pages by line count (~60 lines per page, matching typical docs)
    lines = content.split("\n")
    lines_per_page = 60
    all_pages: list[str] = []

    for chunk_start in range(0, len(lines), lines_per_page):
        if max_pages is not None and len(all_pages) >= max_pages:
            break
        check_timeout(start_time, timeout, f"text pagination page {len(all_pages)}")
        chunk_end = min(chunk_start + lines_per_page, len(lines))
        page_text = "\n".join(lines[chunk_start:chunk_end])
        all_pages.append(page_text)

    # If content was so short it produced no pages, make one page
    if not all_pages:
        all_pages.append(content)

    page_count = len(all_pages)

    # Build paginated output matching Marker format: {page_num} + separator
    parts = []
    for i, page_text in enumerate(all_pages):
        parts.append(f"\n\n{{{i}}}{PAGE_SEPARATOR}\n\n{page_text}")
    markdown_text = "".join(parts)

    total_lines = len(lines)
    json_blocks = {
        "type": "text_direct",
        "file_type": ext,
        "lines": total_lines,
        "pages": page_count,
    }

    metadata = {
        "conversion_method": "text_direct",
        "file_type": ext,
        "file_size_bytes": file_size,
        "total_lines": total_lines,
        "lines_per_page": lines_per_page,
        "marker_version": "1.10.2",
    }

    logger.info(
        f"Text conversion complete: {total_lines} lines, {page_count} pages, "
        f"{len(markdown_text)} chars"
    )

    return markdown_text, page_count, json_blocks, metadata, None


def _csv_to_markdown(raw_content: str, filename: str) -> str:
    """
    Convert CSV content to markdown table format.

    Uses csv.reader for proper handling of quoted fields, commas in values, etc.
    """
    reader = csv.reader(io.StringIO(raw_content))
    rows = list(reader)

    if not rows:
        return "(empty CSV)"

    # First row is header
    header = rows[0]
    col_names = [str(c).strip() if c else f"Col_{i+1}" for i, c in enumerate(header)]

    md_lines = [f"# {filename}", ""]
    md_lines.append("| " + " | ".join(col_names) + " |")
    md_lines.append("| " + " | ".join(["---"] * len(col_names)) + " |")

    for row in rows[1:]:
        cells = [
            str(c).replace("|", "\\|").replace("\n", " ") if c else ""
            for c in row
        ]
        # Pad or truncate to match header columns
        while len(cells) < len(col_names):
            cells.append("")
        cells = cells[:len(col_names)]
        md_lines.append("| " + " | ".join(cells) + " |")

    return "\n".join(md_lines)


def _convert_xls_direct(
    validated_path: Path,
    timeout: int,
    start_time: float,
    max_pages: Optional[int] = None,
) -> Tuple[str, int, Optional[dict], Optional[dict], Optional[Dict[str, str]]]:
    """
    Convert .xls (legacy Excel) files directly to markdown tables using xlrd.

    openpyxl only supports .xlsx. For .xls files, xlrd v2.0.2 is required.
    Paginates rows using SPREADSHEET_ROWS_PER_PAGE, same as _convert_spreadsheet_direct.

    Returns:
        (markdown_text, page_count, json_blocks, metadata, images_b64=None)
    """
    import xlrd

    logger.info(
        f"Legacy .xls file detected — using xlrd direct markdown conversion "
        f"(skipping Marker pipeline)"
    )

    check_timeout(start_time, timeout, "opening .xls file")

    try:
        wb = xlrd.open_workbook(str(validated_path))
    except Exception as e:
        raise OCRFileError(
            f"Cannot open .xls file {validated_path.name}: {e}",
            str(validated_path),
        )

    all_pages: list[str] = []
    total_rows = 0
    sheet_stats: list[dict] = []

    for sheet_idx in range(wb.nsheets):
        if max_pages is not None and len(all_pages) >= max_pages:
            break
        ws = wb.sheet_by_index(sheet_idx)
        sheet_name = ws.name
        check_timeout(start_time, timeout, f"xls sheet {sheet_name}")

        if ws.nrows == 0:
            continue

        sheet_row_count = ws.nrows
        total_rows += sheet_row_count

        # Read header row (row 0)
        header = [ws.cell_value(0, col) for col in range(ws.ncols)]
        col_names = [
            str(c).strip() if c else f"Col_{i+1}" for i, c in enumerate(header)
        ]

        sheet_stats.append({
            "sheet": sheet_name,
            "rows": sheet_row_count,
            "columns": len(col_names),
        })

        # Build markdown header
        md_header = "| " + " | ".join(col_names) + " |"
        md_sep = "| " + " | ".join(["---"] * len(col_names)) + " |"

        # Paginate data rows (skip header row 0)
        data_row_count = sheet_row_count - 1
        page_idx = 0
        for chunk_start in range(0, data_row_count, SPREADSHEET_ROWS_PER_PAGE):
            if max_pages is not None and len(all_pages) >= max_pages:
                break
            check_timeout(start_time, timeout, f"xls page {len(all_pages)}")

            chunk_end = min(chunk_start + SPREADSHEET_ROWS_PER_PAGE, data_row_count)

            page_lines = [f"# {sheet_name}"]
            if page_idx > 0 or wb.nsheets > 1:
                page_lines[0] += f" (rows {chunk_start + 2}-{chunk_end + 1})"
            page_lines.append("")
            page_lines.append(md_header)
            page_lines.append(md_sep)

            for row_idx in range(chunk_start + 1, chunk_end + 1):
                cells = []
                for col_idx in range(ws.ncols):
                    val = ws.cell_value(row_idx, col_idx)
                    cell_type = ws.cell_type(row_idx, col_idx)
                    # Format cell value appropriately
                    if cell_type == xlrd.XL_CELL_EMPTY:
                        cells.append("")
                    elif cell_type == xlrd.XL_CELL_NUMBER:
                        # Display integers without decimal point
                        if val == int(val):
                            cells.append(str(int(val)))
                        else:
                            cells.append(str(val))
                    elif cell_type == xlrd.XL_CELL_DATE:
                        try:
                            dt = xlrd.xldate_as_datetime(val, wb.datemode)
                            cells.append(dt.strftime("%Y-%m-%d %H:%M:%S"))
                        except Exception:
                            cells.append(str(val))
                    elif cell_type == xlrd.XL_CELL_BOOLEAN:
                        cells.append("TRUE" if val else "FALSE")
                    else:
                        cells.append(
                            str(val).replace("|", "\\|").replace("\n", " ")
                        )
                # Pad or truncate to match header columns
                while len(cells) < len(col_names):
                    cells.append("")
                cells = cells[:len(col_names)]
                page_lines.append("| " + " | ".join(cells) + " |")

            all_pages.append("\n".join(page_lines))
            page_idx += 1

    sheet_count = wb.nsheets
    wb.release_resources()

    page_count = len(all_pages) if all_pages else 1
    # Build paginated output matching Marker format
    parts = []
    for i, page_text in enumerate(all_pages):
        parts.append(f"\n\n{{{i}}}{PAGE_SEPARATOR}\n\n{page_text}")
    markdown_text = (
        "".join(parts)
        if parts
        else f"\n\n{{0}}{PAGE_SEPARATOR}\n\n(empty spreadsheet)"
    )

    json_blocks = {
        "type": "xls_direct",
        "sheets": sheet_stats,
        "total_rows": total_rows,
        "pages": page_count,
    }

    metadata = {
        "conversion_method": "xls_direct",
        "total_rows": total_rows,
        "sheets": sheet_count,
        "rows_per_page": SPREADSHEET_ROWS_PER_PAGE,
        "marker_version": "1.10.2",
    }

    logger.info(
        f"XLS conversion complete: {total_rows} rows across "
        f"{len(sheet_stats)} sheet(s), {page_count} pages"
    )

    return markdown_text, page_count, json_blocks, metadata, None


def _convert_spreadsheet_direct(
    validated_path: Path,
    timeout: int,
    start_time: float,
    max_pages: Optional[int] = None,
) -> Tuple[str, int, Optional[dict], Optional[dict], Optional[Dict[str, str]]]:
    """
    Convert large spreadsheets directly to markdown tables without weasyprint.

    For XLSX files with >SPREADSHEET_MAX_ROWS_FOR_MARKER rows, weasyprint
    rendering explodes memory (62+ GB for 11K rows). This function reads the
    spreadsheet with openpyxl and emits paginated markdown tables directly.
    Note: .xls files are routed to _convert_xls_direct() before reaching here.

    Returns:
        (markdown_text, page_count, json_blocks, metadata, images_b64=None)
    """
    from openpyxl import load_workbook

    logger.info(
        f"Large spreadsheet detected — using direct markdown conversion "
        f"(skipping weasyprint/Marker PDF pipeline)"
    )

    wb = load_workbook(str(validated_path), read_only=True, data_only=True)
    all_pages: list[str] = []
    total_rows = 0
    sheet_stats: list[dict] = []

    for sheet_name in wb.sheetnames:
        if max_pages is not None and len(all_pages) >= max_pages:
            break
        check_timeout(start_time, timeout, f"spreadsheet sheet {sheet_name}")
        ws = wb[sheet_name]

        # Read header row
        rows = list(ws.iter_rows(values_only=True))
        if not rows:
            continue

        sheet_row_count = len(rows)
        total_rows += sheet_row_count
        header = rows[0]
        col_names = [str(c) if c is not None else f"Col_{i+1}" for i, c in enumerate(header)]

        sheet_stats.append({
            "sheet": sheet_name,
            "rows": sheet_row_count,
            "columns": len(col_names),
        })

        # Build markdown header
        md_header = "| " + " | ".join(col_names) + " |"
        md_sep = "| " + " | ".join(["---"] * len(col_names)) + " |"

        # Paginate data rows (skip header row)
        data_rows = rows[1:]
        page_idx = 0
        for chunk_start in range(0, len(data_rows), SPREADSHEET_ROWS_PER_PAGE):
            if max_pages is not None and len(all_pages) >= max_pages:
                break
            check_timeout(start_time, timeout, f"spreadsheet page {len(all_pages)}")

            chunk_end = min(chunk_start + SPREADSHEET_ROWS_PER_PAGE, len(data_rows))
            chunk = data_rows[chunk_start:chunk_end]

            page_lines = [f"# {sheet_name}"]
            if page_idx > 0 or len(wb.sheetnames) > 1:
                page_lines[0] += f" (rows {chunk_start + 2}-{chunk_end + 1})"
            page_lines.append("")
            page_lines.append(md_header)
            page_lines.append(md_sep)

            for row in chunk:
                cells = [
                    str(c).replace("|", "\\|").replace("\n", " ") if c is not None else ""
                    for c in row
                ]
                # Pad or truncate to match header columns
                while len(cells) < len(col_names):
                    cells.append("")
                cells = cells[:len(col_names)]
                page_lines.append("| " + " | ".join(cells) + " |")

            all_pages.append("\n".join(page_lines))
            page_idx += 1

    sheet_count = len(wb.sheetnames)
    wb.close()

    page_count = len(all_pages) if all_pages else 1
    # Build paginated output matching Marker format: {page_num} + separator
    parts = []
    for i, page_text in enumerate(all_pages):
        parts.append(f"\n\n{{{i}}}{PAGE_SEPARATOR}\n\n{page_text}")
    markdown_text = (
        "".join(parts)
        if parts
        else f"\n\n{{0}}{PAGE_SEPARATOR}\n\n(empty spreadsheet)"
    )

    # Build json_blocks equivalent
    json_blocks = {
        "type": "spreadsheet_direct",
        "sheets": sheet_stats,
        "total_rows": total_rows,
        "pages": page_count,
    }

    metadata = {
        "conversion_method": "direct_spreadsheet",
        "total_rows": total_rows,
        "sheets": sheet_count,
        "rows_per_page": SPREADSHEET_ROWS_PER_PAGE,
        "marker_version": "1.10.2",
    }

    logger.info(
        f"Spreadsheet conversion complete: {total_rows} rows across "
        f"{len(sheet_stats)} sheet(s), {page_count} pages"
    )

    return markdown_text, page_count, json_blocks, metadata, None


def _is_large_spreadsheet(validated_path: Path) -> bool:
    """Check if an XLSX file exceeds the row threshold for Marker processing.

    Note: .xls files are always routed to _convert_xls_direct() before this check,
    so this function only needs to handle .xlsx files via openpyxl.
    """
    ext = validated_path.suffix.lower()
    if ext != ".xlsx":
        return False
    try:
        from openpyxl import load_workbook
        wb = load_workbook(str(validated_path), read_only=True, data_only=True)
        total_rows = 0
        for sheet_name in wb.sheetnames:
            ws = wb[sheet_name]
            total_rows += ws.max_row or 0
        wb.close()
        if total_rows > SPREADSHEET_MAX_ROWS_FOR_MARKER:
            logger.info(
                f"Spreadsheet has {total_rows} rows (threshold: "
                f"{SPREADSHEET_MAX_ROWS_FOR_MARKER}). Using direct conversion."
            )
            return True
        return False
    except Exception as e:
        logger.warning(f"Could not check spreadsheet size: {e}. Using Marker pipeline.")
        return False


def _run_marker_conversion(
    model_dict: Any,
    validated_path: Path,
    mode: Literal["fast", "balanced", "accurate"],
    timeout: int,
    start_time: float,
    max_pages: Optional[int] = None,
    page_range: Optional[str] = None,
    disable_image_extraction: bool = False,
) -> Tuple[str, int, Optional[dict], Optional[dict], Optional[Dict[str, str]]]:
    """
    Core Marker conversion logic: two-pass (markdown + JSON) using a pre-loaded model_dict.

    Returns:
        (markdown_text, page_count, json_blocks, marker_metadata, images_b64)

    Raises:
        OCRTimeoutError, OCRModelError, OCRGPUError on failure.
    """
    import torch

    from marker.config.parser import ConfigParser
    from marker.converters.pdf import PdfConverter
    from marker.output import text_from_rendered

    # Pre-compute effective page_range from max_pages if needed.
    # Only restrict when max_pages < actual document page count to avoid
    # requesting non-existent pages (Marker validates page_range).
    # Non-PDF files skip this since page count isn't known before Marker converts them.
    max_pages_range: Optional[List[int]] = None
    if max_pages is not None and page_range is None:
        if validated_path.suffix.lower() == ".pdf":
            import pypdfium2

            pdf_doc = pypdfium2.PdfDocument(str(validated_path))
            doc_page_count = len(pdf_doc)
            pdf_doc.close()
            if max_pages < doc_page_count:
                max_pages_range = list(range(max_pages))
                logger.info(
                    f"max_pages={max_pages} < doc pages={doc_page_count}, "
                    f"restricting to first {max_pages} pages"
                )
            else:
                logger.info(
                    f"max_pages={max_pages} >= doc pages={doc_page_count}, "
                    f"processing all pages"
                )

    # Build config for Marker
    cli_options: dict = {
        "output_format": "markdown",
        "paginate_output": True,
    }

    if page_range is not None:
        cli_options["page_range"] = page_range
    if disable_image_extraction:
        cli_options["disable_image_extraction"] = True

    config_parser = ConfigParser(cli_options)
    config_dict = config_parser.generate_config_dict()

    if max_pages_range is not None:
        config_dict["page_range"] = max_pages_range

    # --- Pass 1: Markdown (text + images) ---
    logger.info("Running Marker conversion (markdown pass)...")
    converter_md = PdfConverter(
        artifact_dict=model_dict,
        renderer=config_parser.get_renderer(),
        config=config_dict,
    )

    try:
        rendered_md = converter_md(str(validated_path))
    except (torch.cuda.OutOfMemoryError, MemoryError) as e:
        clear_gpu_memory()
        raise OCRGPUError(
            f"GPU out of memory processing {validated_path.name}: {e}. "
            f"Try mode=fast or reduce page_range.",
            {
                "file_size_bytes": validated_path.stat().st_size,
                "original_error": str(e),
            },
        )
    except Exception as e:
        error_msg = str(e)
        if "cuda" in error_msg.lower() or "out of memory" in error_msg.lower():
            clear_gpu_memory()
            raise OCRGPUError(
                f"GPU error during conversion: {e}",
                {"original_error": error_msg},
            )
        raise OCRModelError(
            f"Marker conversion failed: {e}",
            {"file": str(validated_path), "original_error": error_msg},
        )

    markdown_text, _, md_images = text_from_rendered(rendered_md)
    page_count = converter_md.page_count or 1

    # Get metadata from the rendered output
    marker_metadata = None
    if hasattr(rendered_md, "metadata"):
        marker_metadata = rendered_md.metadata
        if not isinstance(marker_metadata, dict):
            marker_metadata = (
                marker_metadata.model_dump()
                if hasattr(marker_metadata, "model_dump")
                else dict(marker_metadata)
                if marker_metadata
                else None
            )

    logger.info(
        f"Markdown pass complete: {page_count} pages, "
        f"{len(markdown_text)} chars, {len(md_images)} images"
    )

    # F-1: Free Pass 1 intermediate objects before Pass 2 to reclaim RAM
    del converter_md, rendered_md
    gc.collect()

    # Check timeout before JSON pass
    check_timeout(start_time, timeout, "markdown conversion")

    # --- Pass 2: JSON blocks ---
    json_blocks = None
    json_cli_options: dict = {
        "output_format": "json",
    }
    if page_range is not None:
        json_cli_options["page_range"] = page_range
    if disable_image_extraction:
        json_cli_options["disable_image_extraction"] = True

    json_config_parser = ConfigParser(json_cli_options)
    json_config_dict = json_config_parser.generate_config_dict()

    if max_pages_range is not None:
        json_config_dict["page_range"] = max_pages_range

    logger.info("Running Marker conversion (JSON pass, reusing cached models)...")
    try:
        converter_json = PdfConverter(
            artifact_dict=model_dict,
            renderer=json_config_parser.get_renderer(),
            config=json_config_dict,
        )
        rendered_json = converter_json(str(validated_path))
        json_text, _, _ = text_from_rendered(rendered_json)

        if json_text:
            parsed_json = json.loads(json_text)
            if isinstance(parsed_json, dict):
                json_blocks = parsed_json
            elif isinstance(parsed_json, list):
                json_blocks = {"children": parsed_json}
            else:
                json_blocks = {"raw": parsed_json}
            logger.info(
                f"JSON pass complete: "
                f"{len(json_blocks.get('children', []))} top-level blocks"
            )
    except (torch.cuda.OutOfMemoryError, MemoryError) as e:
        clear_gpu_memory()
        logger.warning(
            "OOM during JSON pass (markdown pass succeeded): %s. "
            "Continuing without json_blocks.", e
        )
    except Exception as e:
        error_msg = str(e)
        if "out of memory" in error_msg.lower():
            clear_gpu_memory()
            logger.warning("OOM during JSON pass (markdown pass succeeded): %s. "
                           "Continuing without json_blocks.", e)
        else:
            logger.warning(f"JSON pass failed (non-fatal): {e}")

    # Process images
    images_b64 = None
    if not disable_image_extraction and md_images:
        images_b64 = images_to_base64(md_images)
        if images_b64:
            logger.info(f"Encoded {len(images_b64)} images to base64")

    return markdown_text, page_count, json_blocks, marker_metadata, images_b64


def _build_ocr_result(
    markdown_text: str,
    page_count: int,
    json_blocks: Optional[dict],
    marker_metadata: Optional[dict],
    images_b64: Optional[Dict[str, str]],
    validated_path: Path,
    document_id: str,
    provenance_id: str,
    mode: Literal["fast", "balanced", "accurate"],
    request_id: str,
    start_timestamp: str,
    start_time: float,
    model_load_ms: int,
) -> OCRResult:
    """Build an OCRResult from conversion outputs. Shared by CLI and daemon modes."""
    end_time = time.time()
    end_timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    duration_ms = int((end_time - start_time) * 1000)

    # Parse page offsets
    page_offsets = parse_marker_page_offsets(markdown_text)

    # Compute content hash
    content_hash = compute_content_hash(markdown_text)

    # Compute quality score
    quality_score = compute_quality_score(markdown_text, page_count, json_blocks)

    # Extract PDF metadata
    pdf_meta = extract_pdf_metadata(validated_path)

    # Merge metadata
    merged_metadata: dict = {}
    if marker_metadata:
        merged_metadata.update(marker_metadata)
    merged_metadata["marker_version"] = "1.10.2"
    merged_metadata["processing_mode"] = mode
    merged_metadata["model_load_ms"] = model_load_ms
    merged_metadata["local_gpu"] = True

    ocr_result = OCRResult(
        id=str(uuid.uuid4()),
        provenance_id=provenance_id,
        document_id=document_id,
        extracted_text=markdown_text,
        text_length=len(markdown_text),
        datalab_request_id=request_id,
        datalab_mode=mode,
        parse_quality_score=quality_score,
        page_count=page_count,
        cost_cents=0,
        content_hash=content_hash,
        processing_started_at=start_timestamp,
        processing_completed_at=end_timestamp,
        processing_duration_ms=duration_ms,
        page_offsets=page_offsets,
        images=images_b64 if images_b64 else None,
        json_blocks=json_blocks,
        metadata=merged_metadata if merged_metadata else None,
        extraction_json=None,
        cost_breakdown_full=None,
        extras_features=None,
        doc_title=pdf_meta.get("title"),
        doc_author=pdf_meta.get("author"),
        doc_subject=pdf_meta.get("subject"),
    )

    logger.info(
        f"OCR complete: {page_count} pages, {len(markdown_text)} chars, "
        f"{duration_ms}ms, quality={quality_score}, cost=$0.00"
    )

    return ocr_result


def process_document_with_models(
    model_dict: Any,
    request: dict,
    model_load_ms: int = 0,
) -> OCRResult:
    """
    Process a document using a pre-loaded model_dict. Used by daemon mode.

    The model_dict is created ONCE at daemon startup and reused across all
    work items. This avoids the ~10-30s model loading overhead per document.

    Args:
        model_dict: Pre-loaded Marker model dictionary from create_model_dict()
        request: Work item dict with keys: file_path, document_id, provenance_id,
                 mode, timeout, max_pages, page_range, disable_image_extraction
        model_load_ms: Time taken to load models (for metadata, recorded once at startup)

    Returns:
        OCRResult with extracted text and metadata

    Raises:
        OCRFileError, OCRTimeoutError, OCRModelError, OCRGPUError
    """
    import torch

    file_path = request.get("file_path", "")
    document_id = request.get("document_id", "")
    provenance_id = request.get("provenance_id", "")
    mode = request.get("mode", "balanced")
    timeout = request.get("timeout", 600)
    max_pages = request.get("max_pages")
    page_range = request.get("page_range")
    disable_image_extraction = request.get("disable_image_extraction", False)

    # Validate inputs
    if not file_path:
        raise OCRFileError("file_path is required", "")
    if not document_id:
        raise OCRModelError("document_id is required")
    if not provenance_id:
        raise OCRModelError("provenance_id is required")
    if mode not in ("fast", "balanced", "accurate"):
        raise OCRModelError(f"Invalid mode: {mode}. Must be fast/balanced/accurate")

    validated_path = validate_file(file_path)
    # Set document ID for progress reporting (tqdm interception)
    _progress_reporter.set_document_id(document_id)
    logger.info(
        f"Processing document (daemon): {validated_path} "
        f"(mode={mode}, timeout={timeout}s)"
    )

    start_time = time.time()
    start_timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    request_id = f"local-marker-{uuid.uuid4()}"

    try:
        # Route text files directly (no OCR needed)
        ext = validated_path.suffix.lower()
        if ext in (".txt", ".csv", ".md"):
            markdown_text, page_count, json_blocks, marker_metadata, images_b64 = (
                _convert_text_direct(
                    validated_path=validated_path,
                    timeout=timeout,
                    start_time=start_time,
                    max_pages=max_pages,
                )
            )
        # Route legacy .xls files via xlrd (openpyxl and Marker don't support .xls)
        elif ext == ".xls":
            markdown_text, page_count, json_blocks, marker_metadata, images_b64 = (
                _convert_xls_direct(
                    validated_path=validated_path,
                    timeout=timeout,
                    start_time=start_time,
                    max_pages=max_pages,
                )
            )
        # Route large .xlsx spreadsheets to direct markdown conversion (skip weasyprint)
        elif _is_large_spreadsheet(validated_path):
            markdown_text, page_count, json_blocks, marker_metadata, images_b64 = (
                _convert_spreadsheet_direct(
                    validated_path=validated_path,
                    timeout=timeout,
                    start_time=start_time,
                    max_pages=max_pages,
                )
            )
        else:
            markdown_text, page_count, json_blocks, marker_metadata, images_b64 = (
                _run_marker_conversion(
                    model_dict=model_dict,
                    validated_path=validated_path,
                    mode=mode,
                    timeout=timeout,
                    start_time=start_time,
                    max_pages=max_pages,
                    page_range=page_range,
                    disable_image_extraction=disable_image_extraction,
                )
            )

        # Clear progress reporter after OCR completes
        _progress_reporter.set_document_id("")

        return _build_ocr_result(
            markdown_text=markdown_text,
            page_count=page_count,
            json_blocks=json_blocks,
            marker_metadata=marker_metadata,
            images_b64=images_b64,
            validated_path=validated_path,
            document_id=document_id,
            provenance_id=provenance_id,
            mode=mode,
            request_id=request_id,
            start_timestamp=start_timestamp,
            start_time=start_time,
            model_load_ms=model_load_ms,
        )

    except (OCRFileError, OCRTimeoutError, OCRModelError, OCRGPUError):
        raise
    except (torch.cuda.OutOfMemoryError, MemoryError) as e:
        clear_gpu_memory()
        raise OCRGPUError(
            f"Out of memory: {e}",
            {"original_error": str(e)},
        )
    except Exception as e:
        error_msg = str(e)
        if "cuda" in error_msg.lower() or "gpu" in error_msg.lower() or "out of memory" in error_msg.lower():
            raise OCRGPUError(
                f"Unexpected GPU error: {e}",
                {"original_error": error_msg},
            )
        raise OCRModelError(
            f"Unexpected error during OCR processing: {e}",
            {"original_error": error_msg},
        )


def process_document(
    file_path: str,
    document_id: str,
    provenance_id: str,
    mode: Literal["fast", "balanced", "accurate"] = "balanced",
    timeout: int = 600,
    max_pages: Optional[int] = None,
    page_range: Optional[str] = None,
    disable_image_extraction: bool = False,
) -> OCRResult:
    """
    Process a document through Marker-pdf local GPU OCR.

    This is the original CLI-mode function. It loads models fresh each time.
    For daemon mode, use process_document_with_models() instead.

    Args:
        file_path: Path to document (PDF, image, DOCX, etc.)
        document_id: UUID of the document record in database
        provenance_id: UUID for the OCR_RESULT provenance record
        mode: OCR quality mode (fast/balanced/accurate)
        timeout: Maximum processing time in seconds (default 600s = 10 min)
        max_pages: Maximum pages to process
        page_range: Specific pages, 0-indexed (e.g. "0-5,10")
        disable_image_extraction: Skip image extraction

    Returns:
        OCRResult with extracted text and metadata

    Raises:
        OCRFileError: File not found or unsupported type
        OCRTimeoutError: Processing exceeded timeout
        OCRModelError: Model loading or inference failure
        OCRGPUError: GPU memory or CUDA error
    """
    import torch

    # Validate file first
    validated_path = validate_file(file_path)
    logger.info(
        f"Processing document: {validated_path} (mode={mode}, timeout={timeout}s)"
    )

    # Record timing
    start_time = time.time()
    start_timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())

    # Generate local request ID
    request_id = f"local-marker-{uuid.uuid4()}"

    try:
        # Route text files and .xls directly — no Marker models needed
        ext = validated_path.suffix.lower()
        if ext in (".txt", ".csv", ".md"):
            markdown_text, page_count, json_blocks, marker_metadata, images_b64 = (
                _convert_text_direct(
                    validated_path=validated_path,
                    timeout=timeout,
                    start_time=start_time,
                    max_pages=max_pages,
                )
            )
            return _build_ocr_result(
                markdown_text=markdown_text,
                page_count=page_count,
                json_blocks=json_blocks,
                marker_metadata=marker_metadata,
                images_b64=images_b64,
                validated_path=validated_path,
                document_id=document_id,
                provenance_id=provenance_id,
                mode=mode,
                request_id=request_id,
                start_timestamp=start_timestamp,
                start_time=start_time,
                model_load_ms=0,
            )

        if ext == ".xls":
            markdown_text, page_count, json_blocks, marker_metadata, images_b64 = (
                _convert_xls_direct(
                    validated_path=validated_path,
                    timeout=timeout,
                    start_time=start_time,
                    max_pages=max_pages,
                )
            )
            return _build_ocr_result(
                markdown_text=markdown_text,
                page_count=page_count,
                json_blocks=json_blocks,
                marker_metadata=marker_metadata,
                images_b64=images_b64,
                validated_path=validated_path,
                document_id=document_id,
                provenance_id=provenance_id,
                mode=mode,
                request_id=request_id,
                start_timestamp=start_timestamp,
                start_time=start_time,
                model_load_ms=0,
            )

        # Configure mode env vars BEFORE importing Marker models
        configure_mode(mode)

        # Import Marker components (heavy imports, deferred intentionally)
        logger.info("Loading Marker models...")
        model_load_start = time.time()

        try:
            from marker.models import create_model_dict
        except ImportError as e:
            raise OCRModelError(
                f"Marker-pdf not installed or import failed: {e}",
                {"missing_package": str(e)},
            )

        # Create model dict (downloads models on first call)
        try:
            model_dict = create_model_dict()
        except (torch.cuda.OutOfMemoryError, MemoryError) as e:
            clear_gpu_memory()
            raise OCRGPUError(
                f"Out of memory loading models: {e}",
                {"original_error": str(e)},
            )
        except Exception as e:
            error_msg = str(e)
            if "cuda" in error_msg.lower() or "gpu" in error_msg.lower() or "out of memory" in error_msg.lower():
                raise OCRGPUError(
                    f"GPU error loading models: {e}",
                    {"original_error": error_msg},
                )
            raise OCRModelError(
                f"Failed to load Marker models: {e}",
                {"original_error": error_msg},
            )

        model_load_ms = int((time.time() - model_load_start) * 1000)
        logger.info(f"Models loaded in {model_load_ms}ms")

        # Check timeout after model loading
        check_timeout(start_time, timeout, "model loading")

        # Use shared conversion logic
        markdown_text, page_count, json_blocks, marker_metadata, images_b64 = (
            _run_marker_conversion(
                model_dict=model_dict,
                validated_path=validated_path,
                mode=mode,
                timeout=timeout,
                start_time=start_time,
                max_pages=max_pages,
                page_range=page_range,
                disable_image_extraction=disable_image_extraction,
            )
        )

        return _build_ocr_result(
            markdown_text=markdown_text,
            page_count=page_count,
            json_blocks=json_blocks,
            marker_metadata=marker_metadata,
            images_b64=images_b64,
            validated_path=validated_path,
            document_id=document_id,
            provenance_id=provenance_id,
            mode=mode,
            request_id=request_id,
            start_timestamp=start_timestamp,
            start_time=start_time,
            model_load_ms=model_load_ms,
        )

    except (OCRFileError, OCRTimeoutError, OCRModelError, OCRGPUError):
        # Re-raise our own errors
        raise
    except (torch.cuda.OutOfMemoryError, MemoryError) as e:
        clear_gpu_memory()
        raise OCRGPUError(
            f"Out of memory: {e}",
            {"original_error": str(e)},
        )
    except Exception as e:
        error_msg = str(e)
        # Categorize unexpected errors
        if "cuda" in error_msg.lower() or "gpu" in error_msg.lower() or "out of memory" in error_msg.lower():
            raise OCRGPUError(
                f"Unexpected GPU error: {e}",
                {"original_error": error_msg},
            )
        raise OCRModelError(
            f"Unexpected error during OCR processing: {e}",
            {"original_error": error_msg},
        )
    finally:
        # Always free GPU memory
        clear_gpu_memory()


# =============================================================================
# DAEMON MODE
# =============================================================================


def _load_models_for_daemon(
    default_mode: str,
) -> tuple[Any, int]:
    """
    Load Marker models once for daemon mode.

    Args:
        default_mode: Default OCR mode to configure env vars for.

    Returns:
        (model_dict, model_load_ms)

    Raises:
        OCRModelError, OCRGPUError on failure.
    """
    import torch

    configure_mode(default_mode)

    logger.info("Daemon: Loading Marker models (one-time)...")
    model_load_start = time.time()

    try:
        from marker.models import create_model_dict
    except ImportError as e:
        raise OCRModelError(
            f"Marker-pdf not installed or import failed: {e}",
            {"missing_package": str(e)},
        )

    try:
        model_dict = create_model_dict()
    except (torch.cuda.OutOfMemoryError, MemoryError) as e:
        clear_gpu_memory()
        raise OCRGPUError(
            f"Out of memory loading models: {e}",
            {"original_error": str(e)},
        )
    except Exception as e:
        error_msg = str(e)
        if "cuda" in error_msg.lower() or "gpu" in error_msg.lower() or "out of memory" in error_msg.lower():
            raise OCRGPUError(
                f"GPU error loading models: {e}",
                {"original_error": error_msg},
            )
        raise OCRModelError(
            f"Failed to load Marker models: {e}",
            {"original_error": error_msg},
        )

    model_load_ms = int((time.time() - model_load_start) * 1000)
    logger.info(f"Daemon: Models loaded in {model_load_ms}ms")

    return model_dict, model_load_ms


def daemon_loop(model_dict: Any, model_load_ms: int) -> None:
    """
    Persistent processing loop - models already loaded in GPU memory.

    Protocol (JSON lines over stdin/stdout):
      Input:  {"action": "process", "file_path": "...", "document_id": "...", ...}
      Output: OCRResult as JSON (same format as CLI mode)

      Input:  {"action": "ping"}
      Output: {"status": "ready"}

      Input:  {"action": "shutdown"}
      Output: (none, clean exit)

    On error: {"error": "...", "category": "...", "details": {...}}

    Reads one JSON line per iteration, processes, writes one JSON line response.
    Exits cleanly on stdin EOF (TypeScript process closed pipe) or shutdown action.
    """
    # Track items processed for diagnostics
    items_processed = 0

    # Output ready signal so the TypeScript caller knows models are loaded
    _output_json({
        "status": "ready",
        "models_loaded": True,
        "load_time_ms": model_load_ms,
    })

    logger.info("Daemon: Entering main loop, waiting for work items on stdin...")

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue

        try:
            request = json.loads(line)
        except json.JSONDecodeError as e:
            _output_json({
                "error": f"Invalid JSON: {e}",
                "category": "OCR_USAGE_ERROR",
                "details": {"raw_input": line[:200]},
            })
            continue

        action = request.get("action", "process")

        if action == "ping":
            _output_json({
                "status": "ready",
                "items_processed": items_processed,
            })
            continue

        if action == "shutdown":
            logger.info(
                f"Daemon: Shutdown requested after {items_processed} items processed."
            )
            break

        if action != "process":
            _output_json({
                "error": f"Unknown action: {action}",
                "category": "OCR_USAGE_ERROR",
                "details": {"valid_actions": ["process", "ping", "shutdown"]},
            })
            continue

        # Process the document using pre-loaded models
        try:
            result = process_document_with_models(
                model_dict=model_dict,
                request=request,
                model_load_ms=model_load_ms,
            )
            result_dict = asdict(result)
            _output_json(result_dict)
            del result, result_dict  # F-11: free immediately after serialization
            items_processed += 1
        except LocalOCRError as e:
            logger.error(f"Daemon OCR error [{e.category}]: {e}")
            _output_json({
                "error": str(e),
                "category": e.category,
                "details": e.details,
            })
        except Exception as e:
            logger.exception(f"Daemon unexpected error: {e}")
            _output_json({
                "error": str(e),
                "category": "OCR_INTERNAL_ERROR",
                "details": {"type": type(e).__name__},
            })
        finally:
            # F-3: Free GPU cache between documents to prevent VRAM fragmentation
            clear_gpu_memory()

    # Clean exit: stdin EOF or shutdown action
    logger.info(
        f"Daemon: Exiting after {items_processed} items processed. Freeing GPU memory."
    )
    clear_gpu_memory()


# =============================================================================
# CLI INTERFACE
# =============================================================================


def main() -> None:
    """CLI entry point - outputs JSON to stdout, logs to stderr."""
    parser = argparse.ArgumentParser(
        description="Local GPU OCR Worker (Marker-pdf) - Extract text from documents",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Single-document CLI mode:
  python ocr_worker_local.py --file_path ./doc.pdf --document_id abc --provenance_id def --mode balanced
  python ocr_worker_local.py --file_path ./doc.pdf --document_id abc --provenance_id def --mode fast --max_pages 10

  # Persistent daemon mode (models loaded once, work items via stdin):
  python ocr_worker_local.py --daemon --mode balanced
        """,
    )

    # Daemon mode flag
    parser.add_argument(
        "--daemon",
        action="store_true",
        help="Run in daemon mode: load models once, accept JSON work items on stdin",
    )

    # CLI mode args (required only when NOT in daemon mode)
    parser.add_argument(
        "--file_path", type=str, required=False, help="Path to document file"
    )
    parser.add_argument(
        "--document_id", type=str, required=False, help="Document ID (UUID)"
    )
    parser.add_argument(
        "--provenance_id", type=str, required=False, help="Provenance ID (UUID)"
    )
    parser.add_argument(
        "--mode",
        choices=["fast", "balanced", "accurate"],
        default="balanced",
        help="OCR mode (default: balanced). In daemon mode, sets the default mode.",
    )
    parser.add_argument(
        "--max_pages", type=int, default=None, help="Maximum pages to process"
    )
    parser.add_argument(
        "--page_range",
        type=str,
        default=None,
        help='Page range, 0-indexed (e.g. "0-5,10")',
    )
    parser.add_argument(
        "--disable_image_extraction",
        action="store_true",
        help="Skip image extraction",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=600,
        help="Timeout in seconds (default: 600 = 10 min)",
    )

    args = parser.parse_args()

    # Verify torch is available early
    try:
        import torch  # noqa: F401
    except ImportError:
        output_error(
            "OCR_MODEL_ERROR",
            "PyTorch not installed. Required for Marker-pdf local OCR.",
            {"missing_package": "torch"},
        )

    # ---- DAEMON MODE ----
    if args.daemon:
        # Set up signal handlers for clean shutdown
        def _signal_handler(signum: int, _frame: Any) -> None:
            sig_name = signal.Signals(signum).name
            logger.info(f"Daemon: Received {sig_name}, shutting down...")
            clear_gpu_memory()
            sys.exit(0)

        signal.signal(signal.SIGTERM, _signal_handler)
        signal.signal(signal.SIGINT, _signal_handler)

        try:
            model_dict, model_load_ms = _load_models_for_daemon(args.mode)
            daemon_loop(model_dict, model_load_ms)
        except LocalOCRError as e:
            logger.error(f"Daemon startup error [{e.category}]: {e}")
            _output_json({
                "error": str(e),
                "category": e.category,
                "details": e.details,
            })
            sys.exit(1)
        except Exception as e:
            logger.exception(f"Daemon fatal error: {e}")
            _output_json({
                "error": str(e),
                "category": "OCR_INTERNAL_ERROR",
                "details": {"type": type(e).__name__},
            })
            sys.exit(1)
        return

    # ---- CLI MODE (original behavior) ----
    # Validate required CLI args
    if not args.file_path:
        parser.error("--file_path is required in CLI mode (or use --daemon)")
    if not args.document_id:
        parser.error("--document_id is required in CLI mode (or use --daemon)")
    if not args.provenance_id:
        parser.error("--provenance_id is required in CLI mode (or use --daemon)")

    try:
        result = process_document(
            file_path=args.file_path,
            document_id=args.document_id,
            provenance_id=args.provenance_id,
            mode=args.mode,
            timeout=args.timeout,
            max_pages=args.max_pages,
            page_range=args.page_range,
            disable_image_extraction=args.disable_image_extraction,
        )

        # Output JSON to stdout (compact format for python-shell compatibility)
        json.dump(asdict(result), sys.stdout)
        sys.stdout.write("\n")
        sys.stdout.flush()

    except LocalOCRError as e:
        logger.error(f"OCR error [{e.category}]: {e}")
        output_error(e.category, str(e), e.details)

    except Exception as e:
        logger.exception(f"Fatal unexpected error: {e}")
        output_error(
            "OCR_MODEL_ERROR",
            f"Unexpected fatal error: {e}",
            {"type": type(e).__name__},
        )


if __name__ == "__main__":
    main()
