import os
import cv2
import gradio as gr
import google.generativeai as genai
from ultralytics import YOLO
import tempfile
import torch
import spaces

import numpy as np
from PIL import Image, ImageDraw, ImageFont
import arabic_reshaper
from bidi.algorithm import get_display

# =============================
# Gemini API Key
# =============================
# ⚠️ لا تضع المفتاح داخل كود عام. الأفضل: Secrets/Env
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") or "AIzaSyAvm28ZnTMaZ1Jtg9sYM-EO4qlAN2W4BIQ"
genai.configure(api_key=GEMINI_API_KEY)

SYSTEM_PROMPT = (
    "أنت مساعد ذكي يستقبل كلمات أو جمل قصيرة قادمة من مترجم لغة الإشارة العربية.\n"
    "ملاحظة مهمة: الإدخال قد يكون غير نظيف ويحتوي تكرارًا كبيرًا للحروف بسبب الترجمة، مثل: "
    "\"ااااللللااااممممممممببممم\" (المقصود غالبًا: \"ألم\").\n\n"
    "مهمتك:\n"
    "1) تنظيف النص بإزالة تكرار الحروف غير الطبيعي ودمجها إلى كلمة صحيحة.\n"
    "2) تصحيح الكلمة قدر الإمكان لتصبح عربية واضحة ومفهومة.\n"
    "3) إذا كانت كلمة واحدة فقط مثل \"ألم\" أو \"دوخة\" أو \"غثيان\": اشرح معناها باختصار.\n"
    "4) إذا كانت جملة: أعد صياغتها كنص عربي واضح ومفهوم دون تغيير المعنى.\n"
)


def fix_with_gemini(raw_text: str) -> str:
    if not raw_text:
        return ""
    try:
        model = genai.GenerativeModel("models/gemini-2.5-flash")
        prompt = SYSTEM_PROMPT + f"\n\nالنص الخام:\n«{raw_text}»"
        resp = model.generate_content(prompt)
        return (resp.text or "").strip()
    except Exception as e:
        return f"خطأ في Gemini: {e}"


# =============================
# إعدادات YOLO
# =============================
WEIGHTS_PATH = "best.pt"

IMG_SIZE_VIDEO = 320
IMG_SIZE_IMAGE = 1280

CONF_THRESHOLD = 0.05
IOU_THRESHOLD = 0.70
MAX_DET = 20

MIN_STABLE_FRAMES = 1
FRAME_SKIP = 1
MAX_FRAMES = 1000
WORD_GAP_FRAMES = 10

CENTER_CROP = False

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

yolo_model = None
DEVICE = "cpu"

def get_model():
    global yolo_model, DEVICE

    if yolo_model is None:
        print("🔹 Loading YOLO model...")
        yolo_model = YOLO(WEIGHTS_PATH)
        print("📚 Classes:", yolo_model.names)
        print("🧠 YOLO task:", getattr(yolo_model, "task", "unknown"))

    if torch.cuda.is_available():
        if DEVICE != "cuda":
            DEVICE = "cuda"
            try:
                yolo_model.to(DEVICE)
                print("✅ YOLO model moved to cuda")
            except Exception as e:
                print("⚠️ تعذر نقل الموديل إلى cuda:", e)
    else:
        DEVICE = "cpu"

    return yolo_model


# =============================
# رسم عربي على الفيديو via PIL (✅ Bold + أكبر)
# =============================
FONT_PATH = os.path.join(os.path.dirname(__file__), "NotoNaskhArabic-VariableFont_wght.ttf")

# ✅ تحكم سريع بالحجم/الثخانة
AR_TEXT_SIZE = 46          # كان 36 -> كبرناه شوي
AR_STROKE_W = 3            # يزيد “البولد” (جرّب 2/3/4)
AR_STROKE_RGB = (0, 0, 0)  # حد خارجي أسود ليوضح على الفيديو

