#!/usr/bin/env python3
"""Convert every PDF's pages (from metadata.json) into 3 web-efficient WebP variants:

1. thumb:
   Small thumbnail (~320px max width, WebP quality 70) for fast index/grid browsing.
2. full:
   High-resolution page render (200 DPI, WebP quality 82) preserving original color/grayscale
   and illustrations.
3. bw (clean grayscale / color-aware reading & OCR copy):
   250 DPI page processed to remove scan artefacts:
   - Detects whether the page contains genuine color pictures/diagrams (after white-balancing
     to ignore uniform paper yellowing/sepia tint).
   - If the page has no color pictures/diagrams, outputs a clean anti-aliased grayscale
     image: uneven lighting and gutter shadows are flattened via Gaussian background
     division, paper bleed-through and isolated dust specks (< 4 px) are suppressed to pure
     white (255), dark scanner border bars are stripped, and smooth stroke gradients are
     preserved for comfortable reading and accurate OCR.
   - If the page contains color pictures or diagrams, preserves color (RGB) while still
     white-balancing and cleaning the background illumination.

Outputs for each book are written under its `"directory"` from metadata.json:
  <directory>/thumb/0001.webp
  <directory>/full/0001.webp
  <directory>/bw/0001.webp

Usage:
    python3 convert_pages.py [options]

Options:
    --metadata PATH     Path to metadata.json (default: ./metadata.json)
    --only DIR          Only process the book matching this directory slug (or substring)
    --max-pages N       Only process the first N pages per PDF (useful for testing)
    --workers N         Number of parallel worker processes (default: CPU count)
    --force             Overwrite existing page images
"""
from __future__ import annotations

import argparse
import concurrent.futures
import json
import os
import sys
from pathlib import Path

import fitz  # PyMuPDF
import numpy as np
from PIL import Image
from scipy import ndimage


THUMB_WIDTH = 320
FULL_DPI = 200
BW_DPI = 250


def has_color_content(
    rgb: np.ndarray,
    chroma_thresh: float = 38.0,
    min_color_fraction: float = 0.005,
) -> bool:
    """Return True if an RGB page has genuine color pictures/diagrams (ignoring paper tint)."""
    # Downsample by 4x for fast chroma analysis
    small = rgb[::4, ::4, :].astype(np.float32)

    # White-balance using the 85th percentile per channel (paper background color)
    bg_rgb = np.percentile(small.reshape(-1, 3), 85, axis=0)
    bg_rgb = np.maximum(bg_rgb, 32.0)
    wb = np.clip(small * (255.0 / bg_rgb), 0.0, 255.0)

    c_max = wb.max(axis=2)
    c_min = wb.min(axis=2)
    lum = wb.mean(axis=2)
    chroma = c_max - c_min

    # Ignore very dark ink and near-white paper; smooth locally to ignore edge fringing
    valid = (lum > 30.0) & (lum < 235.0)
    smoothed_chroma = ndimage.uniform_filter(chroma * valid, size=5, mode="nearest")

    color_pixels = np.count_nonzero(smoothed_chroma > chroma_thresh)
    return (color_pixels / float(smoothed_chroma.size)) >= min_color_fraction


