import pandas as pd
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
from sentence_transformers import SentenceTransformer, util
from difflib import get_close_matches
import requests
import io
import torch
# No need for OpenAI import - using free Hugging Face models only
from typing import Optional

# GitHub Excel file URL
GITHUB_EXCEL_URL = "https://raw.githubusercontent.com/3Zzamt0o/HospitalData/main/Hospital%20Data%20(1).xlsx"

def load_hospital_data():
    try:
        response = requests.get(GITHUB_EXCEL_URL)
        response.raise_for_status()
        excel_data = io.BytesIO(response.content)
        df = pd.read_excel(excel_data)
        print("Successfully loaded data from GitHub")
        return df
    except Exception as e:
        print(f"Error loading data from GitHub: {e}")
        return pd.DataFrame({
            'العيادات': ['عيادة الأسنان', 'عيادة الباطنة', 'عيادة العظام', 'عيادة الأطفال'],
            'الدكتور': ['د. أحمد محمد', 'د. محمد علي', 'د. علي حسن', 'د. سارة أحمد'],
            'الأيام': ['السبت والاثنين', 'الأحد والثلاثاء', 'الاثنين والأربعاء', 'الثلاثاء والخميس'],
            'سعر الكشف': ['100 جنيه', '150 جنيه', '120 جنيه', '130 جنيه']
        })

# Load the hospital data
print("Loading hospital data...")
df = load_hospital_data()

# Enhanced Q&A data for lab tests and equipment
data = {
    "السؤال": [
        "ما هي مواعيد معمل التحاليل؟", "ما هي المده المستغرقه لعمل التحاليل؟",
        "ما هي انواع التحاليل الموجوده و المتوفره؟", "هل يوجد تحليل صوره دم؟",
        "ماهي تحاليل صوره الدم؟", "ما هي تحاليل السكر؟", "ما هي تحاليل وظائف الكبد؟",
        "ما هي تحاليل وظائف الكلي؟", "ما هي تحاليل البول و البراز؟",
        "ما هي تحاليل الدهون؟", "ما هي تحاليل الألتهابات و الروماتيزم؟",
        "ما هي تحاليل الفيروسات؟", "ما هي تحاليل الغده الدرقيه؟",
        "هل يوجد حشو اسنان؟", "هل يوجد جهاز سونار؟", "هل يوجد اجهزه تنفس اصطناعي؟",
        "هل يوجد قياس نظر؟", "ما هي الاجهزه المتاحه في عياده الرمد؟",
        "هل يوجد قياس ضغط؟", "هل يوجد سونار؟", "هل يوجد رسم قلب؟"
    ],
    "الجواب": [
        "طول ايام الاسبوع من 9 صباحا الي 9 مسائا", "من 24 ساعه الي 48 ساعه",
        "تحاليل صوره الدم، تحاليل سكر، تحاليل وظائف الكبد، تحاليل وظائف الكلي",
        "نعم يوجد جميع انواع التحاليل بالمستشفي", "CBC، Hb%(Heamoglobin)، Blood Group، Rh",
        "Rbs سكر عشوائي، سكر صائم FBS، PPbs سكر فاطر", "ALT(GPT)، AST(GOT)، Billirubin",
        "Ceatinin الكرياتنين، Urea اليوريا، Urine البول، Uric acid اليوريك اسد",
        "Urine البول، Stool البراز، جرثومه المعده في البراز", "Cholesterol كوليسترول، Triglyceride الدهون الثلاثيه",
        "CRP، ASOT", "HCV Ab، GBS Ag، HIC Ab", "Tsh، T3، t4",
        "نعم يوجد حشو أسنان", "نعم يوجد جهاز سونار", "نعم يوجد أجهزة تنفس إصطناعي",
        "نعم يوجد قياس نظر", "جهاز قاع العين، جهاز قياس النظر، جهاز الكشف من الحساسية والالتهابات",
        "نعم يوجد قياس ضغط", "نعم يوجد سونار", "نعم يوجد رسم قلب"
    ]
}

qa_data = pd.DataFrame(data)

