import gradio as gr
import cv2
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import arabic_reshaper
from bidi.algorithm import get_display
import os
import requests

# 1. Load Model
model = YOLO("best.pt")

# 2. Arabic letter mapping
arabic_letters = {
    "ain": "ع", "al": "لا", "aleff": "ا", "bb": "ب", "dal": "د", "dha": "ذ",
    "dhad": "ض", "fa": "ف", "gaaf": "ق", "ghain": "غ", "ha": "ه", "haa": "ح",
    "jeem": "ج", "kaaf": "ك", "khaa": "خ", "la": "ل", "laam": "ل", "meem": "م",
    "nun": "ن", "ra": "ر", "saad": "ص", "seen": "س", "sheen": "ش", "ta": "ط",
    "taa": "ت", "thaa": "ث", "thal": "ظ", "toot": "ط", "waw": "و", "ya": "ي",
    "yaa": "ي", "zay": "ز"
}

# --- BULLETPROOF FONT DOWNLOAD ---
FONT_PATH = "Amiri-Regular.ttf"
if not os.path.exists(FONT_PATH):
    print("📥 Downloading Arabic font...")
    try:
        res = requests.get("https://github.com/google/fonts/raw/main/ofl/amiri/Amiri-Regular.ttf")
        with open(FONT_PATH, "wb") as f:
            f.write(res.content)
        print("✅ Font ready!")
    except:
        FONT_PATH = None

def predict_live(img):
    if img is None: return None

    # Speed optimization
    results = model(img, conf=0.3, imgsz=320, verbose=False)
    
    pil_img = Image.fromarray(img)
    draw = ImageDraw.Draw(pil_img)

    # Load font
    try:
        font = ImageFont.truetype(FONT_PATH, 60) if FONT_PATH else ImageFont.load_default()
    except:
        font = ImageFont.load_default()

    for r in results:
        for box in r.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cls_name = model.names[int(box.cls[0])]

            # Arabic Transformation
            raw_text = arabic_letters.get(cls_name, cls_name)
            reshaped_text = arabic_reshaper.reshape(raw_text)
            bidi_text = get_display(reshaped_text)

            # Draw
            draw.rectangle([x1, y1, x2, y2], outline=(0, 255, 0), width=5)
            draw.text((x1, y1 - 75), bidi_text, font=font, fill=(0, 255, 0))

    return np.array(pil_img)

demo = gr.Interface(
    fn=predict_live,
    inputs=gr.Image(sources=["webcam"], streaming=True),
    outputs=gr.Image(label="Translator"),
    live=True
)

if __name__ == "__main__":
    demo.launch()