import gradio as gr
from groq import GroqClient

# ==============================
# Initialize Groq client
# ==============================
client = GroqClient(api_key="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")

# ==============================
# System Prompt for Doctor
# ==============================
SYSTEM_PROMPT = """
You are Dr. HealBot, a calm, knowledgeable, and empathetic doctor talking to a patient.

GOAL:
Have a natural conversation — ask 3-4 short medical questions to understand the patient's condition, 
then start giving practical advice including:
- possible over-the-counter medicines (generic name only)
- simple lifestyle or habit changes
- nutrition or exercise guidance
- when to see a real doctor

TONE & STYLE:
- Speak like a real doctor, short and direct sentences (1-2 lines max).
- Be warm but professional.
- Use plain language — no medical jargon unless necessary.
- No bullet points or lists — just natural speech.
- Only one question per response, until enough info is gathered.
- After about 4 patient answers, switch to giving advice.

CONVERSATION FLOW EXAMPLE:
Doctor: How can I help you?
Patient: I’ve had a cough for 2 weeks.
Doctor: Is it dry or with phlegm?
Patient: With phlegm.
Doctor: Do you have fever or chest pain?
Patient: Mild fever.
Doctor: Do you smoke or have allergies?
Patient: I smoke.
Doctor: Sounds like a mild chest infection. You can try paracetamol for fever and warm fluids. 
Cut down on smoking and rest. If symptoms persist beyond 5 days, see a doctor.

ALWAYS END with a gentle reminder: 
"Please consult a qualified doctor if it doesn’t improve or if symptoms worsen."
"""

# ==============================
# Initial greeting
# ==============================
INITIAL_MESSAGE = "How can I help you today?"

# ==============================
# Chat logic
# ==============================
def chat_with_doctor(message, history):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    # Build chat history
    for chat in history:
        if isinstance(chat, dict):
            messages.append(chat)
        elif isinstance(chat, (list, tuple)) and len(chat) == 2:
            if chat[0]:
                messages.append({"role": "user", "content": chat[0]})
            if chat[1]:
                messages.append({"role": "assistant", "content": chat[1]})

    # Add current patient message
    messages.append({"role": "user", "content": message})

    try:
        # Count how many patient turns have occurred
        patient_turns = sum(1 for chat in history if isinstance(chat, (list, tuple)) and chat[0])

        # After 4 patient turns, guide the model to provide recommendations
        if patient_turns >= 4:
            messages.append({
                "role": "system",
                "content": (
                    "Now begin giving specific recommendations based on the patient's symptoms. "
                    "Include possible generic medicines (like paracetamol, ibuprofen, etc.), "
                    "lifestyle and nutrition tips, and when to seek medical attention. "
                    "Keep it short and empathetic, like a real doctor speaking naturally."
                )
            })

        # Generate the response using Groq LLM
        chat_completion = client.chat.completions.create(
            messages=messages,
            model="llama-3.3-70b-versatile",
            temperature=0.6,
            max_tokens=120,
        )

        response = chat_completion.choices[0].message.content.strip()

        # Append to history
        history.append([message, response])
        return history

    except Exception as e:
        error_msg = f"⚠️ Error: {str(e)}. Please check your API connection and try again."
        history.append([message, error_msg])
        return history


def reset_conversation():
    """Reset the chat to start fresh"""
    return []

# ==============================
# Custom CSS
# ==============================
custom_css = """
#chatbot {
    height: 600px;
}
.gradio-container {
    font-family: 'Arial', sans-serif;
}
#warning {
    background-color: #fff3cd;
    border: 1px solid #ffc107;
    border-radius: 8px;
    padding: 15px;
    margin: 10px 0;
    color: #856404;
}
"""

# ==============================
# Gradio Interface
# ==============================
with gr.Blocks(css=custom_css, title="AI Medical Consultant") as demo:
    gr.Markdown(
        """
        # 🏥 AI Medical Consultant
        ### Realistic Doctor-Patient Conversation • Medicine • Lifestyle • Nutrition
        """
    )

    gr.HTML(
        """
        <div id="warning">
            <strong>⚠️ Medical Disclaimer:</strong><br>
            This AI provides general health information only. It is <b>NOT</b> a substitute for
            professional medical advice, diagnosis, or treatment.<br>
            Always consult qualified healthcare providers for medical concerns.<br>
            For emergencies, call your local emergency number immediately.
        </div>
        """
    )

    chatbot = gr.Chatbot(
        value=[[None, INITIAL_MESSAGE]],
        elem_id="chatbot",
        height=600,
        show_label=False,
        avatar_images=(None, "https://api.dicebear.com/7.x/bottts/svg?seed=doctor"),
        type="tuples"
    )

    with gr.Row():
        msg = gr.Textbox(
            placeholder="Describe your symptoms or ask a question...",
            show_label=False,
            scale=9,
            lines=2
        )
        submit_btn = gr.Button("Send 📤", scale=1, variant="primary")

    with gr.Row():
        clear_btn = gr.Button("🔄 Start New Consultation", variant="secondary")

    gr.Markdown(
        """
        ### 💡 Tips for Best Results:
        - Be specific about your symptoms (location, severity, duration)
        - Mention any relevant medical history or medications
        - Ask follow-up questions freely
        """
    )

    # Event Handlers
    msg.submit(chat_with_doctor, [msg, chatbot], [chatbot]).then(
        lambda: gr.update(value=""), None, [msg]
    )
    submit_btn.click(chat_with_doctor, [msg, chatbot], [chatbot]).then(
        lambda: gr.update(value=""), None, [msg]
    )
    clear_btn.click(reset_conversation, None, [chatbot]).then(
        lambda: [[None, INITIAL_MESSAGE]], None, [chatbot]
    )

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