import gradio as gr
import google.generativeai as genai
import os
from datetime import datetime

# Set your Gemini API Key securely
os.environ["GOOGLE_API_KEY"] = "AIzaSyDwMK08RcUvg9JomaJSuX9XT3BGNasQk5M"

# Configure Gemini client
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Using Gemini 2.5 Flash
MODEL_ID = "gemini-2.0-flash-exp"

# System prompt for the AI Doctor
DOCTOR_SYSTEM_PROMPT = """You are Dr. AI Assistant, a professional and empathetic medical doctor conducting a thorough patient consultation. 

YOUR CONSULTATION PROCESS:
1. **Initial Greeting & Personal Details Collection**:
   - Greet the patient warmly
   - Ask for: Name, Age, Gender, Contact details
   - Ask about occupation and lifestyle

2. **Medical History**:
   - Ask about chief complaints (main symptoms)
   - Duration and severity of symptoms
   - Past medical history
   - Current medications
   - Allergies
   - Family medical history
   - Previous surgeries or hospitalizations

3. **Detailed Symptom Analysis**:
   - Ask specific questions about each symptom
   - Pain level (1-10 scale)
   - Aggravating and relieving factors
   - Associated symptoms

4. **Lifestyle Assessment**:
   - Diet and eating habits
   - Exercise routine
   - Sleep patterns
   - Stress levels
   - Smoking/alcohol consumption
   - Water intake

5. **Physical Examination Questions** (virtual):
   - Ask about vital signs if known (BP, temperature, pulse)
   - Visible symptoms
   - Any recent changes in body

6. **Provide Comprehensive Recommendations**:
   - Preliminary assessment
   - Dietary recommendations
   - Lifestyle modifications
   - Exercise suggestions
   - Medication suggestions (with disclaimer)
   - When to seek immediate medical attention
   - Follow-up advice

7. **Generate Summary** (when patient requests):
   - Patient details
   - Chief complaints
   - Medical history
   - Assessment
   - Recommendations (diet, lifestyle, medications)
   - Disclaimer about consulting a physical doctor

IMPORTANT GUIDELINES:
- Be professional, empathetic, and reassuring
- **CRITICAL: Ask ONLY ONE question at a time - NEVER ask multiple questions in a single response**
- Wait for the patient's answer before asking the next question
- Use medical terminology but explain in simple terms
- Always include disclaimer that this is preliminary advice
- Recommend seeing a physical doctor for serious conditions
- Show concern and empathy
- Be thorough but not overwhelming
- Keep the conversation natural and flowing like a real doctor-patient interaction

DISCLAIMER TO INCLUDE:
"⚠️ Please note: This is an AI-powered preliminary consultation. For accurate diagnosis and treatment, please consult a licensed medical professional in person. In case of emergency, call emergency services immediately."

Start by greeting the patient and asking for their name."""

def initialize_chat():
    """Initialize the Gemini chat with doctor system prompt"""
    model = genai.GenerativeModel(
        model_name=MODEL_ID,
        system_instruction=DOCTOR_SYSTEM_PROMPT
    )
    chat = model.start_chat(history=[])
    return chat

# Global chat instance
chat_session = initialize_chat()

def chat_with_doctor(message, history):
    """Handle patient messages and get doctor's response"""
    global chat_session
    
    try:
        # Send patient message to doctor
        response = chat_session.send_message(message)
        doctor_response = response.text
        
        return doctor_response
    
    except Exception as e:
        return f"I apologize, but I encountered an error: {str(e)}. Please try again or restart the consultation."

