import gradio as gr
import os
from groq import Groq
from datetime import datetime
from reportlab.lib.pagesizes import A4
from reportlab.lib.units import inch
from reportlab.lib import colors
from reportlab.platypus import SimpleDocTemplate, Table, TableStyle, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.enums import TA_CENTER, TA_JUSTIFY, TA_RIGHT
import tempfile

# ==============================
# Initialize Groq client
# ==============================
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# ==============================
# Language Configuration
# ==============================
LANGUAGES = {
    "english": {
        "code": "en",
        "name": "English",
        "system_prompt": """
You are Dr. HealBot, a calm, knowledgeable, and empathetic virtual doctor.

GOAL:
Hold a natural, focused conversation with the patient to understand their health issue and offer helpful preliminary medical guidance.

CONVERSATION LOGIC:
- Ask only relevant and concise medical questions necessary for diagnosing the illness.
- Each question should help clarify symptoms or narrow possible causes.
- Stop asking once enough information is collected for a basic assessment.
- Then, provide a structured, friendly, and visually clear medical response using headings, emojis, and bullet points.

FINAL RESPONSE FORMAT:
🩺 **Based on what you've told me...**
Brief summary of what the patient described.

💡 **Possible Causes (Preliminary)**
- List 1–2 possible conditions using phrases like "It could be" or "This sounds like".

💊 **Suggested Over-the-Counter Medicines**
- Generic medicine names only.
- Mention to check packaging or consult a pharmacist for dosage confirmation.

🥗 **Lifestyle & Home Care Tips**
- 2–3 practical suggestions.

⚠️ **When to See a Real Doctor**
- 2–3 warning signs or conditions when urgent medical care is needed.

📅 **Follow-Up Advice**
- Brief recommendation for self-care or follow-up timing.

TONE & STYLE:
- Speak like a real, caring doctor — short, clear, and empathetic.
- Use plain language, no jargon.
- Only one question per turn unless clarification is essential.
- Keep tone warm, calm, and professional.

IMPORTANT:
- Always emphasize that this is preliminary guidance.
- Never make definitive diagnoses.
- For severe symptoms, immediately advise seeking emergency care.
""",
        "initial_message": "👋 Hello! I'm **Dr. HealBot**. How can I help you today? Please describe your main concern or symptom.",
        "consultation_complete": "✅ **Consultation complete!** You can now export this as a PDF or start a new consultation.",
        "placeholder": "💬 Describe your symptoms (e.g., 'I have a fever and headache for 3 days')...",
        "send_btn": "Send 📤",
        "new_consultation": "🔄 New Consultation",
        "export_pdf": "💾 Export PDF",
        "title": "Dr. HealBot - AI Medical Consultant",
        "subtitle": "Your AI Medical Consultation Assistant",
        "pdf_label": "📄 Download Consultation Report",
        "language_label": "Language / اللغة",
        "pdf_title": "Dr. HealBot",
        "pdf_subtitle": "AI Medical Consultation Summary",
        "pdf_disclaimer_title": "IMPORTANT MEDICAL DISCLAIMER",
        "pdf_disclaimer_text": "This consultation provides preliminary AI-generated medical guidance and is NOT a substitute for professional medical care. Always consult a licensed healthcare provider for accurate diagnosis and treatment. In case of emergency, call emergency services immediately.",
        "pdf_transcript": "Consultation Transcript",
        "pdf_patient": "Patient:",
        "pdf_doctor": "Dr. HealBot:",
        "pdf_footer1": "This document is confidential and for personal reference only.",
        "pdf_footer2": "© 2025 Dr. HealBot - AI Medical Consultation Tool"
    },
    "arabic": {
        "code": "ar",
        "name": "العربية",
        "system_prompt": """
أنت دكتور هيل بوت، طبيب افتراضي هادئ وذو معرفة واسعة ومتعاطف.

الهدف:
إجراء محادثة طبيعية ومركزة مع المريض لفهم مشكلته الصحية وتقديم إرشادات طبية أولية مفيدة.

منطق المحادثة:
- اطرح فقط الأسئلة الطبية ذات الصلة والموجزة اللازمة لتشخيص المرض.
- يجب أن يساعد كل سؤال في توضيح الأعراض أو تضييق الأسباب المحتملة.
- توقف عن طرح الأسئلة بمجرد جمع معلومات كافية للتقييم الأساسي.
- بعد ذلك، قدم استجابة طبية منظمة وودية وواضحة بصريًا باستخدام العناوين والرموز التعبيرية والنقاط.

صيغة الاستجابة النهائية:
🩺 **بناءً على ما أخبرتني به...**
ملخص موجز لما وصفه المريض.

💡 **الأسباب المحتملة (أولية)**
- اذكر 1-2 من الحالات المحتملة باستخدام عبارات مثل "قد يكون" أو "يبدو أن هذا".

💊 **الأدوية المقترحة بدون وصفة طبية**
- أسماء الأدوية العامة فقط.
- اذكر ضرورة التحقق من العبوة أو استشارة الصيدلي لتأكيد الجرعة.

🥗 **نصائح نمط الحياة والرعاية المنزلية**
- 2-3 اقتراحات عملية.

⚠️ **متى يجب زيارة طبيب حقيقي**
- 2-3 علامات تحذيرية أو حالات تتطلب رعاية طبية عاجلة.

📅 **نصائح المتابعة**
- توصية موجزة للرعاية الذاتية أو توقيت المتابعة.

الأسلوب والنبرة:
- تحدث كطبيب حقيقي ومهتم - قصير وواضح ومتعاطف.
- استخدم لغة بسيطة، بدون مصطلحات معقدة.
- سؤال واحد فقط في كل دورة ما لم يكن التوضيح ضروريًا.
- حافظ على نبرة دافئة وهادئة ومهنية.

مهم:
- أكد دائمًا أن هذا إرشاد أولي.
- لا تقدم تشخيصات نهائية أبدًا.
- للأعراض الشديدة، انصح فورًا بطلب الرعاية الطارئة.

يجب أن تكون جميع إجاباتك باللغة العربية.
""",
        "initial_message": "👋 مرحباً! أنا **دكتور هيل بوت**. كيف يمكنني مساعدتك اليوم؟ من فضلك صف مشكلتك الرئيسية أو العَرَض.",
        "consultation_complete": "✅ **اكتملت الاستشارة!** يمكنك الآن تصدير هذا كملف PDF أو بدء استشارة جديدة.",
        "placeholder": "💬 صف أعراضك (مثال: 'لدي حمى وصداع منذ 3 أيام')...",
        "send_btn": "إرسال 📤",
        "new_consultation": "🔄 استشارة جديدة",
        "export_pdf": "💾 تصدير PDF",
        "title": "دكتور هيل بوت - استشارة طبية بالذكاء الاصطناعي",
        "subtitle": "مساعدك الطبي بالذكاء الاصطناعي",
        "pdf_label": "📄 تحميل تقرير الاستشارة",
        "language_label": "اللغة / Language",
        "pdf_title": "دكتور هيل بوت",
        "pdf_subtitle": "ملخص الاستشارة الطبية بالذكاء الاصطناعي",
        "pdf_disclaimer_title": "إخلاء المسؤولية الطبية المهم",
        "pdf_disclaimer_text": "توفر هذه الاستشارة إرشادات طبية أولية مُنتَجة بالذكاء الاصطناعي وليست بديلاً عن الرعاية الطبية المهنية. استشر دائمًا مقدم رعاية صحية مرخصًا للحصول على تشخيص ودواء دقيق. في حالة الطوارئ، اتصل بخدمات الطوارئ فورًا.",
        "pdf_transcript": "نص الاستشارة",
        "pdf_patient": "المريض:",
        "pdf_doctor": "دكتور هيل بوت:",
        "pdf_footer1": "هذا المستند سري وللاستخدام الشخصي فقط.",
        "pdf_footer2": "© 2025 دكتور هيل بوت - أداة الاستشارة الطبية بالذكاء الاصطناعي"
    }
}

