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

# Support PyMuPDF (fitz)
HAVE_FITZ = False
try:
    import fitz
    HAVE_FITZ = True
except ImportError:
    pass

class PreprocessingPipeline:
    def __init__(self, config: dict = None):
        self.config = config or {}
        self.do_grayscale = self.config.get("grayscale", True)
        self.do_deskew = self.config.get("deskew", True)
        self.do_denoise = self.config.get("denoise", True)
        self.do_binarize = self.config.get("binarize", False)
        self.target_width = self.config.get("target_width", 1000)
        self.target_height = self.config.get("target_height", 1400)
        self.pdf_dpi = self.config.get("pdf_dpi", 200)

    def pdf_to_image(self, pdf_path: str, page_num: int = 0) -> np.ndarray:
        """Converts the specified page of a PDF file to a NumPy BGR image."""
        if HAVE_FITZ:
            doc = fitz.open(pdf_path)
            if page_num < len(doc):
                pix = doc[page_num].get_pixmap(dpi=self.pdf_dpi)
                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        raise RuntimeError("PDF rendering requires PyMuPDF (fitz).")

    def deskew(self, image: np.ndarray) -> np.ndarray:
        """Detects skew angle and rotates image to deskew."""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
        blur = cv2.GaussianBlur(gray, (9, 9), 0)
        thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]

        # Use morphological dilation to find text line blocks
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (30, 5))
        dilate = cv2.dilate(thresh, kernel, iterations=2)

        contours, _ = cv2.findContours(dilate, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        angles = []
        for c in contours:
            if cv2.contourArea(c) > 1000:
                rect = cv2.minAreaRect(c)
                angle = rect[-1]
                if angle < -45:
                    angle = -(90 + angle)
                elif angle > 45:
                    angle = 90 - angle
                if abs(angle) < 45:
                    angles.append(angle)

        if not angles:
            return image

        median_angle = float(np.median(angles))
        if abs(median_angle) < 0.5:
            return image

        (h, w) = image.shape[:2]
        center = (w // 2, h // 2)
        M = cv2.getRotationMatrix2D(center, median_angle, 1.0)
        rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
        return rotated

    def denoise(self, image: np.ndarray) -> np.ndarray:
        """Applies median filter denoising."""
        if len(image.shape) == 3:
            return cv2.fastNlMeansDenoisingColored(image, None, 5, 5, 7, 21)
        return cv2.medianBlur(image, 3)

    def binarize(self, gray_image: np.ndarray) -> np.ndarray:
        """Applies adaptive Gaussian thresholding."""
        return cv2.adaptiveThreshold(
            gray_image, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
        )

    def resize(self, image: np.ndarray) -> np.ndarray:
        """Resizes image to target dimensions while maintaining aspect ratio."""
        h, w = image.shape[:2]
        scale = min(self.target_width / w, self.target_height / h)
        new_w, new_h = int(w * scale), int(h * scale)
        return cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)

    def process(self, image_or_path) -> np.ndarray:
        """Runs the complete configurable preprocessing pipeline."""
        if isinstance(image_or_path, str):
            if image_or_path.lower().endswith(".pdf"):
                img = self.pdf_to_image(image_or_path)
            else:
                img = cv2.imread(image_or_path)
        elif isinstance(image_or_path, np.ndarray):
            img = image_or_path.copy()
        elif isinstance(image_or_path, Image.Image):
            img = cv2.cvtColor(np.array(image_or_path), cv2.COLOR_RGB2BGR)
        else:
            raise ValueError(f"Unsupported image input type: {type(image_or_path)}")

        if img is None:
            raise ValueError(f"Failed to load image from: {image_or_path}")

        # 1. Deskew
        if self.do_deskew:
            img = self.deskew(img)

        # 2. Denoise
        if self.do_denoise:
            img = self.denoise(img)

        # 3. Grayscale
        if self.do_grayscale and len(img.shape) == 3:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        # 4. Adaptive Binarization
        if self.do_binarize:
            if len(img.shape) == 3:
                img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            img = self.binarize(img)

        # 5. Resize
        if self.target_width and self.target_height:
            img = self.resize(img)

        return img