import gradio as gr
from transformers import pipeline

# Load LLaMA model via Hugging Face pipeline
# You can swap this with any LLaMA model hosted on Hugging Face
# Example: "meta-llama/Llama-2-7b-chat-hf" if you're authenticated
chatbot = pipeline("text-generation", model="tiiuae/falcon-7b-instruct", max_new_tokens=200)

# Disclaimer text
disclaimer = "⚠️ **Disclaimer:** Doctor Twin is an AI health companion and not a substitute for professional medical advice, diagnosis, or treatment. Always consult a qualified healthcare provider for any medical concerns."

# Simulated response handler
def doctor_twin(user_input):
    if not user_input.strip():
        return "Hi! Please describe your symptom or health concern so I can help you better."

    prompt = f"You are Doctor Twin, a friendly virtual health assistant. Greet the user and respond empathetically with general advice. User says: {user_input}"
    
    response = chatbot(prompt)[0]['generated_text'].split("User says:")[1].strip()
    
    # Basic recommendations (simulated)
    recommendations = []
    if any(word in user_input.lower() for word in ['cold', 'cough', 'flu']):
        recommendations = ["Stay hydrated", "Get plenty of rest", "Use warm fluids", "See a doctor if fever persists more than 3 days"]
    elif "headache" in user_input.lower():
        recommendations = ["Drink water", "Rest in a dark, quiet room", "Avoid screen time", "Consult a doctor if persistent"]
    elif "sleep" in user_input.lower():
        recommendations = ["Maintain regular sleep schedule", "Avoid caffeine before bed", "Try relaxation techniques"]
    
    return response, recommendations

# Gradio interface
with gr.Blocks(title="Doctor Twin – Your Friendly Health Companion") as demo:
    gr.Markdown("# 🩺 Doctor Twin – Your Friendly Health Companion")
    gr.Markdown("**Helping you understand symptoms and stay informed – powered by AI.**")
    gr.Markdown(disclaimer)

    with gr.Row():
        with gr.Column():
            chatbot_input = gr.Textbox(label="Describe your symptoms or concern", placeholder="e.g., I have a sore throat and mild fever...")
            submit_btn = gr.Button("Ask Doctor Twin")

        with gr.Column():
            chat_output = gr.Textbox(label="Doctor Twin Says", lines=6)
            recommendation_output = gr.HighlightedText(label="Suggested Recommendations", combine_adjacent=True)

    submit_btn.click(doctor_twin, chatbot_input, outputs=[chat_output, recommendation_output])

# Run the app
if __name__ == "__main__":
    demo.launch()