import os
import cv2
import numpy as np
import gradio as gr

import zipfile
import tempfile
import shutil
from pathlib import Path

from fastapi import FastAPI, UploadFile, File
from fastapi.responses import Response, JSONResponse


# ============================================================
# SETTINGS
# ============================================================

CARD_ASPECT_RATIO = 85.60 / 53.98

MIN_CARD_AREA_RATIO = 0.015
MAX_CARD_AREA_RATIO = 0.92

CARD_MARGIN = 0.030

MAX_DIMENSION = 1600

MIN_OUTPUT_WIDTH = 100
MIN_OUTPUT_HEIGHT = 60

MAX_OUTPUT_WIDTH = 2500
MAX_OUTPUT_HEIGHT = 1600

MAX_CANDIDATES = 2500
MAX_FALLBACK_CONTOURS = 100

EARLY_ACCEPT_SCORE = 0.91

FAST_PATH_ENABLED = True
COLOR_PATH_ENABLED = True

FAST_ACCEPT_SCORE = 0.60
FAST_MAX_DIMENSION = 1200

# ------------------------------------------------------------
# Brightness
# ------------------------------------------------------------
# Small default increase.
#
# 0   = original
# 10  = slightly brighter
# 20  = brighter
# 30  = strong
#
# ZIP processing uses this value.
# ------------------------------------------------------------

DEFAULT_BRIGHTNESS = 10


# ============================================================
# FASTAPI
# ============================================================

app = FastAPI(
    title="Egyptian ID Card Cropper"
)


# ============================================================
# POINT HELPERS
# ============================================================

def distance(a, b):
    return float(np.linalg.norm(a - b))


def order_points(points):

    pts = np.asarray(
        points,
        dtype=np.float32
    ).reshape(4, 2)

    center = np.mean(
        pts,
        axis=0
    )

    angles = np.arctan2(
        pts[:, 1] - center[1],
        pts[:, 0] - center[0]
    )

    pts = pts[np.argsort(angles)]

    sums = (
        pts[:, 0] +
        pts[:, 1]
    )

    tl_index = np.argmin(sums)

    pts = np.roll(
        pts,
        -tl_index,
        axis=0
    )

    remaining = pts[1:]

    br_index = np.argmax(
        remaining[:, 0] +
        remaining[:, 1]
    )

    br = remaining[br_index]

    others = [
        p
        for i, p in enumerate(remaining)
        if i != br_index
    ]

    others = sorted(
        others,
        key=lambda p: p[1]
    )

    tr = others[0]
    bl = others[1]

    tl = pts[0]

    return np.array(
        [
            tl,
            tr,
            br,
            bl
        ],
        dtype=np.float32
    )


def polygon_angle(a, b, c):

    ba = a - b
    bc = c - b

    denominator = (
        np.linalg.norm(ba) *
        np.linalg.norm(bc)
    )

    if denominator <= 1e-8:
        return 0.0

    cosine = (
        np.dot(ba, bc) /
        denominator
    )

    cosine = np.clip(
        cosine,
        -1.0,
        1.0
    )

    return float(
        np.degrees(
            np.arccos(cosine)
        )
    )


# ============================================================
# BORDER CHECK
# ============================================================

def touches_image_border(
    quad,
    image_shape,
    tolerance=3
):

    h, w = image_shape[:2]

    points = order_points(quad)

    for x, y in points:

        if x <= tolerance:
            return True

        if y <= tolerance:
            return True

        if x >= w - 1 - tolerance:
            return True

        if y >= h - 1 - tolerance:
            return True

    return False


# ============================================================
# QUADRILATERAL GEOMETRY
# ============================================================

def quad_geometry(quad):

    ordered = order_points(quad)

    tl, tr, br, bl = ordered

    width_top = distance(tl, tr)
    width_bottom = distance(bl, br)

    height_left = distance(tl, bl)
    height_right = distance(tr, br)

    width = (
        width_top +
        width_bottom
    ) / 2.0

    height = (
        height_left +
        height_right
    ) / 2.0

    if width <= 0 or height <= 0:
        return None

    ratio = (
        max(width, height) /
        min(width, height)
    )

    return {
        "points": ordered,
        "width": width,
        "height": height,
        "ratio": ratio,
        "area": cv2.contourArea(ordered)
    }


def is_reasonable_card(
    quad,
    image_shape
):

    h, w = image_shape[:2]

    image_area = float(h * w)

    geometry = quad_geometry(quad)

    if geometry is None:
        return False

    area = geometry["area"]

    area_ratio = area / image_area

    if area_ratio < MIN_CARD_AREA_RATIO:
        return False

    if area_ratio > MAX_CARD_AREA_RATIO:
        return False

    ratio = geometry["ratio"]

    if ratio < 1.15:
        return False

    if ratio > 2.30:
        return False

    points = geometry["points"]

    if touches_image_border(
        points,
        image_shape,
        tolerance=3
    ):
        return False

    tl, tr, br, bl = points

    angles = [
        polygon_angle(tl, tr, br),
        polygon_angle(tr, br, bl),
        polygon_angle(br, bl, tl),
        polygon_angle(bl, tl, tr)
    ]

    for angle in angles:

        if angle < 35:
            return False

        if angle > 145:
            return False

    return True