def draw_arabic_text(frame_bgr, text, x, y, font_size=AR_TEXT_SIZE, bgr_color=(0, 255, 0)):
    img = Image.fromarray(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img)

    try:
        font = ImageFont.truetype(FONT_PATH, font_size)
    except Exception as e:
        print("⚠️ خطأ تحميل الخط العربي:", e)
        font = ImageFont.load_default()

    shaped = arabic_reshaper.reshape(text)
    rtl_text = get_display(shaped)

    rgb_fill = (bgr_color[2], bgr_color[1], bgr_color[0])

    # ✅ stroke_width يعطيك خط أثخن (زي بولد) + حد خارجي
    draw.text((x, y), rtl_text, font=font, fill=rgb_fill)
    (x, y),
    rtl_text,
    font=font,
    fill=rgb_fill,
    stroke_width=AR_STROKE_W,
    stroke_fill=rgb_fill

    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)


# =============================
# (من كودك الأساسي) تكبير + قص من الوسط
# =============================
def resize_and_center_crop(frame, target: int = 640):
    h, w = frame.shape[:2]
    short_side = min(w, h)
    if short_side <= 0:
        return frame

    scale = target / short_side
    new_w = int(w * scale)
    new_h = int(h * scale)

    frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)

    h, w = frame.shape[:2]
    x1 = max(0, (w - target) // 2)
    y1 = max(0, (h - target) // 2)
    x2 = min(x1 + target, w)
    y2 = min(y1 + target, h)

    crop = frame[y1:y2, x1:x2]

    ch, cw = crop.shape[:2]
    if ch != target or cw != target:
        crop = cv2.resize(crop, (target, target), interpolation=cv2.INTER_AREA)

    return crop


# =============================
# تجهيز الفيديو قبل المعالجة (بدون قص)
# =============================
def preprocess_video(input_path: str, target_short_side: int = 640, target_fps: int = 8) -> str:
    cap = cv2.VideoCapture(input_path)
    if not cap.isOpened():
        print("[preprocess] تعذر فتح الفيديو، سنستخدم الملف الأصلي كما هو.")
        return input_path

    orig_fps = cap.get(cv2.CAP_PROP_FPS)
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    if orig_fps <= 0:
        frame_step = 1
        out_fps = float(target_fps)
    else:
        frame_step = max(1, int(round(orig_fps / target_fps)))
        out_fps = orig_fps / frame_step

    short_side = min(w, h) if min(w, h) > 0 else 1
    scale = float(target_short_side) / float(short_side)
    new_w = max(1, int(w * scale))
    new_h = max(1, int(h * scale))

    fd, tmp_path = tempfile.mkstemp(suffix=".mp4")
    os.close(fd)

    out_w, out_h = (new_w, new_h)

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out = cv2.VideoWriter(tmp_path, fourcc, out_fps, (out_w, out_h))

    frame_idx = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break

        if frame_idx % frame_step == 0:
            if CENTER_CROP:
                processed = resize_and_center_crop(frame, target=target_short_side)
            else:
                processed = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
            out.write(processed)

        frame_idx += 1

    cap.release()
    out.release()
    print(f"[preprocess] orig=({w}x{h}), new=({out_w}x{out_h}), saved={tmp_path}")
    return tmp_path


# =============================
# detect frame (BGR + RTL sorting للصورة)
# =============================
def detect_frame(frame_bgr, imgsz: int, sort_mode: str = "conf"):
    """
    sort_mode:
      - "conf": أعلى ثقة أولاً (مناسب للفيديو)
      - "x": ترتيب أفقي RTL (من يمين لليسار) (مناسب للصورة)
    """
    model = get_model()

    result = model.predict(
        source=frame_bgr,
        conf=CONF_THRESHOLD,
        imgsz=imgsz,
        iou=IOU_THRESHOLD,
        max_det=MAX_DET,
        verbose=False,
        device=DEVICE,
    )[0]

    task = getattr(model, "task", "unknown")
    has_probs = hasattr(result, "probs") and result.probs is not None

    boxes = result.boxes
    num_boxes = 0 if boxes is None else len(boxes)

    max_conf = 0.0
    if boxes is not None and len(boxes) > 0:
        try:
            max_conf = float(boxes.conf.max().item())
        except Exception:
            try:
                max_conf = float(boxes.conf.max())
            except Exception:
                max_conf = 0.0

    print(
        f"[detect_frame] task={task} imgsz={imgsz} conf_thr={CONF_THRESHOLD} "
        f"iou={IOU_THRESHOLD} max_det={MAX_DET} boxes={num_boxes} max_conf={max_conf:.4f} has_probs={has_probs}"
    )

    if boxes is None or len(boxes) == 0:
        try:
            plotted = result.plot()  # numpy BGR
            return [], plotted
        except Exception:
            return [], frame_bgr

    items = []
    for box in list(boxes):
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        cls_id = int(box.cls[0])
        conf = float(box.conf[0]) if hasattr(box, "conf") else 0.0

        if isinstance(model.names, dict):
            eng = model.names.get(cls_id, str(cls_id))
        else:
            eng = model.names[cls_id] if cls_id < len(model.names) else str(cls_id)

        letter = arabic_map.get(eng, eng)
        items.append((x1, y1, x2, y2, conf, letter))

    if sort_mode == "x":
        # ✅ RTL: من اليمين لليسار (نرتب على x1 تنازلي)
        items.sort(key=lambda t: t[0], reverse=True)
    else:
        # أعلى ثقة أولاً
        items.sort(key=lambda t: t[4], reverse=True)

    labels = []
    for x1, y1, x2, y2, conf, letter in items:
        labels.append(letter)
        cv2.rectangle(frame_bgr, (x1, y1), (x2, y2), (0, 255, 0), 2)
        frame_bgr = draw_arabic_text(frame_bgr, letter, x1, max(0, y1 - 55), font_size=AR_TEXT_SIZE)

    return labels, frame_bgr


# =============================
# VIDEO → RAW TEXT + OUTPUT VIDEO + DEBUG
# =============================
def extract_and_render(video_path: str):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return "", None, "تعذر فتح الفيديو في extract_and_render"

    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    out_path = "processed_output.mp4"

    fps = cap.get(cv2.CAP_PROP_FPS)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    if fps <= 0:
        fps = 8.0

    out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))

    word = ""
    words = []
    last_label = None
    last_added = None
    stable = 0
    last_seen = None
    frame_index = 0

    frames_with_dets = 0
    debug_lines = []

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        frame_index += 1
        if frame_index > MAX_FRAMES:
            break

        if FRAME_SKIP > 1 and frame_index % FRAME_SKIP != 0:
            continue

        frame = cv2.flip(frame, 1)

        labels, rendered = detect_frame(frame, imgsz=IMG_SIZE_VIDEO, sort_mode="conf")
        out.write(rendered)

        if labels:
            frames_with_dets += 1
            debug_lines.append(f"frame {frame_index}: {labels}")

            label = labels[0]  # ✅ أعلى ثقة (يد واحدة/حرف واحد)
            last_seen = frame_index

            if label == last_label:
                stable += 1
            else:
                last_label = label
                stable = 1

            if stable >= MIN_STABLE_FRAMES:
                if label != last_added:
                    word += label
                    last_added = label
                stable = 0
        else:
            if word and last_seen and (frame_index - last_seen >= WORD_GAP_FRAMES):
                words.append(word)
                word = ""
                last_label = None
                last_added = None
                stable = 0
                last_seen = None

    cap.release()
    out.release()

    if word:
        words.append(word)

    raw_text = " ".join(words).strip()

    if not debug_lines:
        debug_info = (
            f"total_frames={frame_index}, frames_with_detections=0\n"
            "لم يتم رصد أي صناديق (boxes) من YOLO.\n"
            "افتح Logs وابحث عن [detect_frame] وشوف task/has_probs/max_conf.\n"
            "لو task=classify أو has_probs=True مع boxes=0 غالباً best.pt مو Detection."
        )
    else:
        sample = "\n".join(debug_lines[:30])
        debug_info = (
            f"total_frames={frame_index}, frames_with_detections={frames_with_dets}\n"
            "أمثلة من الفريمات اللي فيها حروف:\n"
            f"{sample}"
        )

    return raw_text, out_path, debug_info


