#!/usr/bin/env python3
"""
Tách file PDF thành ảnh đã tiền xử lý cho OCR.
- Render từng trang với PyMuPDF (fitz)
- Grayscale + tăng contrast
- Resize width ~1200px, giữ tỷ lệ
- Lưu dạng 001.jpg, 002.jpg, ... vào thư mục img/

Usage:
    python3 preprocess_pdf_ocr.py "/path/to/file.pdf"
"""

import sys
import os
from pathlib import Path

try:
    import fitz  # PyMuPDF
except ImportError:
    print("❌ Cần PyMuPDF: pip install pymupdf")
    sys.exit(1)

try:
    from PIL import Image, ImageEnhance
except ImportError:
    print("❌ Cần Pillow: pip install Pillow")
    sys.exit(1)


def pdf_to_ocr_images(pdf_path: str, output_dir: str = "img", target_width: int = 1200, dpi: int = 200):
    pdf_path = os.path.abspath(pdf_path)
    if not os.path.isfile(pdf_path):
        print(f"❌ Không tìm thấy file: {pdf_path}")
        sys.exit(1)

    # Tạo thư mục img bên cạnh file PDF
    base_dir = os.path.dirname(pdf_path)
    img_dir = os.path.join(base_dir, output_dir)
    os.makedirs(img_dir, exist_ok=True)

    print(f"📄 PDF: {os.path.basename(pdf_path)}")
    print(f"📏 DPI: {dpi}, Target width: {target_width}px")
    print(f"🖼️  Output: {img_dir}/")
    print()

    doc = fitz.open(pdf_path)
    total = len(doc)
    digits = len(str(total))
    print(f"📑 Tổng số trang: {total}")
    print()

    success = 0
    for i in range(total):
        page_num = i + 1
        filename = f"{page_num:0{digits}d}.jpg"
        out_path = os.path.join(img_dir, filename)

        print(f"[{page_num:0{digits}d}/{total}] Đang xử lý...", end="", flush=True)

        # Render page → pixmap
        # Dùng matrix để tăng độ phân giải
        zoom = dpi / 72  # fitz mặc định 72 DPI
        mat = fitz.Matrix(zoom, zoom)
        pix = doc[i].get_pixmap(matrix=mat)

        # Pixmap → PIL Image
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)

        # Grayscale
        img = img.convert("L")

        # Tăng contrast (gấp đôi)
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(2.0)

        # Resize width = target_width, giữ tỷ lệ
        w, h = img.size
        if w != target_width:
            ratio = target_width / w
            new_h = int(h * ratio)
            img = img.resize((target_width, new_h), Image.LANCZOS)

        # Sharpening nhẹ
        from PIL import ImageFilter
        img = img.filter(ImageFilter.SHARPEN)

        # Save JPEG quality cao
        img.save(out_path, "JPEG", quality=95, optimize=True)

        w_final, h_final = img.size
        size_kb = os.path.getsize(out_path) / 1024
        print(f" ✅ {w_final}x{h_final}px, {size_kb:.0f}KB")

        success += 1

    doc.close()
    print()
    print(f"✅ Hoàn tất! {success}/{total} trang đã lưu vào:")
    print(f"   {img_dir}")


if __name__ == "__main__":
    if len(sys.argv) > 1:
        pdf_to_ocr_images(sys.argv[1])
    else:
        print("Usage: python3 preprocess_pdf_ocr.py <path/to/file.pdf>")