# ============================================================
# SCORE CARD
# ============================================================

def score_card(
    quad,
    contour_area,
    image_area
):

    geometry = quad_geometry(quad)

    if geometry is None:
        return -1

    ratio = geometry["ratio"]

    ratio_error = abs(
        ratio -
        CARD_ASPECT_RATIO
    )

    ratio_score = max(
        0.0,
        1.0 -
        ratio_error / 0.70
    )

    area_ratio = (
        contour_area /
        image_area
    )

    area_score = min(
        area_ratio / 0.30,
        1.0
    )

    points = geometry["points"]

    tl, tr, br, bl = points

    angles = [
        polygon_angle(tl, tr, br),
        polygon_angle(tr, br, bl),
        polygon_angle(br, bl, tl),
        polygon_angle(bl, tl, tr)
    ]

    angle_error = np.mean([
        abs(angle - 90.0)
        for angle in angles
    ])

    angle_score = max(
        0.0,
        1.0 -
        angle_error / 50.0
    )

    score = (
        ratio_score * 0.55 +
        area_score * 0.25 +
        angle_score * 0.20
    )

    return float(score)


# ============================================================
# FAST DETECTION
# ============================================================

def fast_card_detection(image):

    h, w = image.shape[:2]

    scale = 1.0

    if max(h, w) > FAST_MAX_DIMENSION:

        scale = (
            FAST_MAX_DIMENSION /
            float(max(h, w))
        )

        work = cv2.resize(
            image,
            None,
            fx=scale,
            fy=scale,
            interpolation=cv2.INTER_AREA
        )

    else:
        work = image.copy()

    gray = cv2.cvtColor(
        work,
        cv2.COLOR_BGR2GRAY
    )

    gray = cv2.GaussianBlur(
        gray,
        (3, 3),
        0
    )

    candidates = []

    for low, high in [
        (35, 110),
        (50, 150),
        (70, 180)
    ]:

        edges = cv2.Canny(
            gray,
            low,
            high
        )

        kernel = cv2.getStructuringElement(
            cv2.MORPH_RECT,
            (5, 5)
        )

        edges = cv2.morphologyEx(
            edges,
            cv2.MORPH_CLOSE,
            kernel,
            iterations=1
        )

        contours, _ = cv2.findContours(
            edges,
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE
        )

        candidates.extend(contours)

    _, binary = cv2.threshold(
        gray,
        0,
        255,
        cv2.THRESH_BINARY +
        cv2.THRESH_OTSU
    )

    contours, _ = cv2.findContours(
        binary,
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE
    )

    candidates.extend(contours)

    inverted = cv2.bitwise_not(binary)

    contours, _ = cv2.findContours(
        inverted,
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE
    )

    candidates.extend(contours)

    image_area = float(
        work.shape[0] *
        work.shape[1]
    )

    candidates = sorted(
        candidates,
        key=cv2.contourArea,
        reverse=True
    )[:80]

    best_quad = None
    best_score = -1

    for contour in candidates:

        area = cv2.contourArea(contour)

        if area <= 0:
            continue

        area_ratio = (
            area /
            image_area
        )

        if area_ratio < MIN_CARD_AREA_RATIO:
            continue

        if area_ratio > MAX_CARD_AREA_RATIO:
            continue

        perimeter = cv2.arcLength(
            contour,
            True
        )

        if perimeter <= 0:
            continue

        for epsilon_factor in [
            0.012,
            0.020,
            0.030
        ]:

            approx = cv2.approxPolyDP(
                contour,
                epsilon_factor * perimeter,
                True
            )

            if len(approx) != 4:
                continue

            quad = (
                approx
                .reshape(4, 2)
                .astype(np.float32)
            )

            if not is_reasonable_card(
                quad,
                work.shape
            ):
                continue

            score = score_card(
                quad,
                area,
                image_area
            )

            if score > best_score:

                best_score = score
                best_quad = quad.copy()

                if best_score >= EARLY_ACCEPT_SCORE:
                    break

        if best_score >= EARLY_ACCEPT_SCORE:
            break

    if (
        best_quad is None or
        best_score < FAST_ACCEPT_SCORE
    ):
        return None

    if scale != 1.0:
        best_quad = best_quad / scale

    best_quad[:, 0] = np.clip(
        best_quad[:, 0],
        0,
        w - 1
    )

    best_quad[:, 1] = np.clip(
        best_quad[:, 1],
        0,
        h - 1
    )

    return order_points(best_quad)


# ============================================================
# COLOR DETECTION
# ============================================================