# =============================
# IMAGE → RAW TEXT + OUTPUT IMAGE + DEBUG
# =============================
def process_image(pil_img: Image.Image):
    if pil_img is None:
        return "", None, "لم يتم رفع صورة"

    rgb = np.array(pil_img.convert("RGB"))
    frame_bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)

    labels, rendered_bgr = detect_frame(frame_bgr, imgsz=IMG_SIZE_IMAGE, sort_mode="x")

    raw_text = "".join(labels).strip()
    rendered_rgb = cv2.cvtColor(rendered_bgr, cv2.COLOR_BGR2RGB)
    rendered_pil = Image.fromarray(rendered_rgb)

    if not labels:
        debug_info = (
            "لا يوجد detections في الصورة.\n"
            "افتح Logs وشوف سطر [detect_frame].\n"
            "- إذا task=classify (أو has_probs=True): best.pt مو Detection.\n"
            "- إذا max_conf قريب من 0: جرّب صور أوضح/قريبة أو ارفع imgsz أو خفض conf."
        )
    else:
        debug_info = f"labels={labels}"

    return raw_text, rendered_pil, debug_info


# =============================
# Gradio + @spaces.GPU
# =============================
@spaces.GPU
def run_video(file):
    if file is None:
        return "لم يتم رفع فيديو", "", None, "لم يتم رفع فيديو"

    video_path = file.name
    light_path = preprocess_video(video_path, target_short_side=640, target_fps=8)

    raw, processed_path, debug_info = extract_and_render(light_path)
    pretty = fix_with_gemini(raw) if raw else ""

    if not raw:
        raw = "لم يتم التعرف على أي نص من الإشارات."

    return raw, pretty, processed_path, debug_info


