import gradio as gr
from transformers import BlipProcessor, BlipForQuestionAnswering, MarianMTModel, MarianTokenizer
from PIL import Image, ImageDraw, ImageFont
import torch, uuid, os
from datetime import datetime

# تحميل نموذج BLIP المدرب مسبقًا
blip_model = BlipForQuestionAnswering.from_pretrained("mshsahmed/blip-vqa-finetuned-kvasir-v58849")
processor = BlipProcessor.from_pretrained("mshsahmed/blip-vqa-finetuned-kvasir-v58849")

# تحميل نماذج الترجمة
ar_en_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-ar-en")
ar_en_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-ar-en")

en_ar_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
en_ar_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-ar")

# دوال الترجمة
def translate_ar_to_en(text):
    inputs = ar_en_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = ar_en_model.generate(**inputs)
    return ar_en_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

def translate_en_to_ar(text):
    inputs = en_ar_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = en_ar_model.generate(**inputs)
    return en_ar_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

# قاموس ترجمة طبي يدوي
medical_terms = {
    "colonoscopy": "تنظير القولون ",
    "Ulcerative Colitis": "التهاب القولون التقرحي",
    "Have all polyps been removed?": "هل تم إزالة جميع الاورام الحميدة",
    "polyps": " الاورام الحميدة",
    "What type of polyp is present?": "  ما هو نوع الورم الموجود",
    "gastroscopy": "تنظير المعدة ",
    "polyp": "ورم "

   }

# دالة الترجمة الذكية للإجابات
def translate_answer_medical(answer_en):
    key = answer_en.lower().strip()
    if key in medical_terms:
        return medical_terms[key]
    else:
        return translate_en_to_ar(answer_en)

# ✅ Arabic font helper
def get_font(size=22):
    try:
        return ImageFont.truetype("Amiri-Regular.ttf", size)
    except:
        return ImageFont.load_default()

# ✅ Report generation function
def generate_report_image(image, question_ar, question_en, answer_ar, answer_en):
    width, height = 1000, 700
    background = Image.new("RGB", (width, height), color="white")
    draw = ImageDraw.Draw(background)
    font = get_font(22)
    font_bold = get_font(26)

    draw.text((40, 20), " Medical VQA Report", font=font_bold, fill="black")
    draw.text((40, 60), f"Date:{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", font=font, fill="gray")
    # Header background
    

    img_resized = image.resize((300, 300))
    background.paste(img_resized, (50, 110))

    x, y = 380, 110
    spacing = 70
    lines = [
        f"السؤال بالعربية :{question_ar}",
        
        f"الإجابة بالعربية :{answer_ar}",
        f"Question in English:{question_en}",

        f"Answer in English:{answer_en}",
        f" © 2025 NTI - Artificial Intelligence in Medical Project"
    ]

    for line in lines:
        for subline in line.split("\n"):
            draw.text((x, y), subline, font=font, fill="black")
            y += spacing

    file_name = f"report_{uuid.uuid4().hex[:8]}.png"
    background.save(file_name)
    return file_name

# ✅ Main VQA function
def vqa_multilingual(image, question):
    if not image or not question.strip():
        return "يرجى رفع صورة وكتابة سؤال.", "", "", "", None

    is_arabic = any('\u0600' <= c <= '\u06FF' for c in question)
    question_ar = question.strip() if is_arabic else translate_en_to_ar(question)
    question_en = translate_ar_to_en(question) if is_arabic else question.strip()

    inputs = processor(image, question_en, return_tensors="pt")
    with torch.no_grad():
        output = blip_model.generate(**inputs)
    answer_en = processor.decode(output[0], skip_special_tokens=True).strip()
    answer_ar = translate_answer_medical(answer_en)

    report_image_path = generate_report_image(image, question_ar, question_en, answer_ar, answer_en)

    return (
        question_ar,
        question_en,
        answer_ar,
        answer_en,
        report_image_path
    )

# واجهة Gradio
gr.Interface(
    fn=vqa_multilingual,
    inputs=[
        gr.Image(type="pil", label="🔍 ارفع صورة الأشعة"),
        gr.Textbox(label="💬 أدخل سؤالك (بالعربية أو الإنجليزية)")
    ],
    outputs=[
        gr.Textbox(label="🟠 السؤال بالعربية"),
        gr.Textbox(label="🟢 السؤال بالإنجليزية"),
        gr.Textbox(label="🟠 الإجابة بالعربية"),
        gr.Textbox(label="🟢 الإجابة بالإنجليزية"),
        gr.Image(type="filepath", label="📸 Report")

    ],
    title=" نموذج  ثنائي اللغة (عربي - إنجليزي) خاص بمنظار المعدة ",
    description="ارفع صورة طبية واسأل بالعربية أو الإنجليزية، وستحصل على الإجابة باللغتين."
).launch(share=True)