import gradio as gr
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
import torch
import re
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)

# Configuration
BASE_MODEL_ID = "Qwen/Qwen3-0.6B"
LORA_ADAPTER_ID = "towardsinnovationlab/qwen3-medical"

model = None
tokenizer = None

# --- Guardrail keywords ---
MEDICAL_KEYWORDS = [
    "symptom", "disease", "diagnosis", "treatment", "medication", "drug", "dose",
    "pain", "fever", "headache", "diabetes", "blood pressure", "heart", "cancer",
    "infection", "allergy", "vaccine", "surgery", "mental health", "anxiety",
    "depression", "sleep", "diet", "nutrition", "exercise", "cholesterol",
    "pregnancy", "child", "doctor", "hospital", "prescription", "side effect",
    "vitamin", "supplement", "injury", "wound", "bone", "muscle", "skin",
    "digestive", "stomach", "kidney", "liver", "lung", "breathing", "cough",
    "cold", "flu", "covid", "virus", "bacteria", "antibiotic", "immune",
    "thyroid", "hormone", "stroke", "migraine", "dehydration", "screening",
    "health", "medical", "clinical", "patient", "nurse", "physician", "therapy",
    "chronic", "acute", "preventive", "care", "hygiene", "weight", "obesity",
    "asthma", "arthritis", "spine", "nerve", "eye", "ear", "dental", "teeth"
]

OFF_TOPIC_RESPONSE = (
    "⚠️ I'm a **medical assistant** and can only answer health and medical questions.\n\n"
    "Please ask me about symptoms, diseases, treatments, medications, nutrition, "
    "mental health, or any other medical topic. I'm here to help! 🏥"
)


def is_medical_question(message: str) -> bool:
    """Return True if the message appears to be health/medical related."""
    lowered = message.lower()
    return any(kw in lowered for kw in MEDICAL_KEYWORDS)


def clean_response(text: str) -> str:
    """Remove <think> tags and leftover special tokens from model output."""
    text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
    text = re.sub(r'<\|[^>]+\|>', '', text)
    return text.strip()


@spaces.GPU(duration=120)
def respond(message: str, history: list) -> str:
    """Generate a medical response, rejecting off-topic queries."""
    global model, tokenizer

    # --- Guardrail: block non-medical questions ---
    if not is_medical_question(message):
        return OFF_TOPIC_RESPONSE

    # --- Lazy model loading ---
    if model is None:
        print("Loading tokenizer...")
        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)

        print("Loading base model with 4-bit quantization...")
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
        )
        base_model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL_ID,
            quantization_config=bnb_config,
            device_map="auto",
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )

        print("Loading LoRA adapter...")
        model = PeftModel.from_pretrained(base_model, LORA_ADAPTER_ID, device_map="auto")
        model.eval()
        print("Model loaded!")

    # --- Build conversation ---
    messages = [
        {
            "role": "system",
            "content": (
                "You are a helpful medical assistant. Answer only health and medical questions. "
                "If a question is not related to medicine, health, symptoms, treatments, or wellness, "
                "politely decline and remind the user to ask medical questions only. "
                "Provide accurate, detailed medical information using bullet points where appropriate."
            )
        }
    ]
    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})
    messages.append({"role": "user", "content": message})

    # --- Tokenize ---
    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    if hasattr(inputs, 'input_ids'):
        inputs = inputs['input_ids']

    # --- Generate ---
    with torch.inference_mode():
        outputs = model.generate(
            inputs,
            max_new_tokens=1024,
            min_new_tokens=50,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            top_k=50,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            use_cache=True,
        )

    new_tokens = outputs[0][inputs.shape[1]:]
    clean = tokenizer.decode(new_tokens, skip_special_tokens=True)
    final_response = clean_response(clean)

    # Fallback if response is suspiciously short
    if len(final_response) < 20:
        raw = tokenizer.decode(new_tokens, skip_special_tokens=False)
        return clean or raw

    return final_response


custom_css = """
.examples button {
    background: linear-gradient(135deg, #0891b2 0%, #0e7490 100%) !important;
    border: none !important;
    color: white !important;
    border-radius: 8px !important;
    padding: 12px 16px !important;
    font-weight: 500 !important;
    transition: all 0.2s ease !important;
}
.examples button:hover {
    background: linear-gradient(135deg, #0e7490 0%, #155e75 100%) !important;
    transform: translateY(-2px);
    box-shadow: 0 4px 12px rgba(8, 145, 178, 0.4) !important;
}
.chat-banner {
    background: linear-gradient(135deg, #3b82f6 0%, #2563eb 100%);
    color: white;
    padding: 12px 24px;
    border-radius: 10px;
    text-align: center;
    font-weight: 600;
    font-size: 1.1em;
    margin: 15px 0;
    box-shadow: 0 4px 12px rgba(59, 130, 246, 0.3);
}
"""

with gr.Blocks(css=custom_css) as demo:
    gr.ChatInterface(
        fn=respond,
        title="🏥 Medical Chatbot",
        description=(
            "Powered by Qwen3-Medical (4-bit quantized). "
            "⚠️ For informational purposes only — not a substitute for professional medical advice. "
            "Responses may take up to 60 seconds. Best results in English."
        ),
        examples=[
            "What are the main symptoms of diabetes?",
            "How can I lower my blood pressure naturally?",
            "What are the side effects of ibuprofen?",
            "What causes migraines and how can I prevent them?",
            "How much sleep do adults need each night?",
            "What are the early warning signs of a heart attack?",
            "When should I see a doctor for a persistent cough?",
            "What foods should I avoid if I have high cholesterol?",
            "What are the symptoms of dehydration and how can I prevent it?",
            "How is anxiety different from normal stress, and what are effective coping strategies?",
            "What is high blood sugar (hyperglycemia) and what should I do if I suspect it?",
            "How often should I get routine health screenings?",
        ],
        cache_examples=False,
    )
    gr.HTML('<div class="chat-banner">💬 Ask me any medical or health question!</div>')

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