def color_card_detection(image):

    h, w = image.shape[:2]

    scale = 1.0

    if max(h, w) > FAST_MAX_DIMENSION:

        scale = (
            FAST_MAX_DIMENSION /
            float(max(h, w))
        )

        work = cv2.resize(
            image,
            None,
            fx=scale,
            fy=scale,
            interpolation=cv2.INTER_AREA
        )

    else:
        work = image.copy()

    hsv = cv2.cvtColor(
        work,
        cv2.COLOR_BGR2HSV
    )

    lower = np.array(
        [0, 0, 70],
        dtype=np.uint8
    )

    upper = np.array(
        [179, 150, 255],
        dtype=np.uint8
    )

    mask = cv2.inRange(
        hsv,
        lower,
        upper
    )

    kernel = cv2.getStructuringElement(
        cv2.MORPH_RECT,
        (9, 9)
    )

    mask = cv2.morphologyEx(
        mask,
        cv2.MORPH_CLOSE,
        kernel,
        iterations=2
    )

    mask = cv2.morphologyEx(
        mask,
        cv2.MORPH_OPEN,
        kernel,
        iterations=1
    )

    contours, _ = cv2.findContours(
        mask,
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE
    )

    image_area = float(
        work.shape[0] *
        work.shape[1]
    )

    contours = sorted(
        contours,
        key=cv2.contourArea,
        reverse=True
    )[:50]

    best_quad = None
    best_score = -1

    for contour in contours:

        area = cv2.contourArea(contour)

        if area <= 0:
            continue

        area_ratio = area / image_area

        if area_ratio < MIN_CARD_AREA_RATIO:
            continue

        if area_ratio > MAX_CARD_AREA_RATIO:
            continue

        perimeter = cv2.arcLength(
            contour,
            True
        )

        if perimeter <= 0:
            continue

        for epsilon in [
            0.01,
            0.02,
            0.03,
            0.04
        ]:

            approx = cv2.approxPolyDP(
                contour,
                epsilon * perimeter,
                True
            )

            if len(approx) != 4:
                continue

            quad = (
                approx
                .reshape(4, 2)
                .astype(np.float32)
            )

            if not is_reasonable_card(
                quad,
                work.shape
            ):
                continue

            score = score_card(
                quad,
                area,
                image_area
            )

            if score > best_score:

                best_score = score
                best_quad = quad.copy()

                if best_score >= EARLY_ACCEPT_SCORE:
                    break

        if best_score >= EARLY_ACCEPT_SCORE:
            break

    if (
        best_quad is None or
        best_score < FAST_ACCEPT_SCORE
    ):
        return None

    if scale != 1.0:
        best_quad = best_quad / scale

    best_quad[:, 0] = np.clip(
        best_quad[:, 0],
        0,
        w - 1
    )

    best_quad[:, 1] = np.clip(
        best_quad[:, 1],
        0,
        h - 1
    )

    return order_points(best_quad)


# ============================================================
# EDGE CONTOURS
# ============================================================

def get_contours_from_edges(gray):

    candidates = []

    blur = cv2.GaussianBlur(
        gray,
        (5, 5),
        0
    )

    canny_settings = [
        (20, 80),
        (30, 100),
        (40, 120),
        (50, 150),
        (70, 180),
        (90, 220),
        (110, 240)
    ]

    for low, high in canny_settings:

        edges = cv2.Canny(
            blur,
            low,
            high
        )

        for kernel_size in [
            3,
            5,
            7,
            9
        ]:

            kernel = cv2.getStructuringElement(
                cv2.MORPH_RECT,
                (
                    kernel_size,
                    kernel_size
                )
            )

            closed = cv2.morphologyEx(
                edges,
                cv2.MORPH_CLOSE,
                kernel,
                iterations=1
            )

            contours, _ = cv2.findContours(
                closed,
                cv2.RETR_LIST,
                cv2.CHAIN_APPROX_SIMPLE
            )

            candidates.extend(contours)

    return candidates


# ============================================================
# THRESHOLD CONTOURS
# ============================================================

def get_threshold_contours(gray):

    candidates = []

    blur = cv2.GaussianBlur(
        gray,
        (5, 5),
        0
    )

    _, binary = cv2.threshold(
        blur,
        0,
        255,
        cv2.THRESH_BINARY +
        cv2.THRESH_OTSU
    )

    for image in [
        binary,
        cv2.bitwise_not(binary)
    ]:

        kernel = cv2.getStructuringElement(
            cv2.MORPH_RECT,
            (5, 5)
        )

        image = cv2.morphologyEx(
            image,
            cv2.MORPH_CLOSE,
            kernel,
            iterations=2
        )

        contours, _ = cv2.findContours(
            image,
            cv2.RETR_LIST,
            cv2.CHAIN_APPROX_SIMPLE
        )

        candidates.extend(contours)

    for block_size, c in [
        (21, 5),
        (31, 7),
        (41, 9),
        (51, 11)
    ]:

        adaptive = cv2.adaptiveThreshold(
            blur,
            255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY,
            block_size,
            c
        )

        for image in [
            adaptive,
            cv2.bitwise_not(adaptive)
        ]:

            contours, _ = cv2.findContours(
                image,
                cv2.RETR_LIST,
                cv2.CHAIN_APPROX_SIMPLE
            )

            candidates.extend(contours)

    return candidates


