import gradio as gr
from transformers import BlipProcessor, BlipForQuestionAnswering, MarianMTModel, MarianTokenizer
from PIL import Image
import torch
import uuid
import io
from datetime import datetime
import requests
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import A4
from reportlab.lib.utils import ImageReader

# Load BLIP model
blip_model = BlipForQuestionAnswering.from_pretrained("sharawy53/diploma")
processor = BlipProcessor.from_pretrained("sharawy53/diploma")

# Load translation models
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")

# Manual translation for medical terms
medical_terms = {
    "chest x-ray": "أشعة سينية للصدر",
    "x-ray": "أشعة سينية",
    "ct scan": "تصوير مقطعي محوسب",
    "mri": "تصوير بالرنين المغناطيسي",
    "ultrasound": "تصوير بالموجات فوق الصوتية",
    "normal": "طبيعي",
    "abnormal": "غير طبيعي",
    "brain": "الدماغ",
    "fracture": "كسر",
    "no abnormality detected": "لا توجد شذوذات",
    "left lung": "الرئة اليسرى",
    "right lung": "الرئة اليمنى",
    "cardiomegaly": "تضخم القلب",
    "ct": "المسح المقطعي",
    "CT": "المسح المقطعي",
    "MRI": "تصوير بالرنين المغناطيسي"
}

# Fixes for known questions
question_fixes = {
    "what is the unnatural in this image?": "ما الشيء غير الطبيعي في هذه الصورة؟",
    "what is abnormal in this image?": "ما الشيء غير الطبيعي في هذه الصورة؟",
    "is this image normal?": "هل هذه الصورة طبيعية؟"
}

# Translation functions
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):
    if text.lower().strip() in question_fixes:
        return question_fixes[text.lower().strip()]
    inputs = en_ar_tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = en_ar_model.generate(**inputs)
    translated = en_ar_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
    if "القرآن" in translated or "يتفاعل" in translated:
        return question_fixes.get(text.lower().strip(), "سؤال غير مفهوم")
    return translated

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)

# Main function
def vqa_multilingual_save(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)

    # PDF report generation
    pdf_filename = f"report_{uuid.uuid4().hex[:8]}.pdf"
    c = canvas.Canvas(pdf_filename, pagesize=A4)
    width, height = A4

    # Title
    c.setFont("Helvetica-Bold", 16)
    c.drawString(50, height - 50, "🧠 تقرير الإجابة على سؤال الأشعة الطبية")

    # NTI logo from URL
    try:
        logo_url = "https://www.ntiegypt.sci.eg/moodle/pluginfile.php/1/core_admin/logocompact/300x300/1733222082/NTI%20Logo.png"
        logo_response = requests.get(logo_url)
        logo_image = ImageReader(io.BytesIO(logo_response.content))
        c.drawImage(logo_image, width - 120, height - 100, width=70, preserveAspectRatio=True, mask='auto')
    except:
        pass

    # Timestamp
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    c.setFont("Helvetica", 10)
    c.drawString(50, height - 70, f"📅 التاريخ والوقت: {timestamp}")

    # Image
    img_buffer = io.BytesIO()
    image.save(img_buffer, format="PNG")
    img_buffer.seek(0)
    c.drawImage(ImageReader(img_buffer), 50, height - 350, width=200, preserveAspectRatio=True, mask='auto')

    # Text
    c.setFont("Helvetica", 12)
    text_start = height - 370
    line_height = 20
    lines = [
        f"السؤال بالعربية: {question_ar}",
        f"السؤال بالإنجليزية: {question_en}",
        f"الإجابة بالعربية: {answer_ar}",
        f"الإجابة بالإنجليزية: {answer_en}"
    ]
    for i, line in enumerate(lines):
        c.drawString(50, text_start - i * line_height, line)

    c.save()

    return (
        f"السؤال بالعربية:\n{question_ar}",
        f"السؤال بالإنجليزية:\n{question_en}",
        f"الإجابة بالعربية:\n{answer_ar}",
        f"الإجابة بالإنجليزية:\n{answer_en}",
        pdf_filename
    )

# Gradio interface
gr.Interface(
    fn=vqa_multilingual_save,
    inputs=[
        gr.Image(type="pil", label="🔍 ارفع صورة الأشعة"),
        gr.Textbox(label="💬 أدخل سؤالك (بالعربية أو الإنجليزية)")
    ],
    outputs=[
        gr.Textbox(label="🟠 السؤال بالعربية"),
        gr.Textbox(label="🟢 السؤال بالإنجليزية"),
        gr.Textbox(label="🟠 الإجابة بالعربية"),
        gr.Textbox(label="🟢 الإجابة بالإنجليزية"),
        gr.File(label="📥 تحميل التقرير")
    ],
    title="🧠 نموذج ثنائي اللغة (عربي - إنجليزي) خاص بالأشعة الطبية",
    description="ارفع صورة طبية واسأل بالعربية أو الإنجليزية، وستحصل على الإجابة باللغتين مع إمكانية تحميل تقرير PDF يشمل الإجابة، الترجمة، والتاريخ."
).launch()