"""
Image Preprocessor — OpenCV pipeline for cleaning dirty scans before OCR.

Designed for low-quality input common in African contexts:
poor lighting, skewed documents, noise, low resolution.

IMPORTANT: Tesseract works best with clean grayscale or lightly processed images.
Over-processing (aggressive binarization, morphological ops) destroys text on
low-quality scans. This pipeline uses a TIERED approach:

  Tier 1 (Gentle): Grayscale → Upscale → Light denoise → CLAHE contrast
  Tier 2 (Medium): Adds Otsu binarization
  Tier 3 (Heavy):  Adds morphological shadow removal + adaptive threshold

The OCR engine tries Tier 1 first. If confidence is too low, it escalates.
"""
import cv2
import numpy as np
from pathlib import Path
from PIL import Image
import pytesseract
import os
import sys

# Auto-detect Tesseract on Windows
if sys.platform == "win32":
    _win_paths = [
        r"C:\Program Files\Tesseract-OCR\tesseract.exe",
        r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe",
        os.path.expandvars(r"%LOCALAPPDATA%\Programs\Tesseract-OCR\tesseract.exe"),
    ]
    for _p in _win_paths:
        if os.path.exists(_p):
            pytesseract.pytesseract.tesseract_cmd = _p
            break


def fix_orientation(img_bgr: np.ndarray) -> np.ndarray:
    """
    Detect the correct orientation by actively probing 4 angles (0, 90, 180, 270).
    It runs a fast OCR pass on a high-res center crop of the image to pick the 
    angle with the most valid text and highest confidence.
    """
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    
    # Do NOT downscale! Downscaling destroys text legibility for Tesseract.
    # Instead, crop a 1200x1200 square from the center to keep it fast but high-res.
    h, w = gray.shape[:2]
    crop_size = min(1200, min(h, w))
    start_y = (h - crop_size) // 2
    start_x = (w - crop_size) // 2
    probe_img = gray[start_y:start_y+crop_size, start_x:start_x+crop_size]

    # Binarize to make the probe even faster and clearer
    _, probe_img = cv2.threshold(probe_img, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)

    best_score = -1
    best_angle = 0
    
    angles = [0, 90, 180, 270]
    for angle in angles:
        if angle == 90:
            test_img = cv2.rotate(probe_img, cv2.ROTATE_90_CLOCKWISE)
        elif angle == 180:
            test_img = cv2.rotate(probe_img, cv2.ROTATE_180)
        elif angle == 270:
            test_img = cv2.rotate(probe_img, cv2.ROTATE_90_COUNTERCLOCKWISE)
        else:
            test_img = probe_img
            
        try:
            # psm 11 = Sparse text
            data = pytesseract.image_to_data(test_img, output_type=pytesseract.Output.DICT, config='--psm 11')
            
            total_conf = 0.0
            word_count = 0
            # We will score based on the length of words. Real text has long words. Gibberish is mostly single characters.
            word_length_score = 0
            
            for i in range(len(data['text'])):
                text = data['text'][i].strip()
                conf = data['conf'][i]
                
                try:
                    conf_val = float(conf)
                except (ValueError, TypeError):
                    conf_val = -1.0
                    
                # Only consider reasonably confident detections that contain actual letters
                if conf_val > 50 and text and any(c.isalpha() for c in text):
                    # Clean the word of punctuation to get its true length
                    clean_word = ''.join(c for c in text if c.isalnum())
                    if len(clean_word) > 1: # Ignore single character gibberish
                        # Square the length to heavily reward actual words over random noise
                        word_length_score += (len(clean_word) ** 2)
                        total_conf += conf_val
                        word_count += 1
                        
            avg_conf = (total_conf / word_count) if word_count > 0 else 0
            score = word_length_score * (avg_conf / 100.0)
            
            if score > best_score:
                best_score = score
                best_angle = angle
                
        except Exception as e:
            print(f"Probe failed for angle {angle}: {e}")

    # Fallback to landscape heuristic if probing completely failed (score <= 0)
    if best_score <= 0:
        h, w = img_bgr.shape[:2]
        if w > h:
            # 270 degrees (counter-clockwise) is much more common for landscape scans 
            # (e.g. holding phone sideways with top to the left) than 90 degrees.
            best_angle = 270

    print(f"Active orientation probe decided angle: {best_angle} (Score: {best_score})")

    if best_angle == 90:
        return cv2.rotate(img_bgr, cv2.ROTATE_90_CLOCKWISE)
    elif best_angle == 180:
        return cv2.rotate(img_bgr, cv2.ROTATE_180)
    elif best_angle == 270:
        return cv2.rotate(img_bgr, cv2.ROTATE_90_COUNTERCLOCKWISE)
        
    return img_bgr