# ============================================================
# FIND CARD
# ============================================================

def find_card_contour(image):

    original = image.copy()

    # ========================================================
    # FAST
    # ========================================================

    if FAST_PATH_ENABLED:

        try:

            fast_quad = fast_card_detection(
                original
            )

            if fast_quad is not None:
                return fast_quad

        except Exception as e:

            print(
                "FAST PATH ERROR:",
                repr(e)
            )

    # ========================================================
    # COLOR
    # ========================================================

    if COLOR_PATH_ENABLED:

        try:

            color_quad = color_card_detection(
                original
            )

            if color_quad is not None:
                return color_quad

        except Exception as e:

            print(
                "COLOR PATH ERROR:",
                repr(e)
            )

    # ========================================================
    # ORIGINAL
    # ========================================================

    h, w = original.shape[:2]

    scale = 1.0

    if max(h, w) > MAX_DIMENSION:

        scale = (
            MAX_DIMENSION /
            float(max(h, w))
        )

        work = cv2.resize(
            original,
            None,
            fx=scale,
            fy=scale,
            interpolation=cv2.INTER_AREA
        )

    else:

        work = original.copy()

    gray = cv2.cvtColor(
        work,
        cv2.COLOR_BGR2GRAY
    )

    clahe = cv2.createCLAHE(
        clipLimit=2.0,
        tileGridSize=(8, 8)
    )

    enhanced = clahe.apply(gray)

    candidates = []

    candidates.extend(
        get_contours_from_edges(
            enhanced
        )
    )

    candidates.extend(
        get_threshold_contours(
            enhanced
        )
    )

    if len(candidates) > MAX_CANDIDATES:

        candidates = sorted(
            candidates,
            key=cv2.contourArea,
            reverse=True
        )[:MAX_CANDIDATES]

    image_area = float(
        work.shape[0] *
        work.shape[1]
    )

    best_quad = None
    best_score = -1

    for contour in candidates:

        contour_area = cv2.contourArea(
            contour
        )

        if contour_area <= 0:
            continue

        if contour_area < (
            image_area *
            MIN_CARD_AREA_RATIO
        ):
            continue

        if contour_area > (
            image_area *
            MAX_CARD_AREA_RATIO
        ):
            continue

        perimeter = cv2.arcLength(
            contour,
            True
        )

        if perimeter <= 0:
            continue

        for epsilon_factor in [
            0.008,
            0.010,
            0.012,
            0.015,
            0.018,
            0.020,
            0.025,
            0.030,
            0.035,
            0.040,
            0.050
        ]:

            approx = cv2.approxPolyDP(
                contour,
                epsilon_factor * perimeter,
                True
            )

            if len(approx) != 4:
                continue

            quad = (
                approx
                .reshape(4, 2)
                .astype(np.float32)
            )

            if not is_reasonable_card(
                quad,
                work.shape
            ):
                continue

            score = score_card(
                quad,
                contour_area,
                image_area
            )

            if score > best_score:

                best_score = score
                best_quad = quad.copy()

                if best_score >= EARLY_ACCEPT_SCORE:
                    break

        if best_score >= EARLY_ACCEPT_SCORE:
            break

    # ========================================================
    # FALLBACK
    # ========================================================

    if best_quad is None:

        sorted_candidates = sorted(
            candidates,
            key=cv2.contourArea,
            reverse=True
        )

        for contour in sorted_candidates[
            :MAX_FALLBACK_CONTOURS
        ]:

            contour_area = cv2.contourArea(
                contour
            )

            if contour_area <= (
                image_area *
                MIN_CARD_AREA_RATIO
            ):
                continue

            rect = cv2.minAreaRect(
                contour
            )

            box = cv2.boxPoints(rect)

            box = np.asarray(
                box,
                dtype=np.float32
            )

            if not is_reasonable_card(
                box,
                work.shape
            ):
                continue

            score = score_card(
                box,
                contour_area,
                image_area
            )

            score *= 0.92

            if score > best_score:

                best_score = score
                best_quad = box.copy()

                if best_score >= EARLY_ACCEPT_SCORE:
                    break

    if best_quad is None:
        return None

    if scale != 1.0:
        best_quad = best_quad / scale

    best_quad[:, 0] = np.clip(
        best_quad[:, 0],
        0,
        original.shape[1] - 1
    )

    best_quad[:, 1] = np.clip(
        best_quad[:, 1],
        0,
        original.shape[0] - 1
    )

    return order_points(best_quad)