@spaces.GPU
def run_image(img):
    if img is None:
        return "لم يتم رفع صورة", "", None, "لم يتم رفع صورة"

    raw, rendered_pil, debug_info = process_image(img)
    pretty = fix_with_gemini(raw) if raw else ""

    if not raw:
        raw = "لم يتم التعرف على أي نص من الصورة."

    return raw, pretty, rendered_pil, debug_info


with gr.Blocks() as demo:
    gr.Markdown("## 🤟 ASL → Arabic (YOLO + Gemini) — فيديو + صور (الصورة RTL: يمين → يسار)")

    with gr.Tabs():
        with gr.Tab("Video"):
            inp_v = gr.File(label="ارفع فيديو الإشارة")
            raw_v = gr.Textbox(label="النص الخام", lines=3)
            pretty_v = gr.Textbox(label="النص المحسن (Gemini)", lines=3)
            video_out = gr.Video(label="الفيديو بعد البروسيس")
            debug_v = gr.Textbox(label="Debug info", lines=10)

            btn_v = gr.Button("ابدأ المعالجة (فيديو)")
            btn_v.click(run_video, inputs=[inp_v], outputs=[raw_v, pretty_v, video_out, debug_v])

        with gr.Tab("Image"):
            inp_i = gr.Image(label="ارفع صورة الإشارة", type="pil")
            raw_i = gr.Textbox(label="النص الخام (من الصورة)", lines=3)
            pretty_i = gr.Textbox(label="النص المحسن (Gemini)", lines=3)
            img_out = gr.Image(label="الصورة بعد الديتكشن", type="pil")
            debug_i = gr.Textbox(label="Debug info", lines=8)

            btn_i = gr.Button("ابدأ المعالجة (صورة)")
            btn_i.click(run_image, inputs=[inp_i], outputs=[raw_i, pretty_i, img_out, debug_i])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)