# --- START OF FILE app.py ---

import sys
import os
import re
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import login
from dotenv import load_dotenv

# --- FIX: Add project root to Python's path ---
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, project_root)

# --- Updated Spaces import for Zero-GPU compatibility ---
try:
    import spaces
    print("'spaces' module imported successfully.")
except ImportError:
    print("Warning: 'spaces' module not found. Using dummy decorator for local execution.")
    class DummySpaces:
        def GPU(self, *args, **kwargs):
            def decorator(func):
                print(f"Note: Dummy @GPU decorator used for function '{func.__name__}'.")
                return func
            return decorator
    spaces = DummySpaces()

# --- Step 1: Hugging Face Authentication ---
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("FATAL: Hugging Face token not found. Please set the HF_TOKEN environment variable.")
print("--- Logging in to Hugging Face Hub ---")
login(token=HF_TOKEN)

# --- Step 2: Initialize Model and Tokenizer ---
MODEL_NAME = "Gregniuki/micro_doctor"
print(f"--- Loading model from Hugging Face Hub: {MODEL_NAME} ---")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
print(f"--- Using device: {device}, dtype: {dtype} ---")

try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=dtype, trust_remote_code=True).to(device)
    
    # Critical: Set pad token for open-ended generation
    tokenizer.pad_token = tokenizer.eos_token
    model.config.pad_token_id = model.config.eos_token_id

    model.eval()
    print("--- Model and Tokenizer Loaded Successfully ---")
except Exception as e:
    raise RuntimeError(f"FATAL: Could not load components. Error: {e}")

# --- START: System Prompt Definitions ---

# --- Longer, Detailed Prompts ---
MEDICAL_SYSTEM_PROMPT_DETAILED = """You are a helpful AI medical assistant. Your goal is to be a supportive and reassuring resource, breaking down health topics into clear, easy-to-understand language. You always remember that you are an AI, not a doctor.
Say hello give a plan of action, prepare shedule.
Your responses must be thorough, detailed, and follow this exact three-part structure:

**Explain:**
This is the "what" section.
- First, define the health condition or drug in simple terms.
- Then, thoroughly explain the common causes and risk factors.
- Finally, provide a comprehensive list of symptoms, describing not just the symptom itself, but what it might feel like.

**Guide:**
This is the "what you can do" section.
- Provide detailed, actionable steps for at-home and self-care. This should cover supportive measures like rest, hydration, diet, and relevant lifestyle adjustments.
- You can mention general categories of over-the-counter products that help with symptoms (e.g., "pain relievers," "decongestants"), but you must explicitly state that you cannot recommend specific brands or dosages and they should consult a pharmacist or doctor.

**Recommend and prescribe:**
This is the "when to seek help" section. This is the most important part for user safety.
- Clearly state the circumstances under which a person should seek professional medical advice.
- Create a specific list of 'red-flag' symptoms or situations that warrant seeing a doctor promptly or seeking immediate emergency attention. Be clear about the distinction (e.g., "Contact your doctor if..." vs. "Go to the emergency room if...").
"""
MEDICAL_SYSTEM_PROMPT_CONVERSATIONAL = """You are a friendly and conversational AI health buddy. Your tone should be warm, reassuring, and easy to talk to. 
Avoid overly clinical language. Your primary goal is to make health information accessible and less intimidating. 
Start by greeting the user warmly. Always remind them that you're an AI and not a substitute for a real doctor.
When you explain things, use analogies and simple terms. Focus on practical, encouraging advice.
"""
MEDICAL_SYSTEM_PROMPT_CONCISE = """You are a direct and concise AI medical information provider. Your purpose is to deliver information efficiently and clearly.
- Use bullet points.
- Keep sentences short.
- Structure your response into three sections: 1. Overview, 2. Key Symptoms, 3. When to See a Doctor.
- Do not use conversational filler. Get straight to the facts.
- Always include a disclaimer that you are not a medical professional.
"""
MEDICAL_SYSTEM_PROMPT_EMPATHETIC = """You are an empathetic AI health assistant. Your primary role is to listen and provide supportive, understanding, and gentle guidance.
Start by acknowledging the user's feelings (e.g., "It sounds like you're going through a lot," or "I understand why you're concerned.").
Provide information in a soft and caring tone. 
Focus on well-being and self-care. Frame advice around being kind to oneself.
Crucially, gently but clearly guide them towards professional medical help, emphasizing that it's a positive step for their health.
"""
MEDICAL_SYSTEM_PROMPT_TECHNICAL = """You are a technical AI medical assistant. Your audience is assumed to have some medical literacy (e.g., students, health professionals).
Use precise medical terminology.
Provide detailed explanations of pathophysiology, risk factors, and diagnostic criteria where appropriate.
Structure information logically with clear headings.
Reference general classes of medications or treatments, but do not prescribe.
Always conclude with a strong disclaimer about not being a substitute for clinical consultation.
"""