# ============================================================
# EXPAND CARD
# ============================================================

def expand_quad(
    corners,
    image_shape,
    margin=CARD_MARGIN
):

    h, w = image_shape[:2]

    corners = order_points(corners)

    tl, tr, br, bl = corners

    width_top = distance(tl, tr)
    width_bottom = distance(bl, br)

    height_left = distance(tl, bl)
    height_right = distance(tr, br)

    avg_width = (
        width_top +
        width_bottom
    ) / 2.0

    avg_height = (
        height_left +
        height_right
    ) / 2.0

    pad_x = avg_width * margin
    pad_y = avg_height * margin

    center = np.mean(
        corners,
        axis=0
    )

    expanded = []

    for point in corners:

        direction = point - center

        norm = np.linalg.norm(direction)

        if norm > 0:

            amount = (
                margin *
                0.75 *
                norm
            )

            new_point = (
                point +
                direction / norm *
                amount
            )

        else:

            new_point = point

        expanded.append(new_point)

    expanded = np.asarray(
        expanded,
        dtype=np.float32
    )

    expanded[0] += np.array(
        [
            -pad_x * 0.15,
            -pad_y * 0.15
        ],
        dtype=np.float32
    )

    expanded[1] += np.array(
        [
            pad_x * 0.15,
            -pad_y * 0.15
        ],
        dtype=np.float32
    )

    expanded[2] += np.array(
        [
            pad_x * 0.15,
            pad_y * 0.15
        ],
        dtype=np.float32
    )

    expanded[3] += np.array(
        [
            -pad_x * 0.15,
            pad_y * 0.15
        ],
        dtype=np.float32
    )

    expanded[:, 0] = np.clip(
        expanded[:, 0],
        0,
        w - 1
    )

    expanded[:, 1] = np.clip(
        expanded[:, 1],
        0,
        h - 1
    )

    return order_points(expanded)


# ============================================================
# PERSPECTIVE CROP
# ============================================================

def perspective_crop(
    image,
    corners
):

    corners = order_points(corners)

    tl, tr, br, bl = corners

    width_top = distance(tl, tr)
    width_bottom = distance(bl, br)

    height_left = distance(tl, bl)
    height_right = distance(tr, br)

    output_width = int(
        round(
            max(
                width_top,
                width_bottom
            )
        )
    )

    output_height = int(
        round(
            max(
                height_left,
                height_right
            )
        )
    )

    if output_width < MIN_OUTPUT_WIDTH:
        return None

    if output_height < MIN_OUTPUT_HEIGHT:
        return None

    output_width = min(
        output_width,
        MAX_OUTPUT_WIDTH
    )

    output_height = min(
        output_height,
        MAX_OUTPUT_HEIGHT
    )

    destination = np.array(
        [
            [0, 0],
            [output_width - 1, 0],
            [
                output_width - 1,
                output_height - 1
            ],
            [
                0,
                output_height - 1
            ]
        ],
        dtype=np.float32
    )

    source = np.array(
        [
            tl,
            tr,
            br,
            bl
        ],
        dtype=np.float32
    )

    matrix = cv2.getPerspectiveTransform(
        source,
        destination
    )

    cropped = cv2.warpPerspective(
        image,
        matrix,
        (
            output_width,
            output_height
        ),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE
    )

    if cropped is None or cropped.size == 0:
        return None

    return cropped


# ============================================================
# ORIENTATION
# ============================================================

def normalize_card_orientation(image):

    if image is None:
        return None

    h, w = image.shape[:2]

    if h <= 0 or w <= 0:
        return image

    # IMPORTANT:
    # We only rotate if the resulting image is clearly portrait.
    # No flip is ever performed.
    if h > w:
        image = cv2.rotate(
            image,
            cv2.ROTATE_90_CLOCKWISE
        )

    return image


# ============================================================
# BRIGHTNESS
# ============================================================

def adjust_brightness(
    image,
    brightness
):

    if image is None:
        return None

    try:
        brightness = float(brightness)
    except Exception:
        brightness = 0

    brightness = max(
        -100,
        min(100, brightness)
    )

    if abs(brightness) < 0.01:
        return image.copy()

    # --------------------------------------------------------
    # ONLY brightness.
    #
    # No sharpening.
    # No saturation modification.
    # No contrast modification.
    # No resizing.
    # No flipping.
    #
    # This preserves the original details much better.
    # --------------------------------------------------------

    result = cv2.convertScaleAbs(
        image,
        alpha=1.0,
        beta=brightness
    )

    return result


# ============================================================
# CROP
# ============================================================