def clean_page_for_reading_and_ocr(
    rgb: np.ndarray,
    sauvola_window: int = 41,
    sauvola_k: float = 0.20,
    min_speck_area: int = 4,
    border_px: int = 12,
) -> tuple[Image.Image, bool]:
    """Clean scan artefacts; return (grayscale Image unless color pictures exist, is_color)."""
    is_color = has_color_content(rgb)

    if is_color:
        # Clean color page: normalize per-channel background illumination & strip border bars
        f_rgb = rgb.astype(np.float32)
        bg = np.stack(
            [ndimage.gaussian_filter(f_rgb[:, :, c], sigma=45.0, mode="nearest") for c in range(3)],
            axis=2,
        )
        bg = np.maximum(bg, 32.0)
        norm_rgb = np.clip((f_rgb / bg) * 255.0, 0.0, 255.0)
        # Push near-white paper background to pure 255 for clean look & better compression
        lum = norm_rgb.mean(axis=2)
        paper_mask = lum > 238.0
        norm_rgb[paper_mask] = 255.0
        if border_px > 0 and norm_rgb.shape[0] > 4 * border_px and norm_rgb.shape[1] > 4 * border_px:
            norm_rgb[:border_px, :, :] = 255.0
            norm_rgb[-border_px:, :, :] = 255.0
            norm_rgb[:, :border_px, :] = 255.0
            norm_rgb[:, -border_px:, :] = 255.0
        return Image.fromarray(norm_rgb.astype(np.uint8), mode="RGB"), True

    # Grayscale path: convert RGB -> luminance
    f = (
        0.299 * rgb[:, :, 0].astype(np.float32)
        + 0.587 * rgb[:, :, 1].astype(np.float32)
        + 0.114 * rgb[:, :, 2].astype(np.float32)
    )

    # 1. Estimate slow-varying paper background illumination and normalize to 255.
    bg = ndimage.gaussian_filter(f, sigma=35.0, mode="nearest")
    bg = np.maximum(bg, 1.0)
    norm = np.clip((f / bg) * 255.0, 0.0, 255.0)

    # 2. Local Sauvola threshold to identify foreground (text + grayscale diagrams/shading)
    # vs. paper background / bleed-through.
    mean = ndimage.uniform_filter(norm, size=sauvola_window, mode="nearest")
    mean_sq = ndimage.uniform_filter(norm * norm, size=sauvola_window, mode="nearest")
    var = np.maximum(mean_sq - mean * mean, 0.0)
    std = np.sqrt(var)
    thresh = mean * (1.0 + sauvola_k * ((std / 128.0) - 1.0))

    fg = (norm < thresh) | (norm < 220.0)

    # 3. Strip dark scanner edge bars along the outer page border.
    if border_px > 0 and fg.shape[0] > 4 * border_px and fg.shape[1] > 4 * border_px:
        fg[:border_px, :] = False
        fg[-border_px:, :] = False
        fg[:, :border_px] = False
        fg[:, -border_px:] = False

    # 4. Connected-component despeckling: remove tiny isolated salt-and-pepper specks
    # without eroding thin strokes or Kannada vowel signs.
    if min_speck_area > 1:
        labels, num = ndimage.label(fg)
        if num > 0:
            counts = np.bincount(labels.ravel())
            remove_mask = counts < min_speck_area
            remove_mask[0] = False
            fg[remove_mask[labels]] = False

    # 5. Contrast-stretch foreground grayscale while setting background to pure white (255).
    # Soft ramp near threshold keeps anti-aliased glyph edges smooth.
    cleaned = np.where(fg, norm, 255.0)
    # Darken ink core slightly for crisp reading/OCR while preserving midtones
    cleaned = np.clip((cleaned - 25.0) * (255.0 / 215.0), 0.0, 255.0)
    # Quantize to 32 gray levels so lossless/near-lossless WebP compresses tightly
    quantized = (np.round(cleaned / 8.0) * 8.0).clip(0.0, 255.0).astype(np.uint8)
    return Image.fromarray(quantized, mode="L"), False


def render_page_pixmap(page: fitz.Page, dpi: int, colorspace: fitz.Colorspace) -> Image.Image:
    zoom = dpi / 72.0
    mat = fitz.Matrix(zoom, zoom)
    pix = page.get_pixmap(matrix=mat, colorspace=colorspace, alpha=False)
    mode = "L" if colorspace == fitz.csGRAY else "RGB"
    return Image.frombytes(mode, (pix.width, pix.height), pix.samples)