# --- Shorter, Task-Specific Prompts ---
PROMPT_SYMPTOM_CHECKER = "You are a symptom checker AI. Ask clarifying questions to narrow down possibilities. Provide a list of potential conditions, from most to least likely, and strongly advise consulting a doctor for a real diagnosis."
PROMPT_MED_EXPLAINER = "You are an AI medication explainer. The user will provide a drug name. Explain: 1. Use, 2. Common side effects, 3. Important warnings. State this is not a substitute for a pharmacist's advice."
PROMPT_FIRST_AID = "You are a first-aid assistant AI. Provide clear, step-by-step instructions for minor injuries. Start every response by checking for emergency signs (e.g., severe bleeding) and advise calling 911 if present."
PROMPT_DIET_ADVISOR = "You are an AI dietary advisor for specific conditions. Provide general, evidence-based dietary recommendations. Do not create specific meal plans. Emphasize consulting a registered dietitian."
PROMPT_MENTAL_HEALTH = "You are a supportive mental health AI. Offer calming techniques and general wellness advice. If the user expresses thoughts of self-harm, provide a crisis hotline number immediately and advise seeking professional help."
PROMPT_FITNESS_GUIDE = "You are an AI fitness guide. Provide safe, general exercise recommendations. Always include a warning to consult a doctor before starting a new exercise program."
PROMPT_TERM_DEFINER = "You are a medical dictionary AI. The user will provide a medical term. Define it in simple, easy-to-understand language."
PROMPT_LAB_INTERPRETER = "You are an AI lab test explainer. Explain what a specific lab test measures and what results might generally indicate. State clearly that you cannot interpret specific results and they must be discussed with their doctor."
PROMPT_LIFESTYLE_COACH = "You are a healthy lifestyle coach AI. Provide encouraging tips on sleep, stress management, and hydration. Keep your tone positive and motivational."
PROMPT_APPOINTMENT_PREP = "You are an AI pre-appointment assistant. Help the user prepare for a doctor's visit by advising them to write down: 1. Symptoms, 2. Questions for the doctor, and 3. A list of their current medications."


PROMPT_OPTIONS = {
    # Detailed Personas
    "Detailed Medical Assistant": MEDICAL_SYSTEM_PROMPT_DETAILED,
    "Conversational Health Buddy": MEDICAL_SYSTEM_PROMPT_CONVERSATIONAL,
    "Concise & To-the-Point": MEDICAL_SYSTEM_PROMPT_CONCISE,
    "Empathetic & Supportive": MEDICAL_SYSTEM_PROMPT_EMPATHETIC,
    "Technical & In-Depth": MEDICAL_SYSTEM_PROMPT_TECHNICAL,
    # Task-Specific Helpers
    "Symptom Checker": PROMPT_SYMPTOM_CHECKER,
    "Medication Explainer": PROMPT_MED_EXPLAINER,
    "First-Aid Helper": PROMPT_FIRST_AID,
    "Dietary Advisor": PROMPT_DIET_ADVISOR,
    "Mental Health Support": PROMPT_MENTAL_HEALTH,
    "Fitness & Exercise Guide": PROMPT_FITNESS_GUIDE,
    "Medical Terminology Definer": PROMPT_TERM_DEFINER,
    "Lab Test Interpreter": PROMPT_LAB_INTERPRETER,
    "Healthy Lifestyle Coach": PROMPT_LIFESTYLE_COACH,
    "Pre-Appointment Planner": PROMPT_APPOINTMENT_PREP,
    # Custom Option
    "Custom": ""
}

# --- END: System Prompt Definitions ---