def crop_id_card(
    image,
    brightness=0
):

    if image is None:
        return None

    if len(image.shape) == 2:

        image = cv2.cvtColor(
            image,
            cv2.COLOR_GRAY2BGR
        )

    if (
        len(image.shape) == 3 and
        image.shape[2] == 4
    ):

        image = cv2.cvtColor(
            image,
            cv2.COLOR_BGRA2BGR
        )

    original_h, original_w = image.shape[:2]

    original_area = (
        original_h *
        original_w
    )

    corners = find_card_contour(
        image
    )

    if corners is None:
        return None

    corners = expand_quad(
        corners,
        image.shape,
        CARD_MARGIN
    )

    cropped = perspective_crop(
        image,
        corners
    )

    if cropped is None:
        return None

    x, y, cw, ch = cv2.boundingRect(
        corners.astype(np.float32)
    )

    bbox_area = cw * ch

    if bbox_area >= (
        original_area * 0.97
    ):

        # The image is probably already cropped.
        # Keep the original rather than returning failure.
        cropped = image.copy()

    cropped = normalize_card_orientation(
        cropped
    )

    cropped = adjust_brightness(
        cropped,
        brightness
    )

    return cropped


# ============================================================
# IMAGE DECODING
# ============================================================

def read_image_file(path):

    image = cv2.imread(
        str(path),
        cv2.IMREAD_COLOR
    )

    return image


def decode_bytes(data):

    if not data:
        return None

    npimg = np.frombuffer(
        data,
        dtype=np.uint8
    )

    return cv2.imdecode(
        npimg,
        cv2.IMREAD_COLOR
    )


# ============================================================
# IMAGE ENCODING
# ============================================================

def get_output_extension(
    original_name
):

    ext = Path(
        original_name
    ).suffix.lower()

    # JPEG output for JPEG input.
    if ext in [
        ".jpg",
        ".jpeg"
    ]:
        return ".jpg"

    # PNG stays PNG.
    if ext == ".png":
        return ".png"

    # Other formats are converted to JPEG.
    return ".jpg"


def encode_image(
    image,
    extension
):

    extension = extension.lower()

    if extension in [
        ".jpg",
        ".jpeg"
    ]:

        ok, buffer = cv2.imencode(
            ".jpg",
            image,
            [
                int(
                    cv2.IMWRITE_JPEG_QUALITY
                ),
                95
            ]
        )

    elif extension == ".png":

        ok, buffer = cv2.imencode(
            ".png",
            image,
            [
                int(
                    cv2.IMWRITE_PNG_COMPRESSION
                ),
                3
            ]
        )

    else:

        ok, buffer = cv2.imencode(
            ".jpg",
            image,
            [
                int(
                    cv2.IMWRITE_JPEG_QUALITY
                ),
                95
            ]
        )

    if not ok:
        return None

    return buffer.tobytes()


# ============================================================
# GRADIO IMAGE READER
# ============================================================

def read_gradio_image(file):

    if file is None:
        return None

    if isinstance(file, str):

        return cv2.imread(
            file,
            cv2.IMREAD_COLOR
        )

    if isinstance(file, np.ndarray):

        image = file.copy()

        if len(image.shape) == 3:

            image = cv2.cvtColor(
                image,
                cv2.COLOR_RGB2BGR
            )

        return image

    if hasattr(file, "read"):

        data = file.read()

    else:

        data = file

    if not isinstance(data, bytes):
        return None

    return decode_bytes(data)


# ============================================================
# SINGLE IMAGE PROCESSOR
# ============================================================

def process_image(
    file,
    brightness,
    rotation
):

    try:

        if file is None:

            return (
                None,
                "Please upload an ID image."
            )

        image = read_gradio_image(
            file
        )

        if image is None:

            return (
                None,
                "Could not decode image."
            )

        cropped = crop_id_card(
            image,
            brightness
        )

        if cropped is None:

            return (
                None,
                "❌ ID card was not detected."
            )

        # ----------------------------------------------------
        # USER ROTATION
        # ----------------------------------------------------

        rotation = int(rotation)

        if rotation == 90:

            cropped = cv2.rotate(
                cropped,
                cv2.ROTATE_90_CLOCKWISE
            )

        elif rotation == 180:

            cropped = cv2.rotate(
                cropped,
                cv2.ROTATE_180
            )

        elif rotation == 270:

            cropped = cv2.rotate(
                cropped,
                cv2.ROTATE_90_COUNTERCLOCKWISE
            )

        cropped_rgb = cv2.cvtColor(
            cropped,
            cv2.COLOR_BGR2RGB
        )

        return (
            cropped_rgb,
            "✅ ID card cropped successfully."
        )

    except Exception as e:

        print(
            "PROCESS ERROR:",
            repr(e)
        )

        return (
            None,
            "❌ Error: " + str(e)
        )


# ============================================================
# ZIP PROCESSING
# ============================================================

