import cv2
import os
import numpy as np
import fitz  # PyMuPDF
from deskew import determine_skew

def img_pdf_check(path):
    """Return OpenCV image from image file or first page of PDF."""
    if not isinstance(path, str) or not os.path.exists(path):
        return {"status": "Failed", "message": "Invalid or missing file"}

    ext = os.path.splitext(path)[1].lower()

    # ---- Image ----
    if ext in {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp"}:
        img = cv2.imread(path)
        return img if img is not None else {
            "status": "Failed", "message": "Invalid image file"
        }

    # ---- PDF (first page only) ----
    if ext == ".pdf":
        try:
            with fitz.open(path) as doc:
                pix = doc[0].get_pixmap(dpi=300)

            img = np.frombuffer(pix.samples, np.uint8).reshape(pix.height, pix.width, pix.n)
            if pix.n == 3:
                img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            elif pix.n == 4:
                img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGR)

            return img

        except Exception as e:
            return {"status": "Failed", "message": f"PDF error: {e}"}

    return {"status": "Failed", "message": "Unsupported file type"}

def ensure_grayscale(img):
    """If image has 3 channels (color), convert to grayscale."""
    if len(img.shape) == 3 and img.shape[2] == 3:
        return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return img

def resize_image(image, max_dim):
    """Resize image while keeping aspect ratio."""
    height, width = image.shape[:2]
    if max(height, width) <= max_dim:
        return image

    scale = max_dim / max(height, width)
    new_width = max(1, int(width * scale))
    new_height = max(1, int(height * scale))

    return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)

def rotate_image(image, angle):
    """Rotate image CLOCKWISE by the given angle (0, 90, 180, 270)."""
    if angle == 0:
        return image
    elif angle == 90:
        return cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
    elif angle == 180:
        return cv2.rotate(image, cv2.ROTATE_180)
    elif angle == 270:
        return cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
    return image

def tilt_image(gray_img):
    """Deskew the image using deskew library."""
    angle = determine_skew(gray_img)
    (h, w) = gray_img.shape[:2]
    center = (w // 2, h // 2)

    M = cv2.getRotationMatrix2D(center, angle, 1.0)
    corrected = cv2.warpAffine(
        gray_img,
        M,
        (w, h),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE
    )
    return corrected

def resize_v2(img, max_height=1000, max_width=600, scale=2):
    """Resize function for scanned copy."""
    h, w = img.shape[:2]
    if h < max_height and w < max_width:
        return cv2.resize(img, None, fx=scale, fy=scale, interpolation=cv2.INTER_LANCZOS4)
    else:
        return cv2.resize(img, None, fx=0.5, fy=0.5, interpolation=cv2.INTER_LANCZOS4)