def _ensure_min_resolution(gray_img: np.ndarray, min_width: int = 2000) -> np.ndarray:
    """
    Upscale image if it's too small for reliable OCR.
    Target: at least 2000px wide (roughly 300 DPI for a standard page).
    """
    h, w = gray_img.shape[:2]
    if w < min_width:
        scale = min_width / w
        new_w = int(w * scale)
        new_h = int(h * scale)
        return cv2.resize(gray_img, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
    return gray_img


def _deskew(img: np.ndarray) -> np.ndarray:
    """
    Detect and correct skew angle using minAreaRect on contours.
    Falls back to no correction if skew detection fails.
    """
    try:
        # If the image isn't binary, threshold it just for skew detection
        if len(img.shape) == 2:
            _, thresh = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        else:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

        coords = np.column_stack(np.where(thresh > 0))
        if len(coords) < 100:
            return img

        angle = cv2.minAreaRect(coords)[-1]

        if angle < -45:
            angle = -(90 + angle)
        else:
            angle = -angle

        # Only correct if skew is significant but not extreme
        if abs(angle) < 0.5 or abs(angle) > 15:
            return img

        h, w = img.shape[:2]
        center = (w // 2, h // 2)
        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        rotated = cv2.warpAffine(
            img, rotation_matrix, (w, h),
            flags=cv2.INTER_CUBIC,
            borderMode=cv2.BORDER_REPLICATE,
        )
        return rotated
    except Exception:
        return img


def preprocess_gentle(img: np.ndarray) -> np.ndarray:
    """
    Tier 1 — Gentle preprocessing. Best for most scanned documents.
    
    Steps: Grayscale → Upscale → Light Denoise → CLAHE → Deskew
    
    Does NOT binarize. Tesseract handles grayscale images very well and its
    internal Otsu thresholding is often better than ours.
    """
    # Convert to grayscale
    if len(img.shape) == 3:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    else:
        gray = img

    # Upscale small images
    gray = _ensure_min_resolution(gray)

    # Light denoise — gentler than before (h=6 vs h=10)
    denoised = cv2.fastNlMeansDenoising(gray, None, h=6, templateWindowSize=7, searchWindowSize=21)

    # CLAHE contrast enhancement — gentler settings
    clahe = cv2.createCLAHE(clipLimit=1.5, tileGridSize=(8, 8))
    enhanced = clahe.apply(denoised)

    # Deskew
    result = _deskew(enhanced)

    return result


def preprocess_medium(img: np.ndarray) -> np.ndarray:
    """
    Tier 2 — Medium preprocessing. For documents with moderate shadows.
    
    Steps: Gentle pipeline + Otsu binarization
    """
    # Start with the gentle pipeline
    enhanced = preprocess_gentle(img)

    # Add Otsu binarization (let OpenCV pick the optimal threshold)
    _, binary = cv2.threshold(enhanced, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    return binary


def preprocess_heavy(img: np.ndarray) -> np.ndarray:
    """
    Tier 3 — Heavy preprocessing. For severely degraded scans with deep shadows.
    
    Steps: Grayscale → Upscale → Morphological Background Division → CLAHE → Denoise → Adaptive Threshold → Deskew
    """
    if len(img.shape) == 3:
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    else:
        gray = img
        
    gray = _ensure_min_resolution(gray)

    # Morphological Background Division (flattens shadows & wrinkles)
    kernel_size = 45
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
    background = cv2.morphologyEx(gray, cv2.MORPH_DILATE, kernel)
    diff = 255 - cv2.absdiff(gray, background)

    # CLAHE
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    enhanced = clahe.apply(diff)

    # Denoise
    denoised = cv2.fastNlMeansDenoising(enhanced, None, h=10, templateWindowSize=7, searchWindowSize=21)

    # Adaptive threshold
    binary = cv2.adaptiveThreshold(
        denoised, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 51, 15
    )

    result = _deskew(binary)
    return result


def preprocess_image(image_path: str | Path, output_path: str | Path | None = None) -> np.ndarray:
    """
    Full preprocessing pipeline for a dirty scan or photo.
    
    Uses the GENTLE tier by default (best for Tesseract).
    The OCR engine's adaptive pipeline will escalate if needed.
    """
    img = cv2.imread(str(image_path))
    if img is None:
        raise ValueError(f"Could not load image: {image_path}")
        
    img = fix_orientation(img)
    result = preprocess_gentle(img)

    if output_path:
        cv2.imwrite(str(output_path), result)

    return result


def preprocess_for_ocr(image_path: str | Path) -> str:
    """
    Preprocess an image and save the cleaned version to a temp file.
    Returns the path to the cleaned image for OCR consumption.
    """
    output_path = Path(str(image_path)).with_suffix(".cleaned.png")
    preprocess_image(image_path, output_path)
    return str(output_path)