#!/usr/bin/env python3
"""
ocr-advanced.py — Multi-variant, multi-PSM OCR pipeline for omnius.

Implements a full preprocessing + OCR + cross-reference pipeline:
  1. Load image → grayscale → 2x upscale
  2. Generate 7 preprocessing variants (two adaptive windows, OTSU, two fixed
     thresholds, two sharpen kernels, denoise)
  3. Run Tesseract with PSM 4, 6, 11 on each variant (up to 21 passes)
  4. Score results using combined heuristic (confidence * coverage + line bonus)
  5. Optionally extract regions (header/body/footer) with cross-reference
  6. Output as JSON, text, CSV, or write all formats to an output directory

Usage:
  python3 ocr-advanced.py <image_or_dir> [options]

Single image:
  python3 ocr-advanced.py photo.jpg --output json
  python3 ocr-advanced.py scan.png --output-dir ./ocr_out --regions

Batch directory:
  python3 ocr-advanced.py ./images/ --output-dir ./ocr_out --batch

Output (JSON to stdout):
  {
    "text": "best extracted text",
    "confidence": 85.2,
    "variant": "otsu_psm6",
    "lines": 42,
    "all_variants": { ... },
    "regions": { ... }
  }
"""

import sys
import os
import json
import csv
import argparse
from pathlib import Path

def check_deps():
    """Check that required Python packages are available."""
    missing = []
    try:
        import cv2
    except ImportError:
        missing.append("opencv-python-headless")
    try:
        import numpy
    except ImportError:
        missing.append("numpy")
    try:
        import pytesseract
    except ImportError:
        missing.append("pytesseract")
    try:
        from PIL import Image
    except ImportError:
        missing.append("Pillow")

    if missing:
        print(json.dumps({
            "error": f"Missing Python packages: {', '.join(missing)}. "
                     f"Install with: pip install {' '.join(missing)}",
            "missing": missing,
        }))
        sys.exit(1)

check_deps()

import cv2
import numpy as np
import pytesseract
from PIL import Image

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".tiff", ".tif", ".bmp", ".webp"}


# ---------------------------------------------------------------------------
# Image preprocessing variants
# ---------------------------------------------------------------------------

def to_grayscale(img):
    """Convert BGR to grayscale if needed."""
    if len(img.shape) == 3:
        return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return img


def upscale_2x(gray):
    """2x bicubic upscale for better OCR character recognition."""
    h, w = gray.shape
    return cv2.resize(gray, (w * 2, h * 2), interpolation=cv2.INTER_CUBIC)


def variant_adaptive_wide(gray):
    """Adaptive Gaussian threshold — wide window (31px), handles gradual lighting."""
    return cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY, 31, 10
    )


def variant_adaptive_fine(gray):
    """Adaptive Gaussian threshold — fine window (11px), catches small text detail."""
    return cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY, 11, 2
    )


def variant_otsu(gray):
    """OTSU threshold — optimal global threshold for bimodal images."""
    _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return binary


def variant_fixed_140(gray):
    """Fixed threshold 140 — standard cutoff for dark text on light paper."""
    _, binary = cv2.threshold(gray, 140, 255, cv2.THRESH_BINARY)
    return binary


def variant_fixed_150(gray):
    """Fixed threshold 150 — slightly brighter cutoff for lighter scans."""
    _, binary = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)
    return binary