def process_zip(
    zip_file,
    brightness
):

    if zip_file is None:

        return (
            None,
            "Please upload a ZIP file."
        )

    work_dir = tempfile.mkdtemp(
        prefix="id_crop_"
    )

    input_dir = os.path.join(
        work_dir,
        "input"
    )

    output_dir = os.path.join(
        work_dir,
        "output"
    )

    os.makedirs(input_dir)
    os.makedirs(output_dir)

    try:

        # ----------------------------------------------------
        # Get ZIP path
        # ----------------------------------------------------

        if isinstance(zip_file, str):

            zip_path = zip_file

        elif hasattr(zip_file, "name"):

            zip_path = zip_file.name

        else:

            return (
                None,
                "Invalid ZIP file."
            )

        # ----------------------------------------------------
        # Extract ZIP
        # ----------------------------------------------------

        with zipfile.ZipFile(
            zip_path,
            "r"
        ) as z:

            z.extractall(
                input_dir
            )

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

        image_files = []

        for root, dirs, files in os.walk(
            input_dir
        ):

            for filename in files:

                path = Path(
                    root
                ) / filename

                if path.suffix.lower() in supported:

                    image_files.append(
                        path
                    )

        if not image_files:

            return (
                None,
                "No supported images were found in the ZIP."
            )

        processed = 0
        failed = 0

        # ----------------------------------------------------
        # Process every image
        # ----------------------------------------------------

        for source_path in image_files:

            try:

                image = read_image_file(
                    source_path
                )

                if image is None:

                    failed += 1
                    continue

                cropped = crop_id_card(
                    image,
                    brightness
                )

                if cropped is None:

                    # If detection fails, keep original.
                    # This is important because the user requested
                    # that already-cropped / difficult images
                    # should not disappear.
                    cropped = image.copy()

                    cropped = adjust_brightness(
                        cropped,
                        brightness
                    )

                # ------------------------------------------------
                # Preserve directory structure.
                # ------------------------------------------------

                relative = source_path.relative_to(
                    input_dir
                )

                relative_parent = relative.parent

                output_parent = (
                    Path(output_dir) /
                    relative_parent
                )

                output_parent.mkdir(
                    parents=True,
                    exist_ok=True
                )

                # ------------------------------------------------
                # IMPORTANT:
                # Never save ZIP results as WEBP.
                #
                # JPEG -> JPG
                # PNG  -> PNG
                # Everything else -> JPG
                # ------------------------------------------------

                output_extension = get_output_extension(
                    source_path.name
                )

                output_name = (
                    source_path.stem +
                    "_cropped" +
                    output_extension
                )

                output_path = (
                    output_parent /
                    output_name
                )

                encoded = encode_image(
                    cropped,
                    output_extension
                )

                if encoded is None:

                    failed += 1
                    continue

                with open(
                    output_path,
                    "wb"
                ) as f:

                    f.write(encoded)

                processed += 1

            except Exception as e:

                failed += 1

                print(
                    "ZIP IMAGE ERROR:",
                    source_path,
                    repr(e)
                )

        # ----------------------------------------------------
        # Create output ZIP
        # ----------------------------------------------------

        output_zip_base = os.path.join(
            work_dir,
            "cropped_images"
        )

        output_zip = shutil.make_archive(
            output_zip_base,
            "zip",
            output_dir
        )

        message = (
            f"✅ ZIP processing complete. "
            f"{processed} images processed."
        )

        if failed:
            message += (
                f" {failed} images could not be "
                f"processed and were skipped."
            )

        return (
            output_zip,
            message
        )

    except zipfile.BadZipFile:

        return (
            None,
            "❌ Invalid ZIP file."
        )

    except Exception as e:

        print(
            "ZIP PROCESS ERROR:",
            repr(e)
        )

        return (
            None,
            "❌ ZIP error: " + str(e)
        )

    finally:

        # --------------------------------------------------------
        # Do not delete output ZIP here.
        #
        # Gradio needs the file to remain available.
        #
        # The temporary directory will be cleaned by the
        # operating system / environment.
        # --------------------------------------------------------

        pass


# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks(
    title="Egyptian ID Card Cropper"
) as interface:

    gr.Markdown(
        """
        # 🇪🇬 Egyptian ID Card Cropper

        Upload an Egyptian ID photograph or a ZIP containing
        multiple images.

        ### Features

        - Detects the physical ID card
        - Handles perspective
        - Handles tilted cards
        - Keeps already-cropped IDs
        - Removes background around the ID
        - Does not flip the image
        - Does not sharpen or destroy details
        - Brightness can be controlled
        - Rotation can be controlled
        - ZIP results are saved as JPEG/PNG, never WebP
        """
    )

    # ========================================================
    # SINGLE IMAGE
    # ========================================================

    gr.Markdown(
        "## Single Image"
    )

    input_image = gr.Image(
        type="filepath",
        label="Upload ID Image"
    )

    process_button = gr.Button(
        "Crop ID Card",
        variant="primary"
    )

    # --------------------------------------------------------
    # NEW CROPPED IMAGE
    # --------------------------------------------------------

    output_image = gr.Image(
        type="numpy",
        label="New Cropped ID Card"
    )

    # --------------------------------------------------------
    # CONTROLS UNDER CROPPED IMAGE
    # --------------------------------------------------------

    gr.Markdown(
        "### Image Controls"
    )

    brightness_slider = gr.Slider(
        minimum=-50,
        maximum=50,
        value=10,
        step=1,
        label="Brightness",
        info="0 = original. Positive values make the ID brighter."
    )

    rotation_dropdown = gr.Dropdown(
        choices=[
            0,
            90,
            180,
            270
        ],
        value=0,
        label="Rotate Image",
        info="Rotation is applied clockwise."
    )

    status = gr.Textbox(
        label="Status",
        interactive=False
    )

    process_button.click(
        fn=process_image,
        inputs=[
            input_image,
            brightness_slider,
            rotation_dropdown
        ],
        outputs=[
            output_image,
            status
        ]
    )

    input_image.change(
        fn=process_image,
        inputs=[
            input_image,
            brightness_slider,
            rotation_dropdown
        ],
        outputs=[
            output_image,
            status
        ]
    )

    # --------------------------------------------------------
    # Re-process when controls change
    # --------------------------------------------------------

    brightness_slider.change(
        fn=process_image,
        inputs=[
            input_image,
            brightness_slider,
            rotation_dropdown
        ],
        outputs=[
            output_image,
            status
        ]
    )

    rotation_dropdown.change(
        fn=process_image,
        inputs=[
            input_image,
            brightness_slider,
            rotation_dropdown
        ],
        outputs=[
            output_image,
            status
        ]
    )

    # ========================================================
    # ZIP
    # ========================================================

    gr.Markdown(
        """
        ---
        ## 📦 Process a ZIP

        Upload a ZIP containing your ID images.

        The application will process every image and create:

        **cropped_images.zip**

        JPEG images remain JPEG, PNG images remain PNG,
        and unsupported image formats are converted to JPEG.
        """
    )

    zip_input = gr.File(
        type="filepath",
        label="Upload ZIP File",
        file_types=[".zip"]
    )

    zip_brightness = gr.Slider(
        minimum=0,
        maximum=50,
        value=DEFAULT_BRIGHTNESS,
        step=1,
        label="ZIP Brightness",
        info="All images in the resulting ZIP receive this brightness increase."
    )

    zip_button = gr.Button(
        "Process ZIP",
        variant="primary"
    )

    zip_output = gr.File(
        label="Download Cropped ZIP"
    )

    zip_status = gr.Textbox(
        label="ZIP Status",
        interactive=False
    )

    zip_button.click(
        fn=process_zip,
        inputs=[
            zip_input,
            zip_brightness
        ],
        outputs=[
            zip_output,
            zip_status
        ]
    )


# ============================================================
# MOUNT GRADIO
# ============================================================

app = gr.mount_gradio_app(
    app,
    interface,
    path="/"
)


# ============================================================
# API - SINGLE IMAGE
# ============================================================

@app.post(
    "/crop-id"
)
async def crop_id_endpoint(
    file: UploadFile = File(...)
):

    try:

        contents = await file.read()

        if not contents:

            return JSONResponse(
                status_code=400,
                content={
                    "success": False,
                    "error": "Empty file"
                }
            )

        image = decode_bytes(
            contents
        )

        if image is None:

            return JSONResponse(
                status_code=400,
                content={
                    "success": False,
                    "error": "Invalid image"
                }
            )

        cropped = crop_id_card(
            image,
            DEFAULT_BRIGHTNESS
        )

        if cropped is None:

            return JSONResponse(
                status_code=422,
                content={
                    "success": False,
                    "error":
                        "ID card could not be detected"
                }
            )

        ok, buffer = cv2.imencode(
            ".jpg",
            cropped,
            [
                int(
                    cv2.IMWRITE_JPEG_QUALITY
                ),
                95
            ]
        )

        if not ok:

            return JSONResponse(
                status_code=500,
                content={
                    "success": False,
                    "error":
                        "Could not encode cropped image"
                }
            )

        return Response(
            content=buffer.tobytes(),
            media_type="image/jpeg"
        )

    except Exception as e:

        print(
            "API ERROR:",
            repr(e)
        )

        return JSONResponse(
            status_code=500,
            content={
                "success": False,
                "error": str(e)
            }
        )


# ============================================================
# HEALTH
# ============================================================

@app.get(
    "/health"
)
def health():

    return {
        "status": "ok",
        "service": "Egyptian ID Card Cropper"
    }


# ============================================================
# STARTUP
# ============================================================

if __name__ == "__main__":

    import uvicorn

    port = int(
        os.environ.get(
            "PORT",
            "7860"
        )
    )

    print(
        f"Starting Egyptian ID Card Cropper on port {port}"
    )

    uvicorn.run(
        app,
        host="0.0.0.0",
        port=port
    )