# Expanded medical knowledge base with more comprehensive information
medical_kb = {
    "الجفاف": {
        "symptoms": ["عطش شديد", "فم جاف", "بول أصفر داكن", "إعياء", "دوخة"],
        "treatment": "شرب الكثير من السوائل، تناول الأطعمة الغنية بالماء، تجنب التعرض المباشر للشمس، استشارة الطبيب إذا استمرت الأعراض",
        "prevention": "شرب الماء بانتظام، تناول الفواكه والخضروات، تجنب المشروبات الكحولية والكافيين الزائد"
    },
    "الصداع": {
        "symptoms": ["ألم في الرأس", "حساسية للضوء", "غثيان"],
        "treatment": "الراحة في مكان هادئ، شرب الماء بكثرة، تناول مسكن خفيف، تجنب الضوء القوي والضوضاء",
        "prevention": "النوم الكافي، تجنب التوتر، تناول وجبات منتظمة"
    },
    "الحمى": {
        "symptoms": ["ارتفاع درجة الحرارة", "قشعريرة", "تعرق", "صداع"],
        "treatment": "الراحة التامة، شرب السوائل بكثرة، خفض درجة الحرارة بالكمادات، استشارة الطبيب إذا تجاوزت 39 درجة",
        "emergency": "استشارة طبية فورية إذا تجاوزت 40 درجة أو صاحبتها أعراض شديدة"
    },
    "السعال": {
        "symptoms": ["كحة جافة أو مع بلغم", "احتقان في الحلق"],
        "treatment": "شرب الماء الدافئ مع العسل، استخدام مرطب الهواء، تجنب المثيرات، الراحة وتناول الأدوية المناسبة",
        "prevention": "تجنب التدخين، تجنب الغبار والملوثات"
    }
}

# Initialize models
print("Loading models...")
try:
    # Try to load Arabic BERT model
    model_id = "aubmindlab/bert-base-arabertv2"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForSequenceClassification.from_pretrained(model_id)
    print("Arabic model loaded successfully")
except Exception as e:
    print(f"Warning: Could not load Arabic model: {e}")
    tokenizer = None
    model = None

try:
    # Load sentence transformer for semantic search
    embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
    print("Sentence transformer loaded successfully")
except Exception as e:
    print(f"Error loading sentence transformer: {e}")
    embedder = None

# Try to load free Hugging Face models for medical questions
try:
    # Option 1: Arabic GPT model (best for Arabic medical questions)
    llm_pipeline = pipeline(
        "text-generation",
        model="aubmindlab/aragpt2-base",  # Free Arabic GPT model
        device=0 if torch.cuda.is_available() else -1,
        max_length=512,
        pad_token_id=50256
    )
    print("Arabic GPT model loaded successfully")
except Exception as e:
    print(f"Could not load Arabic model, trying alternative: {e}")
    try:
        # Option 2: Small English model (fallback)
        llm_pipeline = pipeline(
            "text-generation",
            model="distilgpt2",  # Small, fast, free model
            device=0 if torch.cuda.is_available() else -1,
            max_length=256
        )
        print("DistilGPT2 model loaded successfully")
    except Exception as e2:
        print(f"Could not load any text generation model: {e2}")
        # Option 3: Try question-answering model as fallback
        try:
            llm_pipeline = pipeline(
                "question-answering",
                model="distilbert-base-cased-distilled-squad",
                device=0 if torch.cuda.is_available() else -1
            )
            print("Question-answering model loaded as fallback")
        except Exception as e3:
            print(f"Warning: Could not load any LLM pipeline: {e3}")
            llm_pipeline = None

# Prepare passages from hospital data
passages = []
for _, row in df.iterrows():
    if pd.isna(row['الدكتور']) and pd.isna(row['العيادات']):
        continue
    text = f"العيادة: {row['العيادات']}، الدكتور: {row['الدكتور']}، الأيام: {row['الأيام']}، السعر: {row['سعر الكشف']}"
    passages.append(text)

# Create embeddings for passages
if embedder:
    try:
        corpus_embeddings = embedder.encode(passages, convert_to_tensor=True)
        print("Corpus embeddings created successfully")
    except Exception as e:
        print(f"Error creating embeddings: {e}")
        corpus_embeddings = None
else:
    corpus_embeddings = None