def process_pdf(
    entry: dict,
    base_dir: Path,
    max_pages: int | None = None,
    force: bool = False,
) -> tuple[str, int, int, str | None]:
    """Process all pages of a single PDF entry from metadata.json.

    Returns (directory_slug, converted_pages, total_pages, error_msg).
    """
    rel_source = entry["source"]
    out_slug = entry["directory"]
    pdf_path = base_dir / rel_source
    out_dir = base_dir / out_slug

    if not pdf_path.is_file():
        return (out_slug, 0, 0, f"missing source PDF: {rel_source}")

    thumb_dir = out_dir / "thumb"
    full_dir = out_dir / "full"
    bw_dir = out_dir / "bw"
    thumb_dir.mkdir(parents=True, exist_ok=True)
    full_dir.mkdir(parents=True, exist_ok=True)
    bw_dir.mkdir(parents=True, exist_ok=True)

    try:
        doc = fitz.open(pdf_path)
    except Exception as exc:  # noqa: BLE001
        return (out_slug, 0, 0, f"failed to open PDF: {exc}")

    total_pages = len(doc)
    limit = min(total_pages, max_pages) if max_pages else total_pages
    converted = 0

    try:
        for idx in range(limit):
            page_num = idx + 1
            fname = f"{page_num:04d}.webp"
            thumb_path = thumb_dir / fname
            full_path = full_dir / fname
            bw_path = bw_dir / fname

            if not force and thumb_path.is_file() and full_path.is_file() and bw_path.is_file():
                continue

            page = doc.load_page(idx)

            # 1 & 2: Full resolution + Thumbnail (derived from full render)
            if force or not full_path.is_file() or not thumb_path.is_file():
                full_img = render_page_pixmap(page, FULL_DPI, fitz.csRGB)
                if force or not full_path.is_file():
                    full_img.save(full_path, format="WEBP", quality=82, method=4)
                if force or not thumb_path.is_file():
                    w, h = full_img.size
                    thumb_h = max(1, int(round(h * (THUMB_WIDTH / float(w)))))
                    thumb_img = full_img.resize((THUMB_WIDTH, thumb_h), Image.Resampling.LANCZOS)
                    thumb_img.save(thumb_path, format="WEBP", quality=70, method=4)

            # 3: Clean grayscale (or color if page has color pictures/diagrams) for reading & OCR
            if force or not bw_path.is_file():
                rgb_img = render_page_pixmap(page, BW_DPI, fitz.csRGB)
                bw_img, is_color = clean_page_for_reading_and_ocr(np.asarray(rgb_img))
                if is_color:
                    bw_img.save(bw_path, format="WEBP", quality=78, method=4)
                else:
                    bw_img.save(bw_path, format="WEBP", quality=75, method=4)

            converted += 1
    except Exception as exc:  # noqa: BLE001
        doc.close()
        return (out_slug, converted, total_pages, f"page {idx + 1} error: {exc}")

    doc.close()
    return (out_slug, converted, total_pages, None)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--metadata", type=Path, default=Path("metadata.json"), help="Path to metadata.json")
    p.add_argument("--only", type=str, default=None, help="Only process books whose directory matches this substring")
    p.add_argument("--max-pages", type=int, default=None, help="Maximum pages to convert per PDF")
    p.add_argument("--workers", type=int, default=max(1, (os.cpu_count() or 4) // 2), help="Parallel worker processes")
    p.add_argument("--force", action="store_true", help="Re-generate existing page files")
    return p.parse_args()


def main() -> int:
    args = parse_args()
    meta_path = args.metadata.resolve()
    base_dir = meta_path.parent

    entries = json.loads(meta_path.read_text(encoding="utf-8"))
    if args.only:
        entries = [e for e in entries if args.only in e["directory"] or args.only in e["source"]]

    if not entries:
        print("No matching entries found in metadata.json.")
        return 0

    print(f"Processing {len(entries)} PDF(s) with {args.workers} worker(s)...")
    errors = 0

    if args.workers <= 1 or len(entries) == 1:
        for i, entry in enumerate(entries, 1):
            slug, conv, total, err = process_pdf(entry, base_dir, args.max_pages, args.force)
            status = f"ERROR ({err})" if err else f"{conv} converted / {total} total pages"
            print(f"[{i}/{len(entries)}] {slug}: {status}", flush=True)
            if err:
                errors += 1
    else:
        with concurrent.futures.ProcessPoolExecutor(max_workers=args.workers) as pool:
            futs = {
                pool.submit(process_pdf, entry, base_dir, args.max_pages, args.force): entry
                for entry in entries
            }
            for i, fut in enumerate(concurrent.futures.as_completed(futs), 1):
                slug, conv, total, err = fut.result()
                status = f"ERROR ({err})" if err else f"{conv} converted / {total} total pages"
                print(f"[{i}/{len(entries)}] {slug}: {status}", flush=True)
                if err:
                    errors += 1

    return 1 if errors else 0


if __name__ == "__main__":
    sys.exit(main())