# ==============================
# Chat Logic (Now with per-user state)
# ==============================
def chat_with_doctor(message, history, language):
    """Process user message and generate AI response"""
    if not message.strip():
        return history, ""

    lang_config = LANGUAGES[language]
    
    # Build message history
    messages = [{"role": "system", "content": lang_config["system_prompt"]}]
    for chat in history:
        if isinstance(chat, (list, tuple)) and len(chat) == 2:
            if chat[0] and chat[0] != "(System)" and not chat[0].startswith("✅"):
                messages.append({"role": "user", "content": chat[0]})
            if chat[1] and chat[1] != "(System)" and not chat[1].startswith("✅"):
                messages.append({"role": "assistant", "content": chat[1]})
    
    messages.append({"role": "user", "content": message})

    try:
        # Count patient turns
        patient_turns = sum(1 for chat in history 
                          if isinstance(chat, (list, tuple)) 
                          and chat[0] 
                          and chat[0] != "(System)"
                          and not chat[0].startswith("✅"))
        
        # After 4-5 exchanges, encourage conclusion
        if patient_turns >= 4:
            messages.append({
                "role": "system",
                "content": "You have gathered sufficient information. Provide your complete structured medical assessment now."
            })

        # API call
        chat_completion = client.chat.completions.create(
            messages=messages,
            model="llama-3.3-70b-versatile",
            temperature=0.7,
            max_tokens=800,
            top_p=0.9,
        )

        response = chat_completion.choices[0].message.content.strip()
        history.append([message, response])
        
        # Detect completion
        if "🩺" in response and ("Based on what you've told me" in response or "بناءً على ما أخبرتني به" in response):
            history.append([None, lang_config["consultation_complete"]])
        
        return history, ""  # Return empty string to clear input

    except Exception as e:
        error_message = f"⚠️ **Error:** {str(e)}"
        history.append([message, error_message])
        return history, ""