def generate_medical_advice_with_llm(question: str) -> Optional[str]:
    """
    Generate medical advice using free Hugging Face models
    """
    try:
        if llm_pipeline:
            # Check what type of pipeline we have
            if hasattr(llm_pipeline, 'task') and llm_pipeline.task == 'question-answering':
                # Using QA model - need context
                medical_context = """
                الرعاية الصحية العامة تشمل النظافة الشخصية، التغذية المتوازنة، شرب الماء بكثرة، 
                النوم الكافي، التمارين الرياضية المنتظمة، وتجنب التدخين. للأعراض البسيطة مثل الصداع 
                يمكن الراحة وشرب الماء، وللحمى يمكن استخدام الكمادات الباردة وشرب السوائل.
                """
                response = llm_pipeline(question=question, context=medical_context)
                return response['answer']
            
            else:
                # Using text generation model
                # Create Arabic medical prompt
                if "arabic" in str(llm_pipeline.model).lower() or "arab" in str(llm_pipeline.model).lower():
                    prompt = f"""سؤال طبي: {question}
النصيحة الطبية: """
                else:
                    # English prompt for English models
                    prompt = f"""Medical Question: {question}
Medical Advice: """
                
                # Generate response
                response = llm_pipeline(
                    prompt, 
                    max_length=len(prompt.split()) + 100,  # Dynamic length
                    num_return_sequences=1,
                    temperature=0.7,
                    do_sample=True,
                    pad_token_id=llm_pipeline.tokenizer.eos_token_id
                )
                
                # Extract generated text
                generated_text = response[0]['generated_text']
                answer = generated_text.replace(prompt, "").strip()
                
                # If it's in English, add Arabic translation note
                if not any(ord(char) > 127 for char in answer):  # Check if contains Arabic
                    answer = f"{answer}\n\n(تم إنشاء هذه الإجابة باللغة الإنجليزية من نموذج الذكاء الاصطناعي)"
                
                return answer
            
    except Exception as e:
        print(f"Error generating LLM response: {e}")
    
    return None

def get_comprehensive_medical_answer(question: str) -> str:
    """
    Enhanced function to provide comprehensive medical answers
    """
    # Check for symptoms in the enhanced knowledge base
    for symptom, info in medical_kb.items():
        if symptom in question or any(s in question for s in info.get("symptoms", [])):
            response = f"🩺 **معلومات طبية عن {symptom}:**\n\n"
            
            if "symptoms" in info:
                response += f"**الأعراض الشائعة:** {', '.join(info['symptoms'])}\n\n"
            
            if "treatment" in info:
                response += f"**العلاج والرعاية:** {info['treatment']}\n\n"
            
            if "prevention" in info:
                response += f"**الوقاية:** {info['prevention']}\n\n"
            
            if "emergency" in info:
                response += f"⚠️ **تحذير:** {info['emergency']}\n\n"
            
            response += "⚠️ **ملاحظة هامة:** هذه النصائح عامة، يرجى استشارة الطبيب للتشخيص الدقيق والعلاج المناسب."
            return response
    
    # If not found in knowledge base, try to generate answer with LLM
    llm_response = generate_medical_advice_with_llm(question)
    if llm_response:
        return f"🤖 **استشارة طبية ذكية:**\n\n{llm_response}\n\n⚠️ **تنبيه:** هذه نصائح عامة من الذكاء الاصطناعي وليست بديلاً عن الاستشارة الطبية المتخصصة."
    
    # Fallback to general medical advice
    return """
    🩺 **استشارة طبية عامة:**
    
    عذراً، لا يمكنني تقديم معلومات محددة عن هذا السؤال الطبي. 
    
    **نصائح عامة للصحة:**
    • اشرب كمية كافية من الماء يومياً
    • تناول غذاء متوازن غني بالفواكه والخضروات
    • احصل على قسط كافٍ من النوم (7-8 ساعات)
    • مارس التمارين الرياضية بانتظام
    • تجنب التدخين والكحول
    
    ⚠️ **للحصول على استشارة طبية دقيقة، يرجى:**
    • زيارة طبيب مختص
    • الاتصال بخط الاستشارة الطبية
    • زيارة أقرب عيادة أو مستشفى
    """

def is_medical_question(question: str) -> bool:
    """Check if the question is medical-related"""
    medical_indicators = [
        "مرض", "علاج", "دواء", "ألم", "وجع", "أعراض", "صحة", "طبي", "مريض",
        "اشعر", "اعاني", "يؤلمني", "يوجعني", "حمى", "صداع", "سعال", "إسهال",
        "إمساك", "غثيان", "دوخة", "تعب", "إعياء", "التهاب", "عدوى", "فيروس",
        "جرثومة", "حساسية", "ضغط", "سكر", "قلب", "معدة", "كبد", "كلى"
    ]
    return any(indicator in question for indicator in medical_indicators)

