import os
import cv2
import numpy as np
import gradio as gr
from ultralytics import YOLO

# 1. Load Models - Fixed paths to match Docker structure
plate_model = YOLO("car_plate_best.pt")
ocr_model = YOLO("model_ocr.pt")

# Explicit mapping to ARABIC-INDIC digits (٠١٢٣٤٥٦٧٨٩)
arabic_mapping = {
    "aain": "ع", "alef": "أ", "a": "أ", "alf": "أ", "baa": "ب",
    "daal": "د", "dal": "د", "e": "ع", "faa": "ف", "geem": "ج",
    "haa": "هـ", "kaaf": "ك", "laam": "ل", "meem": "م", "noon": "ن",
    "qaf": "ق", "raa": "ر", "sad": "ص", "seen": "س", "taa": "ط",
    "waaw": "و", "waw": "و", "yaa": "ى", "zay": "ز", "dad": "ض",
    "0": "٠", "1": "١", "2": "٢", "3": "٣", "4": "٤",
    "5": "٥", "6": "٦", "7": "٧", "8": "٨", "9": "٩",
}

def predict_plate(img):
    if img is None:
        return None, "", ""

    try:
        # Gradio sends RGB, YOLO likes BGR
        img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
        display_img = img_cv.copy()
        
        # 1. Plate Detection
        plate_results = plate_model.predict(source=img_cv, conf=0.30, verbose=False)
        best_plate_box = None
        max_conf = -1
        
        for box in plate_results[0].boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = float(box.conf)
            cv2.rectangle(display_img, (x1, y1), (x2, y2), (0, 255, 0), 3)
            if conf > max_conf:
                max_conf = conf
                best_plate_box = (x1, y1, x2, y2)

        letters_str, numbers_str = "", ""
        
        if best_plate_box:
            x1, y1, x2, y2 = best_plate_box
            h, w = img_cv.shape[:2]
            margin = 20 # Increased margin for better OCR on HF
            cx1, cy1, cx2, cy2 = max(0, x1-margin), max(0, y1-margin), min(w, x2+margin), min(h, y2+margin)
            plate_crop = img_cv[cy1:cy2, cx1:cx2]

            if plate_crop.size > 0:
                # Optimized OCR params for Hugging Face CPU/GPU
                ocr_results = ocr_model.predict(source=plate_crop, conf=0.25, iou=0.45, agnostic_nms=True, verbose=False)
                
                detections = []
                for obox in ocr_results[0].boxes:
                    cid = int(obox.cls)
                    ename = ocr_results[0].names[cid]
                    
                    # Ignore 'License Plate' label
                    if ename.lower() == "license plate":
                        continue
                        
                    achr = arabic_mapping.get(ename, ename)
                    ox1, _, ox2, _ = obox.xyxy[0]
                    detections.append(((ox1 + ox2) / 2, achr))
                
                # Spatial sort Left-to-Right
                detections.sort(key=lambda d: d[0])
                
                arabic_digits_set = set("٠١٢٣٤٥٦٧٨٩")
                
                det_letters = [d[1] for d in detections if d[1] not in arabic_digits_set]
                det_numbers = [d[1] for d in detections if d[1] in arabic_digits_set]
                
                # Reverse for RTL display
                det_numbers.reverse() 
                det_letters.reverse()
                
                letters_str = " ".join(det_letters)
                numbers_str = " ".join(det_numbers)

        return cv2.cvtColor(display_img, cv2.COLOR_BGR2RGB), numbers_str, letters_str
    
    except Exception:
        return None, "Error", "Error"

# UI Setup
with gr.Blocks(theme=gr.themes.Default()) as demo:
    gr.Markdown("# 🚗 Arabic License Plate OCR")
    
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="Upload or Camera", type="pil")
            btn = gr.Button("Extract Plate Info", variant="primary")
            
        with gr.Column():
            output_img = gr.Image(label="Detection Result")
            with gr.Row():
                # Numbers (Left) | Letters (Right)
                out_numbers = gr.Textbox(label="Numbers (Arabic Script)", interactive=False)
                out_letters = gr.Textbox(label="Arabic Letters", interactive=False)

    btn.click(fn=predict_plate, inputs=input_img, outputs=[output_img, out_numbers, out_letters], api_name="predict_plate")

if __name__ == "__main__":
    # CRITICAL: Added server_name and server_port for Hugging Face
    demo.launch(server_name="0.0.0.0", server_port=7860)