def generate_consultation_summary(history):
    """Generate a comprehensive summary of the consultation"""
    global chat_session
    
    summary_request = """Please generate a comprehensive consultation summary in the following format:

📋 **MEDICAL CONSULTATION SUMMARY**
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

**Patient Information:**
- Name: [name]
- Age: [age]
- Gender: [gender]
- Date of Consultation: [current date]

**Chief Complaints:**
[List main symptoms and concerns]

**Medical History:**
[Relevant medical history]

**Assessment:**
[Your preliminary assessment]

**Recommendations:**

1. **Dietary Recommendations:**
   [Specific diet advice]

2. **Lifestyle Modifications:**
   [Lifestyle changes needed]

3. **Exercise & Physical Activity:**
   [Exercise recommendations]

4. **Medication Suggestions:**
   [If any - with disclaimer]

5. **Follow-up Care:**
   [When to seek medical attention]

**Important Notes:**
⚠️ This is a preliminary AI consultation. Please consult a licensed medical doctor for accurate diagnosis and treatment.

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Generated by Dr. AI Assistant"""
    
    try:
        response = chat_session.send_message(summary_request)
        return response.text
    except Exception as e:
        return f"Error generating summary: {str(e)}"

def restart_consultation():
    """Restart the consultation with a fresh session"""
    global chat_session
    chat_session = initialize_chat()
    return [], "Consultation restarted. Please tell me your name to begin."

# Create Gradio Interface
with gr.Blocks(title="AI Doctor Consultation", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🩺 AI Doctor Consultation System
    ### Professional Medical Consultation with Dr. AI Assistant
    
    Welcome! I'm Dr. AI Assistant, here to conduct a thorough medical consultation. 
    I'll ask about your symptoms, medical history, lifestyle, and provide comprehensive recommendations.
    
    **⚠️ Important:** This is a preliminary consultation. Always consult a licensed physician for proper diagnosis and treatment.
    """)
    
    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                height=500,
                show_label=False,
                avatar_images=(None, "https://api.dicebear.com/7.x/bottts/svg?seed=doctor"),
                bubble_full_width=False
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    placeholder="Type your message here...",
                    show_label=False,
                    scale=4
                )
                send_btn = gr.Button("Send 📤", scale=1, variant="primary")
            
            with gr.Row():
                summary_btn = gr.Button("📋 Generate Consultation Summary", variant="secondary")
                clear_btn = gr.Button("🔄 Restart Consultation", variant="stop")
        
        with gr.Column(scale=1):
            gr.Markdown("""
            ### 📌 Consultation Guide
            
            **What I'll ask about:**
            - Personal details
            - Chief complaints
            - Medical history
            - Current medications
            - Lifestyle habits
            - Diet & exercise
            
            **I'll provide:**
            - Preliminary assessment
            - Dietary advice
            - Lifestyle changes
            - Exercise plan
            - Medication suggestions
            - Follow-up guidance
            
            **Commands:**
            - Type normally to chat
            - Click "Generate Summary" for full report
            - Click "Restart" for new consultation
            
            ⚕️ *Be honest and detailed for best advice*
            """)
            
            summary_output = gr.Textbox(
                label="📋 Consultation Summary",
                placeholder="Click 'Generate Summary' after consultation",
                lines=15,
                max_lines=20
            )
    
    # Auto-start conversation
    def start_conversation():
        initial_message = "Hello! I'm Dr. AI Assistant. I'm here to help you today. Let's start with some basic information.\n\n👤 May I have your name, please?"
        return [[None, initial_message]]
    
    demo.load(
        start_conversation,
        outputs=[chatbot]
    )
    
    # Handle send message
    def respond(message, chat_history):
        if not message.strip():
            return "", chat_history
        
        # Add user message
        chat_history.append((message, None))
        
        # Get doctor response
        bot_response = chat_with_doctor(message, chat_history)
        chat_history[-1] = (message, bot_response)
        
        return "", chat_history
    
    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    send_btn.click(respond, [msg, chatbot], [msg, chatbot])
    
    # Generate summary
    summary_btn.click(
        lambda history: generate_consultation_summary(history),
        inputs=[chatbot],
        outputs=[summary_output]
    )
    
    # Restart consultation
    def restart_and_update():
        new_history, greeting = restart_consultation()
        return [[None, greeting]], ""
    
    clear_btn.click(
        restart_and_update,
        outputs=[chatbot, summary_output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(share=True)