#!/usr/bin/env python3
"""
Process all extracted raw_pages with adaptive threshold + morphology,
then combine into PDF, then clean up PNGs.
"""
import os, sys, subprocess
from pathlib import Path

try:
    import cv2
    import numpy as np
except ImportError:
    print("❌ Cần opencv-python")
    sys.exit(1)

BASE = Path(__file__).parent
RAW = BASE / "raw_pages"
CLEANED = BASE / "cleaned"
PDF_SRC = BASE / "pdf" / "so-tay-mahavihara.pdf"
OUTPUT = BASE / "so-tay-mahavihara-preprocessed.pdf"

os.makedirs(CLEANED, exist_ok=True)

# Params from pipeline doc (same as demo)
params = {
    'bilateral_d': 9,
    'bilateral_sigmaColor': 75,
    'bilateral_sigmaSpace': 75,
    'adaptive_block_size': 21,
    'adaptive_C': 25,
    'morph_kernel': (3, 3),
    'min_area': 10,
}

png_files = sorted([f for f in os.listdir(RAW) if f.endswith(".png")])
print(f"📄 Tổng số: {len(png_files)} trang")

# Step 1: Preprocess all
print("\n🔧 STEP 1: Adaptive Threshold + Morphology")
success = 0
for f in png_files:
    inp = str(RAW / f)
    out = str(CLEANED / f)
    print(f"   {f} ...", end=" ", flush=True)

    img = cv2.imread(inp, cv2.IMREAD_GRAYSCALE)
    if img is None:
        print("❌ skip")
        continue

    # Bilateral filter — giữ nét chữ, làm mờ nhiễu nền
    bf = cv2.bilateralFilter(img, d=params['bilateral_d'],
                              sigmaColor=params['bilateral_sigmaColor'],
                              sigmaSpace=params['bilateral_sigmaSpace'])

    # Adaptive threshold
    binary = cv2.adaptiveThreshold(bf, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                    cv2.THRESH_BINARY,
                                    params['adaptive_block_size'],
                                    params['adaptive_C'])

    # Morphology: close (fill small holes)
    kernel = np.ones(params['morph_kernel'], np.uint8)
    cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)

    # Remove small specks
    if params['min_area'] > 0:
        num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(~cleaned, 8, cv2.CV_32S)
        for i in range(1, num_labels):
            if stats[i, cv2.CC_STAT_AREA] < params['min_area']:
                cleaned[labels == i] = 255

    cv2.imwrite(out, cleaned)
    in_kb = os.path.getsize(inp) / 1024
    out_kb = os.path.getsize(out) / 1024
    print(f"✅ {in_kb:.0f}KB → {out_kb:.0f}KB")
    success += 1

print(f"\n   ✅ {success}/{len(png_files)} trang processed")

# Step 2: Combine into PDF
print("\n🔧 STEP 2: Gộp PNG → PDF")
cleaned_pngs = sorted([str(CLEANED / f) for f in os.listdir(CLEANED) if f.endswith(".png")])
if not cleaned_pngs:
    print("❌ Không có ảnh nào trong cleaned/")
    sys.exit(1)

cmd = ["img2pdf"] + cleaned_pngs + ["-o", str(OUTPUT)]
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode == 0:
    out_kb = os.path.getsize(OUTPUT) / 1024
    orig_kb = os.path.getsize(PDF_SRC) / 1024
    pct = (1 - out_kb / orig_kb) * 100
    print(f"   ✅ {OUTPUT.name} ({out_kb:.0f} KB)")
    print(f"   📊 PDF gốc: {orig_kb:.0f} KB → PDF sạch: {out_kb:.0f} KB (giảm {pct:.0f}%)")
else:
    print(f"   ❌ Lỗi img2pdf: {result.stderr}")
    sys.exit(1)

# Step 3: Backup & Remove PNGs
print("\n🔧 STEP 3: Backup → xóa file PNG")
backup_dir = BASE / "_backup"
raw_backup = backup_dir / "raw_pages"
cleaned_backup = backup_dir / "cleaned"

# Backup raw_pages
if any(Path(RAW).iterdir()):
    os.makedirs(raw_backup, exist_ok=True)
    for f in os.listdir(RAW):
        os.rename(str(RAW / f), str(raw_backup / f))
    print(f"   ✅ raw_pages/ → _backup/raw_pages/")

# Backup cleaned
if any(Path(CLEANED).iterdir()):
    os.makedirs(cleaned_backup, exist_ok=True)
    for f in os.listdir(CLEANED):
        os.rename(str(CLEANED / f), str(cleaned_backup / f))
    print(f"   ✅ cleaned/ → _backup/cleaned/")

# Remove empty dirs
os.rmdir(RAW)
os.rmdir(CLEANED)
print(f"   ✅ raw_pages/ và cleaned/ đã xóa (đã backup vào _backup/)")

print("\n" + "=" * 50)
print("✅ HOÀN TẤT!")
print(f"   📄 Output: {OUTPUT}")
print("=" * 50)