def reset_conversation(language):
    """Reset conversation"""
    lang_config = LANGUAGES[language]
    return [[None, lang_config["initial_message"]]]


def change_language(language):
    """Change interface language"""
    lang_config = LANGUAGES[language]
    
    return (
        [[None, lang_config["initial_message"]]],  # chatbot
        gr.update(placeholder=lang_config["placeholder"]),  # msg
        gr.update(value=lang_config["send_btn"]),  # submit_btn
        gr.update(value=lang_config["new_consultation"]),  # clear_btn
        gr.update(value=lang_config["export_pdf"]),  # export_btn
        gr.update(label=lang_config["pdf_label"]),  # pdf_output
    )


# ==============================
# PDF Export
# ==============================
def export_conversation_to_pdf(history, language):
    """Generate PDF with Arabic support"""
    if not history or len(history) <= 1:
        return None

    lang_config = LANGUAGES[language]
    is_rtl = (language == "arabic")

    temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", mode='wb')
    pdf_path = temp_pdf.name
    temp_pdf.close()

    doc = SimpleDocTemplate(
        pdf_path,
        pagesize=A4,
        rightMargin=0.75*inch,
        leftMargin=0.75*inch,
        topMargin=0.75*inch,
        bottomMargin=0.75*inch,
    )

    elements = []
    styles = getSampleStyleSheet()
    
    title_align = TA_CENTER
    text_align = TA_RIGHT if is_rtl else TA_JUSTIFY
    
    title_style = ParagraphStyle(
        'CustomTitle',
        parent=styles['Heading1'],
        fontSize=24,
        textColor=colors.HexColor('#1e40af'),
        spaceAfter=6,
        alignment=title_align,
        fontName='Helvetica-Bold'
    )

    subtitle_style = ParagraphStyle(
        'CustomSubtitle',
        parent=styles['Normal'],
        fontSize=11,
        textColor=colors.HexColor('#4b5563'),
        alignment=title_align,
        spaceAfter=20
    )

    heading_style = ParagraphStyle(
        'CustomHeading',
        parent=styles['Heading2'],
        fontSize=13,
        textColor=colors.HexColor('#1e40af'),
        spaceAfter=10,
        spaceBefore=12,
        fontName='Helvetica-Bold',
        alignment=text_align
    )

    content_style = ParagraphStyle(
        'Content',
        parent=styles['Normal'],
        fontSize=9,
        leftIndent=35 if not is_rtl else 0,
        rightIndent=0 if not is_rtl else 35,
        spaceAfter=10,
        alignment=text_align,
    )

    # Header
    elements.append(Paragraph(lang_config["pdf_title"], title_style))
    elements.append(Paragraph(lang_config["pdf_subtitle"], subtitle_style))
    
    # Metadata
    metadata = [
        ['Generated:', datetime.now().strftime('%B %d, %Y at %I:%M %p')],
        ['Document Type:', 'AI Medical Consultation Record'],
    ]
    
    metadata_table = Table(metadata, colWidths=[2.5*inch, 3.5*inch])
    metadata_table.setStyle(TableStyle([
        ('BACKGROUND', (0, 0), (0, -1), colors.HexColor('#f0f4f8')),
        ('FONTNAME', (0, 0), (0, -1), 'Helvetica-Bold'),
        ('FONTSIZE', (0, 0), (-1, -1), 9),
        ('PADDING', (0, 0), (-1, -1), 8),
        ('GRID', (0, 0), (-1, -1), 0.5, colors.HexColor('#e0e0e0')),
    ]))
    elements.append(metadata_table)
    elements.append(Spacer(1, 0.2*inch))

    # Disclaimer
    disclaimer_style = ParagraphStyle(
        'Disclaimer',
        parent=styles['Normal'],
        fontSize=9,
        textColor=colors.HexColor('#dc2626'),
        alignment=title_align,
        fontName='Helvetica-Bold',
        spaceAfter=8
    )
    elements.append(Paragraph(f"⚠️ {lang_config['pdf_disclaimer_title']}", disclaimer_style))
    
    disclaimer_text_style = ParagraphStyle(
        'DisclaimerText',
        parent=styles['Normal'],
        fontSize=8,
        textColor=colors.HexColor('#991b1b'),
        alignment=text_align,
        spaceAfter=15
    )
    elements.append(Paragraph(lang_config['pdf_disclaimer_text'], disclaimer_text_style))

    # Transcript
    elements.append(Paragraph(f"📋 {lang_config['pdf_transcript']}", heading_style))

    for user_msg, bot_msg in history:
        if user_msg and user_msg is not None and not user_msg.startswith("✅"):
            elements.append(Paragraph(f"👤 <b>{lang_config['pdf_patient']}</b>", content_style))
            elements.append(Paragraph(user_msg, content_style))

        if bot_msg and bot_msg is not None and not bot_msg.startswith("✅"):
            elements.append(Paragraph(f"👨‍⚕️ <b>{lang_config['pdf_doctor']}</b>", content_style))
            elements.append(Paragraph(bot_msg, content_style))

    # Footer
    elements.append(Spacer(1, 0.3*inch))
    footer_style = ParagraphStyle(
        'Footer',
        parent=styles['Normal'],
        fontSize=8,
        textColor=colors.HexColor('#6b7280'),
        alignment=title_align,
    )
    elements.append(Paragraph(lang_config['pdf_footer1'], footer_style))
    elements.append(Paragraph(lang_config['pdf_footer2'], footer_style))

    doc.build(elements)
    return pdf_path


