import gradio as gr
from llama_cpp import Llama
import os

# Initialize the model
print("Loading model...")
llm = Llama(
    model_path="Medical-Guide-COT-llama3.2-1B.Q5_K_M.gguf",
    n_ctx=2048,
    n_threads=4,
    n_gpu_layers=0  # CPU only for Spaces
)

def chat(message, history):
    """
    Chat function that maintains conversation history
    """
    # Build the prompt with conversation history
    prompt = ""
    
    # Add system message
    prompt += "You are a helpful medical assistant. Provide informative responses about medical topics, but always remind users to consult healthcare professionals for actual medical advice.\n\n"
    
    # Add conversation history
    for user_msg, bot_msg in history:
        prompt += f"User: {user_msg}\n"
        prompt += f"Assistant: {bot_msg}\n"
    
    # Add current message
    prompt += f"User: {message}\n"
    prompt += "Assistant:"
    
    # Generate response
    response = llm(
        prompt,
        max_tokens=512,
        temperature=0.7,
        top_p=0.9,
        stop=["User:", "\n\n\n"],
        echo=False
    )
    
    bot_message = response['choices'][0]['text'].strip()
    
    return bot_message

# Custom CSS
custom_css = """
.gradio-container {
    max-width: 900px !important;
}
#chatbot {
    height: 500px !important;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🏥 Medical Guide Chatbot
        
        This chatbot uses the Medical-Guide-COT-llama3.2-1B model to answer medical questions.
        
        ⚠️ **Disclaimer**: This is for informational purposes only. Always consult with qualified healthcare professionals for medical advice.
        """
    )
    
    chatbot = gr.Chatbot(
        elem_id="chatbot",
        bubble_full_width=False,
        avatar_images=(None, "🏥")
    )
    
    with gr.Row():
        msg = gr.Textbox(
            placeholder="Ask a medical question...",
            show_label=False,
            scale=4
        )
        submit = gr.Button("Send", scale=1, variant="primary")
    
    with gr.Row():
        clear = gr.Button("Clear Conversation")
    
    gr.Markdown(
        """
        ### Example Questions:
        - What are the symptoms of diabetes?
        - How can I improve my sleep quality?
        - What is hypertension and how is it treated?
        """
    )
    
    # Event handlers
    msg.submit(chat, inputs=[msg, chatbot], outputs=[chatbot])
    submit.click(chat, inputs=[msg, chatbot], outputs=[chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)
    msg.submit(lambda: "", None, msg)
    submit.click(lambda: "", None, msg)

if __name__ == "__main__":
    demo.launch()