import os
import re
import cv2
import pandas as pd
import gradio as gr
import pytesseract
import spaces

from PIL import Image


# =========================================================
# ZERO GPU STARTUP CHECK
# =========================================================

@spaces.GPU
def zerogpu_probe():
    return "OK"


# =========================================================
# OCR CORRECTION
# =========================================================

def correct_ocr_mistakes(text):

    if not text:
        return ""

    text = text.upper()

    replacements = {
        "O": "0",
        "I": "1",
        "L": "1",
        "S": "5",
        "B": "8",
        "G": "6",
        "Z": "2"
    }

    for k, v in replacements.items():
        text = text.replace(k, v)

    return text


# =========================================================
# EXTRACT LAST 3 DIGITS
# =========================================================

def extract_last_3_digits(val):

    if not val:
        return ""

    val = correct_ocr_mistakes(val)

    digits = re.findall(r"\d", val)

    if len(digits) >= 3:
        return "".join(digits[-3:])

    return ""


# =========================================================
# PARSE OCR TEXT
# =========================================================

def parse_ocr_text(text):

    awb_val = ""
    packet_val = ""
    ship_to_val = ""

    if not text:
        return "", "", ""

    lines = [
        line.strip()
        for line in text.split("\n")
        if line.strip()
    ]

    for i, line in enumerate(lines):

        # -------------------------------------------------
        # AWB
        # -------------------------------------------------

        if (
            re.search(r"AW[B8]\s*[:\-]", line, re.IGNORECASE)
            or "AWB" in line.upper()
        ):

            m = re.search(
                r"AW[B8]\s*[:\-]?\s*([A-Z0-9_]+)",
                line,
                re.IGNORECASE
            )

            if m and len(m.group(1)) > 4:
                awb_val = m.group(1)

            elif i + 1 < len(lines):
                awb_val = lines[i + 1]


        # -------------------------------------------------
        # PACKET ID
        # -------------------------------------------------

        if re.search(
            r"Packet\s*[I1][D0]",
            line,
            re.IGNORECASE
        ):

            m = re.search(
                r"Packet\s*[I1][D0]\s*[:\-]?\s*([A-Z0-9_]+)",
                line,
                re.IGNORECASE
            )

            if m and len(m.group(1)) > 4:
                packet_val = m.group(1)

            elif i + 1 < len(lines):
                packet_val = lines[i + 1]


        # -------------------------------------------------
        # SHIP TO
        # -------------------------------------------------

        if re.search(
            r"Ship\s*T[o0]",
            line,
            re.IGNORECASE
        ):

            m = re.search(
                r"Ship\s*T[o0]\s*[:\-]?\s*(.+)",
                line,
                re.IGNORECASE
            )

            if m and m.group(1).strip():
                ship_to_val = m.group(1).strip()

            elif i + 1 < len(lines):
                ship_to_val = lines[i + 1].strip()


    return (
        extract_last_3_digits(awb_val),
        extract_last_3_digits(packet_val),
        ship_to_val
    )


# =========================================================
# IMAGE PREPROCESSING
# =========================================================

def preprocess_image(image_path):

    img = cv2.imread(image_path)

    if img is None:
        return []

    # Grayscale
    gray = cv2.cvtColor(
        img,
        cv2.COLOR_BGR2GRAY
    )

    # Resize
    gray = cv2.resize(
        gray,
        None,
        fx=2,
        fy=2,
        interpolation=cv2.INTER_CUBIC
    )

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

    # Adaptive threshold
    adaptive_thresh = cv2.adaptiveThreshold(
        denoised,
        255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,
        31,
        2
    )

    return [
        Image.fromarray(adaptive_thresh),
        Image.fromarray(gray)
    ]


# =========================================================
# PROCESS IMAGES
# =========================================================