def check_lab_or_radiology(question):
    """Enhanced lab and radiology checking"""
    keywords = ['أشعة', 'تحاليل', 'رنين', 'سونار', 'تحليل', 'مختبر', 'معمل', 'فحص']
    if any(word in question for word in keywords):
        try:
            matches = get_close_matches(question, qa_data['السؤال'].tolist(), n=1, cutoff=0.3)
            if matches:
                matched_q = matches[0]
                answer = qa_data[qa_data['السؤال'] == matched_q]['الجواب'].values[0]
                return answer
        except Exception as e:
            print(f"Error in lab/radiology check: {e}")
    return None

def answer_question_from_excel(user_question, k=3):
    """Enhanced question answering with better medical support"""
    
    # Check if it's a medical question first
    if is_medical_question(user_question):
        medical_answer = get_comprehensive_medical_answer(user_question)
        if "استشارة طبية عامة" not in medical_answer:
            return medical_answer
    
    # Check in qa_data for lab tests and equipment
    qa_answer = check_lab_or_radiology(user_question)
    if qa_answer:
        return f"💡 **الرد:** {qa_answer}"

    # Enhanced clinic mapping with more variations
    clinic_mapping = {
        "عيادة الباطنة": ["باطنة", "الباطنة", "باطنه", "داخلية", "امراض باطنة"],
        "عيادة الجلدية": ["جلد", "جلدية", "الجلدية", "امراض جلدية", "تناسلية"],
        "عيادة المسالك": ["مسالك", "المسالك", "مسالك بولية", "بولية", "تناسلية"],
        "عيادة النساء": ["نسا", "نساء", "توليد", "ولادة", "النساء والتوليد"],
        "عيادة الأنف والأذن": ["انف", "اذن", "حنجرة", "الأنف", "الأذن", "أنف وأذن"],
        "عيادة الرمد": ["رمد", "الرمد", "عيون", "العيون", "بصريات"],
        "عيادة الأسنان": ["اسنان", "الاسنان", "الأسنان", "سنان", "فم"],
        "الحضانة": ["حضان", "حضانة", "الحضانة", "حضانات", "عناية مركزة"],
        "المعمل": ["معمل", "المعمل", "تحاليل", "مختبر"],
        "عيادة الأطفال": ["اطفال", "الاطفال", "الأطفال", "طب الاطفال"],
        "عيادة الجهاز الهضمي": ["هضم", "هضمي", "الجهاز الهضمي", "معدة", "أمعاء"],
        "عيادة التخاطب": ["تخاطب", "التخاطب", "نطق", "كلام"],
        "عيادة العظام": ["عظام", "العظام", "مفاصل", "كسور", "روماتيزم"],
        "عيادة القلب": ["قلب", "القلب", "قلبية", "دورة دموية"],
        "عيادة الأعصاب": ["أعصاب", "الأعصاب", "عصبية", "مخ وأعصاب"]
    }

    # Check for specific clinic mentions
    target_clinic = None
    for clinic, keywords in clinic_mapping.items():
        if any(keyword in user_question for keyword in keywords):
            target_clinic = clinic
            break

    # If a specific clinic is mentioned, show all information for that clinic
    if target_clinic:
        clinic_info = [p for p in passages if target_clinic in p]
        if clinic_info:
            return f"🏥 **معلومات العيادة:** \n" + "\n".join(clinic_info)

    # For questions about doctors or clinics
    if any(word in user_question for word in ["عيادة", "عيادات", "دكتور", "دكاترة", "دكتورة", "مواعيد", "طبيب"]):
        return f"🏥 **جميع العيادات المتاحة:** \n" + "\n".join([p for p in passages if 'عيادة' in p])
    
    # Use semantic search if available
    if embedder and corpus_embeddings is not None:
        try:
            question_embedding = embedder.encode(user_question, convert_to_tensor=True)
            hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=k)
            matched_passages = [passages[hit['corpus_id']] for hit in hits[0]]
            
            if matched_passages:
                hospital_info = "\n".join(matched_passages)
                return f"💡 **الرد:** {hospital_info}"
        except Exception as e:
            print(f"Error in semantic search: {e}")
    
    # If it's a medical question but not covered above, try LLM
    if is_medical_question(user_question):
        return get_comprehensive_medical_answer(user_question)
    
    # Default response
    return f"💡 **الرد:** عذراً، لا يمكنني العثور على إجابة محددة لسؤالك. يرجى التواصل مع المستشفى مباشرة للحصول على المعلومات المطلوبة."