def variant_sharpen_laplacian_otsu(gray):
    """Laplacian sharpen + OTSU — aggressive edge enhancement."""
    kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
    sharpened = cv2.filter2D(gray, -1, kernel)
    _, binary = cv2.threshold(sharpened, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return binary


def variant_sharpen_unsharp_otsu(gray):
    """Unsharp mask sharpen + OTSU — gentler enhancement, better for photos."""
    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
    sharpened = cv2.filter2D(gray, -1, kernel)
    _, binary = cv2.threshold(sharpened, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return binary


def variant_denoise_otsu(gray):
    """Denoise + OTSU — removes JPEG artifacts and photo noise."""
    denoised = cv2.fastNlMeansDenoising(gray, h=10, templateWindowSize=7, searchWindowSize=21)
    _, binary = cv2.threshold(denoised, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return binary


ALL_VARIANTS = {
    "adaptive_wide":   variant_adaptive_wide,
    "adaptive_fine":   variant_adaptive_fine,
    "otsu":            variant_otsu,
    "fixed_140":       variant_fixed_140,
    "fixed_150":       variant_fixed_150,
    "sharpen_lap":     variant_sharpen_laplacian_otsu,
    "sharpen_unsharp": variant_sharpen_unsharp_otsu,
    "denoise":         variant_denoise_otsu,
}

PSM_MODES = {
    4: "single_block",
    6: "default",
    11: "sparse",
}


# ---------------------------------------------------------------------------
# OCR execution
# ---------------------------------------------------------------------------

def run_tesseract(binary_img, language="eng", psm=6):
    """Run Tesseract on a preprocessed binary image.
    Returns (text, confidence, line_count)."""
    pil_img = Image.fromarray(binary_img)
    config = f"--psm {psm}"

    try:
        text = pytesseract.image_to_string(pil_img, lang=language, config=config).strip()
    except Exception:
        return "", 0.0, 0

    line_count = len([l for l in text.split("\n") if l.strip()])

    # Get confidence via image_to_data
    try:
        data = pytesseract.image_to_data(
            pil_img, lang=language, config=config,
            output_type=pytesseract.Output.DICT,
        )
        confs = [int(c) for c in data["conf"] if int(c) >= 0]
        avg_conf = sum(confs) / len(confs) if confs else 0.0
    except Exception:
        avg_conf = 0.0

    return text, avg_conf, line_count


def compute_score(text, confidence, line_count):
    """Combined scoring heuristic:
      - confidence * sqrt(char_count)  — rewards quality and coverage
      - + line_count * 10              — bonus for structured output (more lines = better parse)
    The agent discovered that line-count is a strong proxy for successful parsing
    on structured documents like invoices and forms."""
    char_count = len(text)
    if char_count == 0:
        return 0
    return confidence * (char_count ** 0.5) + line_count * 10


def extract_region(gray, y_start_pct, y_end_pct, x_start_pct=0, x_end_pct=100):
    """Extract a region from the image by percentage coordinates."""
    h, w = gray.shape
    y1 = int(h * y_start_pct / 100)
    y2 = int(h * y_end_pct / 100)
    x1 = int(w * x_start_pct / 100)
    x2 = int(w * x_end_pct / 100)
    return gray[y1:y2, x1:x2]


def extract_pixel_region(gray, x, y, w, h):
    """Extract a region by pixel coordinates."""
    return gray[y:y+h, x:x+w]


# ---------------------------------------------------------------------------
# Output writers
# ---------------------------------------------------------------------------

def write_txt(text, output_path):
    """Write plain text output."""
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(text)


def write_csv(text, output_path):
    """Write CSV with line numbers."""
    lines = [l for l in text.split("\n") if l.strip()]
    with open(output_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["Line_Number", "Extracted_Text"])
        for i, line in enumerate(lines, 1):
            writer.writerow([i, line])


def write_pdf(text, output_path):
    """Write searchable PDF using reportlab (if available)."""
    try:
        from reportlab.lib.pagesizes import letter
        from reportlab.pdfgen import canvas as pdf_canvas
        from reportlab.lib.units import inch
    except ImportError:
        return False

    lines = [l for l in text.split("\n") if l.strip()]
    c = pdf_canvas.Canvas(output_path, pagesize=letter)
    width, height = letter
    y = height - 1 * inch

    for line in lines:
        if y < 1 * inch:
            c.showPage()
            y = height - 1 * inch
        c.drawString(1 * inch, y, line)
        y -= 14

    c.save()
    return True


def write_all_outputs(text, base_name, output_dir):
    """Write TXT + CSV + PDF to output directory."""
    os.makedirs(output_dir, exist_ok=True)

    txt_path = os.path.join(output_dir, f"{base_name}.txt")
    write_txt(text, txt_path)

    csv_path = os.path.join(output_dir, f"{base_name}.csv")
    write_csv(text, csv_path)

    pdf_path = os.path.join(output_dir, f"{base_name}.pdf")
    pdf_ok = write_pdf(text, pdf_path)

    return {
        "txt": txt_path,
        "csv": csv_path,
        "pdf": pdf_path if pdf_ok else None,
    }


# ---------------------------------------------------------------------------
# Main pipeline
# ---------------------------------------------------------------------------

def run_pipeline(image_path, language="eng", do_regions=False, debug_dir=None,
                 single_psm=None, pixel_region=None, output_dir=None):
    """Run the full multi-variant, multi-PSM OCR pipeline."""

    # Load image
    img = cv2.imread(image_path)
    if img is None:
        return {"error": f"Could not load image: {image_path}"}

    h_orig, w_orig = img.shape[:2]
    gray = to_grayscale(img)

    # Upscale 2x
    gray_2x = upscale_2x(gray)

    # If a pixel region is specified, crop before processing
    if pixel_region:
        rx, ry, rw, rh = pixel_region
        # Scale region coords to match 2x upscale
        gray_2x = extract_pixel_region(gray_2x, rx * 2, ry * 2, rw * 2, rh * 2)

    # Determine PSM modes to test
    psm_modes = [single_psm] if single_psm else [4, 6, 11]

    # Generate all variants and run OCR
    all_results = {}
    best_key = None
    best_score = -1

    for vname, vfunc in ALL_VARIANTS.items():
        try:
            binary = vfunc(gray_2x)
        except Exception:
            continue

        # Save debug images
        if debug_dir:
            os.makedirs(debug_dir, exist_ok=True)
            cv2.imwrite(os.path.join(debug_dir, f"full_{vname}.png"), binary)

        for psm in psm_modes:
            key = f"{vname}_psm{psm}"
            text, confidence, line_count = run_tesseract(binary, language, psm)
            char_count = len(text)
            score = compute_score(text, confidence, line_count)

            all_results[key] = {
                "text": text,
                "chars": char_count,
                "lines": line_count,
                "confidence": round(confidence, 1),
                "score": round(score, 1),
            }

            if score > best_score:
                best_score = score
                best_key = key

    if not best_key:
        return {"error": "All OCR variants failed to produce output"}

    best = all_results[best_key]
    result = {
        "text": best["text"],
        "confidence": best["confidence"],
        "variant": best_key,
        "chars": best["chars"],
        "lines": best["lines"],
        "score": best["score"],
        "image_size": f"{w_orig}x{h_orig}",
        "variants_tested": len(all_results),
        "all_variants": all_results,
    }

    # Region-based OCR
    if do_regions:
        regions = {}
        region_defs = {
            "header": (0, 35),
            "body": (30, 80),
            "footer": (75, 100),
        }

        for rname, (y_start, y_end) in region_defs.items():
            region_gray = extract_region(gray_2x, y_start, y_end)

            if debug_dir:
                cv2.imwrite(os.path.join(debug_dir, f"region_{rname}.png"), region_gray)

            # Test all variants on each region for best accuracy
            region_best = ""
            region_best_score = -1

            for vname in ["otsu", "denoise", "adaptive_fine", "sharpen_unsharp"]:
                if vname not in ALL_VARIANTS:
                    continue
                try:
                    binary = ALL_VARIANTS[vname](region_gray)
                except Exception:
                    continue

                if debug_dir:
                    cv2.imwrite(os.path.join(debug_dir, f"region_{rname}_{vname}.png"), binary)

                text, conf, lc = run_tesseract(binary, language, 6)
                score = compute_score(text, conf, lc)
                if score > region_best_score:
                    region_best_score = score
                    region_best = text

            regions[rname] = region_best

        result["regions"] = regions

    if debug_dir:
        result["debug_dir"] = debug_dir

    # Write output files if output_dir specified
    if output_dir:
        base_name = Path(image_path).stem
        files = write_all_outputs(best["text"], base_name, output_dir)
        result["output_files"] = files

    return result


def run_batch(images_dir, language="eng", do_regions=False, debug_dir=None,
              output_dir=None):
    """Process all images in a directory."""
    images_dir = os.path.abspath(images_dir)
    if not os.path.isdir(images_dir):
        return {"error": f"Not a directory: {images_dir}"}

    out_dir = output_dir or os.path.join(images_dir, "ocr_out")
    os.makedirs(out_dir, exist_ok=True)

    batch_results = {}
    image_files = sorted(
        f for f in os.listdir(images_dir)
        if Path(f).suffix.lower() in IMAGE_EXTENSIONS
    )

    if not image_files:
        return {"error": f"No image files found in {images_dir}"}

    for img_file in image_files:
        img_path = os.path.join(images_dir, img_file)
        img_debug = os.path.join(debug_dir, Path(img_file).stem) if debug_dir else None
        result = run_pipeline(
            img_path,
            language=language,
            do_regions=do_regions,
            debug_dir=img_debug,
            output_dir=out_dir,
        )
        # Compact per-image result (omit all_variants for batch summary)
        batch_results[img_file] = {
            "text": result.get("text", ""),
            "confidence": result.get("confidence", 0),
            "variant": result.get("variant", ""),
            "chars": result.get("chars", 0),
            "lines": result.get("lines", 0),
            "output_files": result.get("output_files"),
            "error": result.get("error"),
        }

    # Write summary
    summary_path = os.path.join(out_dir, "OCR_PROCESSING_SUMMARY.md")
    with open(summary_path, "w", encoding="utf-8") as f:
        f.write("# OCR Processing Summary Report\n\n")
        f.write(f"**Source:** `{images_dir}`\n\n")
        f.write("## Processed Documents\n\n")
        f.write("| Document | Lines | Chars | Confidence | Variant |\n")
        f.write("|----------|-------|-------|------------|----------|\n")
        for img, data in batch_results.items():
            if data.get("error"):
                f.write(f"| {img} | ERROR | - | - | {data['error']} |\n")
            else:
                f.write(
                    f"| {img} | {data['lines']} | {data['chars']} "
                    f"| {data['confidence']}% | {data['variant']} |\n"
                )

    return {
        "batch": True,
        "images_processed": len(batch_results),
        "output_dir": out_dir,
        "summary": summary_path,
        "results": batch_results,
    }


def main():
    parser = argparse.ArgumentParser(
        description="Advanced multi-variant OCR pipeline for omnius"
    )
    parser.add_argument(
        "image",
        help="Path to image file, or directory for --batch mode",
    )
    parser.add_argument("--language", "-l", default="eng",
                        help="OCR language (default: eng)")
    parser.add_argument("--regions", action="store_true",
                        help="Also OCR header/body/footer regions")
    parser.add_argument("--debug-dir",
                        help="Save preprocessed images to this directory")
    parser.add_argument("--psm", type=int, choices=[4, 6, 11],
                        help="Use single PSM mode instead of all 3")
    parser.add_argument("--region",
                        help="Crop region before OCR: x,y,w,h in pixels")
    parser.add_argument("--output", choices=["json", "text"], default="json",
                        help="Stdout output format (default: json)")
    parser.add_argument("--output-dir",
                        help="Write TXT + CSV + PDF outputs to this directory")
    parser.add_argument("--batch", action="store_true",
                        help="Process all images in a directory")

    args = parser.parse_args()

    # Batch mode
    if args.batch or os.path.isdir(args.image):
        result = run_batch(
            args.image,
            language=args.language,
            do_regions=args.regions,
            debug_dir=args.debug_dir,
            output_dir=args.output_dir,
        )
        if args.output == "text":
            if "error" in result:
                print(f"ERROR: {result['error']}", file=sys.stderr)
                sys.exit(1)
            print(f"Processed {result['images_processed']} images → {result['output_dir']}")
        else:
            print(json.dumps(result, indent=2))
        sys.exit(0)

    # Single image mode
    if not os.path.isfile(args.image):
        print(json.dumps({"error": f"File not found: {args.image}"}))
        sys.exit(1)

    pixel_region = None
    if args.region:
        try:
            pixel_region = tuple(int(x) for x in args.region.split(","))
            if len(pixel_region) != 4:
                raise ValueError
        except ValueError:
            print(json.dumps({"error": "Region must be x,y,w,h (4 integers)"}))
            sys.exit(1)

    result = run_pipeline(
        args.image,
        language=args.language,
        do_regions=args.regions,
        debug_dir=args.debug_dir,
        single_psm=args.psm,
        pixel_region=pixel_region,
        output_dir=args.output_dir,
    )

    if args.output == "text":
        if "error" in result:
            print(f"ERROR: {result['error']}", file=sys.stderr)
            sys.exit(1)
        print(result["text"])
    else:
        print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()