def handle_export_pdf(history, language):
    pdf_path = export_conversation_to_pdf(history, language)
    return gr.update(value=pdf_path, visible=True) if pdf_path else gr.update(visible=False)


# ==============================
# CSS
# ==============================
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
@import url('https://fonts.googleapis.com/css2?family=Noto+Sans+Arabic:wght@400;500;600;700&display=swap');

.gradio-container {
    font-family: 'Inter', 'Noto Sans Arabic', sans-serif;
    max-width: 950px !important;
    margin: 0 auto !important;
    padding: 20px !important;
    background-color: #f9fafb;
}

#chatbot {
    height: 60vh !important;
    border-radius: 16px;
    border: 1.5px solid #e5e7eb;
    box-shadow: 0 4px 12px rgba(0,0,0,0.05);
    background: white;
}

.rtl { direction: rtl; text-align: right; }
footer { display: none !important; }

@media (max-width: 768px) {
    #chatbot { height: 70vh !important; }
}
"""

# ==============================
# Gradio Interface with Session State
# ==============================
with gr.Blocks(css=custom_css, title="Dr. HealBot") as demo:
    # CRITICAL: Use gr.State() for per-user session data
    chat_history = gr.State([[None, LANGUAGES["english"]["initial_message"]]])
    
    with gr.Row():
        language_selector = gr.Radio(
            choices=["english", "arabic"],
            value="english",
            label=LANGUAGES["english"]["language_label"],
            interactive=True
        )
    
    title_md = gr.Markdown(
        f"<h1 style='text-align:center; color:#1e40af;'>🏥 {LANGUAGES['english']['title']}</h1>"
        f"<p style='text-align:center; color:#6b7280;'>{LANGUAGES['english']['subtitle']}</p>",
        elem_id="title_section"
    )

    chatbot = gr.Chatbot(
        elem_id="chatbot",
        show_label=False,
        type="tuples",
        bubble_full_width=False,
        render_markdown=True,
        height=500,
        show_copy_button=True,
        rtl=False
    )

    with gr.Row():
        msg = gr.Textbox(
            placeholder=LANGUAGES["english"]["placeholder"],
            show_label=False,
            scale=8,
            lines=2,
        )
        submit_btn = gr.Button(LANGUAGES["english"]["send_btn"], scale=1, variant="primary")

    with gr.Row():
        clear_btn = gr.Button(LANGUAGES["english"]["new_consultation"], variant="secondary")
        export_btn = gr.Button(LANGUAGES["english"]["export_pdf"], variant="secondary")

    pdf_output = gr.File(label=LANGUAGES["english"]["pdf_label"], visible=False)

    # Event handlers - ALL use chat_history State
    language_selector.change(
        change_language,
        inputs=[language_selector],
        outputs=[chat_history, msg, submit_btn, clear_btn, export_btn, pdf_output]
    ).then(
        lambda h: h,
        inputs=[chat_history],
        outputs=[chatbot]
    )
    
    msg.submit(
        chat_with_doctor,
        inputs=[msg, chat_history, language_selector],
        outputs=[chat_history, msg]
    ).then(
        lambda h: h,
        inputs=[chat_history],
        outputs=[chatbot]
    )
    
    submit_btn.click(
        chat_with_doctor,
        inputs=[msg, chat_history, language_selector],
        outputs=[chat_history, msg]
    ).then(
        lambda h: h,
        inputs=[chat_history],
        outputs=[chatbot]
    )
    
    clear_btn.click(
        reset_conversation,
        inputs=[language_selector],
        outputs=[chat_history]
    ).then(
        lambda h: h,
        inputs=[chat_history],
        outputs=[chatbot]
    )
    
    export_btn.click(
        handle_export_pdf,
        inputs=[chat_history, language_selector],
        outputs=[pdf_output]
    )
    
    # Initialize chatbot display from state
    demo.load(
        lambda h: h,
        inputs=[chat_history],
        outputs=[chatbot]
    )

# ==============================
# Launch
# ==============================
if __name__ == "__main__":
    demo.launch(
        share=True,
        show_error=True,
        server_name="0.0.0.0",
        server_port=7860,
    )