def process_images(files, progress=gr.Progress()):

    if not files:

        return (
            pd.DataFrame(
                columns=[
                    "AWB",
                    "Packet ID",
                    "Ship To"
                ]
            ),
            None
        )


    results = []


    # OCR configurations
    ocr_configs = [
        "--psm 6",
        "--psm 11"
    ]


    total_files = len(files)


    # =====================================================
    # LOOP THROUGH FILES
    # =====================================================

    for index, file_obj in enumerate(files):

        progress(
            (index + 1) / total_files,
            desc=f"Processing {index + 1}/{total_files}"
        )


        # -------------------------------------------------
        # GET FILE PATH
        # -------------------------------------------------

        if isinstance(file_obj, str):

            filepath = file_obj

        elif hasattr(file_obj, "name"):

            filepath = file_obj.name

        else:

            filepath = str(file_obj)


        # -------------------------------------------------
        # PREPROCESS IMAGE
        # -------------------------------------------------

        preprocessed_imgs = preprocess_image(
            filepath
        )


        final_awb = ""
        final_packet = ""
        final_ship = ""

        found_all = False


        # =================================================
        # OCR
        # =================================================

        for img in preprocessed_imgs:

            if found_all:
                break


            for config in ocr_configs:

                try:

                    text = pytesseract.image_to_string(
                        img,
                        config=config
                    )

                except Exception:

                    continue


                # Parse OCR
                awb, packet, ship = parse_ocr_text(
                    text
                )


                # Save AWB
                if not final_awb and awb:
                    final_awb = awb


                # Save Packet ID
                if not final_packet and packet:
                    final_packet = packet


                # Save Ship To
                if not final_ship and ship:
                    final_ship = ship


                # Stop if everything found
                if (
                    final_awb
                    and final_packet
                    and final_ship
                ):

                    found_all = True
                    break


        # =================================================
        # ADD RESULT
        # =================================================

        results.append({
            "AWB": final_awb,
            "Packet ID": final_packet,
            "Ship To": final_ship
        })


    # =====================================================
    # CREATE DATAFRAME
    # =====================================================

    df = pd.DataFrame(
        results,
        columns=[
            "AWB",
            "Packet ID",
            "Ship To"
        ]
    )


    # =====================================================
    # CREATE EXCEL
    # =====================================================

    output_excel = os.path.abspath(
        "RVP_Data.xlsx"
    )

    df.to_excel(
        output_excel,
        index=False
    )


    return df, output_excel


# =========================================================
# GRADIO INTERFACE
# =========================================================

with gr.Blocks(
    title="RVP Label OCR Extractor"
) as interface:


    # -----------------------------------------------------
    # TITLE
    # -----------------------------------------------------

    gr.Markdown(
        "# 📦 RVP Label OCR Extractor"
    )


    gr.Markdown(
        "Upload RVP label images to extract "
        "the last 3 digits of AWB, "
        "last 3 digits of Packet ID, "
        "and Ship To."
    )


    # -----------------------------------------------------
    # FILE UPLOAD
    # -----------------------------------------------------

    file_input = gr.File(
        file_count="multiple",
        file_types=[
            ".jpg",
            ".jpeg",
            ".png"
        ],
        label="Upload RVP Photos",
        type="filepath"
    )


    # -----------------------------------------------------
    # EXTRACT BUTTON
    # -----------------------------------------------------

    extract_btn = gr.Button(
        "Extract Data",
        variant="primary"
    )


    # -----------------------------------------------------
    # RESULTS
    # -----------------------------------------------------

    output_df = gr.Dataframe(
        headers=[
            "AWB",
            "Packet ID",
            "Ship To"
        ],
        datatype=[
            "str",
            "str",
            "str"
        ],
        label="Results",
        interactive=False
    )


    # -----------------------------------------------------
    # EXCEL DOWNLOAD
    # -----------------------------------------------------

    output_file = gr.File(
        label="Download Excel File"
    )


    # -----------------------------------------------------
    # BUTTON EVENT
    # -----------------------------------------------------

    extract_btn.click(
        fn=process_images,
        inputs=file_input,
        outputs=[
            output_df,
            output_file
        ]
    )


# =========================================================
# LAUNCH
# =========================================================

if __name__ == "__main__":

    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False
    )