# --- Step 3: Core Chat Function with "Show Thoughts" logic ---
@spaces.GPU
@torch.no_grad()
def medical_chat_response(user_message, history, additional_info, show_thoughts, system_prompt):
    """
    Generate a medical assistant response, with an option to show the model's
    internal "think" block.
    """
    print(f"\n--- NEW REQUEST (Show Thoughts: {show_thoughts}) ---")
    
    # 1. Prepare messages for the model, now using the selected system prompt
    full_message = user_message
    if additional_info:
        full_message += f"\n\nAdditional context: {additional_info}"

    messages = [{"role": "system", "content": system_prompt}]
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": full_message})
    
    # 2. Apply chat template and generate response
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([prompt], return_tensors="pt", add_special_tokens=False).to(device)
    
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=8192,  # Increased token limit
        do_sample=True,
        temperature=0.7,
        top_p=0.95,
        top_k=50,
        pad_token_id=tokenizer.eos_token_id
    )
    
    # 3. Decode the full generation
    input_ids_len = model_inputs.input_ids.shape[1]
    new_ids = generated_ids[0][input_ids_len:]
    full_response = tokenizer.decode(new_ids, skip_special_tokens=True).strip()
    print(f"--- FULL MODEL OUTPUT ---\n{full_response}")

    # 4. Conditionally process the response based on the checkbox
    if show_thoughts:
        # If checkbox is ticked, format the <think> block nicely with Markdown
        formatted_response = full_response.replace("<think>", "**Model's Thoughts:**\n```\n").replace("</think>", "\n```\n")
        print(f"--- SENDING FORMATTED RESPONSE TO CHAT ---\n{formatted_response}")
        return formatted_response
    else:
        # If checkbox is not ticked, remove the <think> block entirely
        cleaned_response = full_response
        if "</think>" in cleaned_response:
            cleaned_response = cleaned_response.split("</think>", 1)[-1].strip()
        print(f"--- SENDING CLEANED RESPONSE TO CHAT ---\n{cleaned_response}")
        return cleaned_response

# --- Step 4: Create and Launch the Gradio App ---
print("\n--- Initializing Medical Assistant Chat Interface ---")

# Custom CSS for medical theme
medical_css = """
.main-container {
    background-color: #f0f8ff;
}
.chatbot-container {
    border-radius: 10px;
    box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1);
}
"""

with gr.Blocks(theme="soft", css=medical_css) as demo:
    gr.Markdown("# Medical Assistant Chat")
    gr.Markdown("Ask health-related questions to our AI medical assistant. Remember: This is not a substitute for professional medical advice.")

    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="Chat History", height=500)
            
            with gr.Row():
                msg = gr.Textbox(
                    label="Your Message",
                    placeholder="What are the common symptoms of the flu?",
                    scale=4
                )
            
            with gr.Row():
                additional_info_box = gr.Textbox(
                    label="Additional Information (Optional)",
                    placeholder="Any additional context or specific details...",
                    lines=2
                )
                show_thoughts_checkbox = gr.Checkbox(label="Show model's thought process", value=False)
            
            clear = gr.ClearButton([msg, chatbot])

        with gr.Column(scale=1):
            gr.Markdown("### Assistant Persona")
            prompt_selector = gr.Dropdown(
                label="Choose a Persona or Task",
                choices=list(PROMPT_OPTIONS.keys()),
                value="Detailed Medical Assistant"
            )
            custom_prompt_box = gr.Textbox(
                label="Custom System Prompt",
                placeholder="Enter your own system prompt here...",
                lines=10,
                interactive=False  # Disabled by default
            )

    # Logic to enable/disable custom prompt box
    def update_prompt_box(selection):
        if selection == "Custom":
            return gr.update(interactive=True, value="You are a helpful assistant.")
        else:
            return gr.update(interactive=False, value=PROMPT_OPTIONS.get(selection, ""))
    
    prompt_selector.change(fn=update_prompt_box, inputs=prompt_selector, outputs=custom_prompt_box)
    
    def respond(message, chat_history, additional_info, show_thoughts, prompt_choice, custom_prompt):
        # Determine which prompt to use
        if prompt_choice == "Custom":
            system_prompt = custom_prompt
        else:
            system_prompt = PROMPT_OPTIONS.get(prompt_choice, MEDICAL_SYSTEM_PROMPT_DETAILED)
        
        bot_message = medical_chat_response(message, chat_history, additional_info, show_thoughts, system_prompt)
        chat_history.append((message, bot_message))
        return "", chat_history

    msg.submit(
        respond,
        [msg, chatbot, additional_info_box, show_thoughts_checkbox, prompt_selector, custom_prompt_box],
        [msg, chatbot]
    )
    
    gr.Examples(
        examples=[
            ["What are the common symptoms of the flu?"],
            ["How can I manage mild headaches at home?"],
            ["What is a healthy diet for heart health?"],
            ["When should I see a doctor for a fever?"],
            ["What are some ways to reduce stress?"]
        ],
        inputs=msg
    )
    
    gr.Markdown("""
    <div style="font-size: 0.8em; color: #666; font-style: italic; margin-top: 10px;">
    <strong>Disclaimer:</strong> This AI assistant provides general health information and is not a substitute for professional medical advice, diagnosis, or treatment. 
    Always seek the advice of your physician or other qualified health provider with any questions you may have regarding a medical condition. 
    If you think you may have a medical emergency, call your doctor or emergency services immediately.
    </div>
    """)

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