def ask_question(question):
    """Main function to handle question answering with comprehensive error handling"""
    try:
        if not question or question.strip() == "":
            return "❌ خطأ: السؤال مطلوب"
        
        question = question.strip()
        
        # First check lab/radiology
        answer = check_lab_or_radiology(question)
        if answer:
            return f"💡 **الرد:** {answer}"
        
        # Then use the comprehensive answering system
        return answer_question_from_excel(question)
        
    except Exception as e:
        return f"❌ خطأ: {str(e)}"

# Enhanced Gradio interface
with gr.Blocks(
    title="نظام الأسئلة والأجوبة الطبي المتطور",
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        font-family: 'Arial', sans-serif;
        direction: rtl;
        text-align: right;
    }
    .gr-button {
        background: linear-gradient(45deg, #4CAF50, #45a049);
        color: white;
        font-weight: bold;
        border-radius: 10px;
    }
    .gr-textbox {
        text-align: right;
        direction: rtl;
        border-radius: 10px;
    }
    .header-style {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        padding: 20px;
        border-radius: 15px;
        margin-bottom: 20px;
    }
    """
) as demo:
    
    gr.Markdown(
        """
        # 🏥 نظام الأسئلة والأجوبة الطبي المتطور
        ### اسأل أي سؤال طبي واحصل على إجابة ذكية ومفصلة!
        
        **🔥 الميزات الجديدة:**
        - 🤖 ذكاء اصطناعي متقدم للأسئلة الطبية
        - 📚 قاعدة معرفة طبية شاملة
        - 🏥 معلومات مفصلة عن العيادات والأطباء
        - 🧪 تفاصيل كاملة عن التحاليل والفحوصات
        
        **يمكنك السؤال عن:**
        - 🩺 الأعراض والأمراض الشائعة
        - 💊 النصائح العلاجية والوقائية
        - 🏥 مواعيد العيادات والأطباء
        - 🧪 التحاليل الطبية والمعمل
        - 📋 أسعار الكشف والخدمات
        - ⚕️ أي استفسار طبي عام
        """,
        elem_id="header"
    )
    
    with gr.Row():
        with gr.Column(scale=4):
            question_input = gr.Textbox(
                label="اكتب سؤالك الطبي هنا",
                placeholder="مثال: ما علاج الصداع؟ أو ما مواعيد عيادة الأطفال؟",
                lines=4,
                text_align="right",
                rtl=True
            )
        with gr.Column(scale=1):
            submit_btn = gr.Button(
                "🔍 اسأل السؤال", 
                variant="primary",
                size="lg"
            )
    
    answer_output = gr.Textbox(
        label="الإجابة الطبية",
        lines=15,
        interactive=False,
        text_align="right",
        rtl=True
    )
    
    # Enhanced example questions
    gr.Examples(
        examples=[
            ["ما علاج الحمى عند الأطفال؟"],
            ["أعاني من صداع مستمر، ما النصيحة؟"],
            ["ما هي أعراض الجفاف وكيف أعالجه؟"],
            ["ما هي مواعيد عيادة القلب؟"],
            ["هل يوجد تحليل فيروس كورونا؟"],
            ["مين أفضل دكتور عظام في المستشفى؟"],
            ["ما هي تحاليل وظائف الكبد المطلوبة؟"],
            ["كيف أعالج السعال الجاف؟"]
        ],
        inputs=question_input,
        label="أمثلة على الأسئلة الطبية"
    )
    
    # Set up the interaction
    submit_btn.click(
        fn=ask_question,
        inputs=question_input,
        outputs=answer_output
    )
    
    # Also allow Enter key to submit
    question_input.submit(
        fn=ask_question,
        inputs=question_input,
        outputs=answer_output
    )
    
    # Add disclaimer
    gr.Markdown(
        """
        ---
        ⚠️ **إخلاء مسؤولية:** هذا النظام يقدم معلومات طبية عامة فقط وليس بديلاً عن الاستشارة الطبية المتخصصة. 
        يرجى استشارة طبيب مختص للحصول على تشخيص وعلاج دقيق.
        """,
        elem_id="disclaimer"
    )

if __name__ == "__main__":
    print("Starting Enhanced Hospital Q&A System...")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        debug=True,
        show_error=True
    )