import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model = None
tokenizer = None

def load_model():
    global model, tokenizer
    if model is None:
        print("[MediAssist] Loading...")
        model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForCausalLM.from_pretrained(
            model_id, torch_dtype=torch.float32, low_cpu_mem_usage=True
        )
        model.eval()
        print("[MediAssist] Ready!")

# For Urdu/Hindi we translate the question to English, get answer, keep it simple
# TinyLlama only works well in English — this is honest and works
SYSTEM_EN = """You are MediAssist, a medical assistant for rural Pakistan communities.
Answer clearly in English with this exact format:
🔍 What this might be:
- cause 1
- cause 2
🏠 Home care steps:
- step 1
- step 2
- step 3
🚨 Go to doctor immediately if:
- warning 1
- warning 2
Keep under 120 words. Never diagnose. Always suggest seeing a doctor for serious issues."""

SYSTEM_UR = """You are a medical assistant. The user is asking in Urdu.
First translate their question to English, answer in English with this format:
🔍 What this might be:
- cause
🏠 Home care:
- step 1
- step 2
🚨 See doctor if:
- warning
Then write: "اردو خلاصہ:" and give a 2-line Urdu summary of your answer.
Keep total response under 150 words."""

SYSTEM_HI = """You are a medical assistant. Answer in simple Hindi with this format:
🔍 यह क्या हो सकता है:
- कारण
🏠 घर पर करें:
- कदम 1
- कदम 2
🚨 डॉक्टर के पास जाएं अगर:
- चेतावनी
100 शब्दों से कम रखें।"""

SYSTEMS = {
    "English": SYSTEM_EN,
    "اردو": SYSTEM_UR,
    "हिन्दी": SYSTEM_HI,
}

DISCLAIMERS = {
    "English": "\n\n⚠️ *For informational purposes only. Always consult a real doctor.*",
    "اردو":    "\n\n⚠️ *صرف معلوماتی مقاصد کے لیے۔ ہمیشہ ڈاکٹر سے مشورہ کریں۔*",
    "हिन्दी":  "\n\n⚠️ *केवल जानकारी के लिए। हमेशा डॉक्टर से सलाह लें।*",
}

def respond(message, history, language):
    load_model()
    system = SYSTEMS.get(language, SYSTEM_EN)
    chat = [{"role": "system", "content": system}]
    for user_msg, bot_msg in history:
        if user_msg: chat.append({"role": "user",      "content": user_msg})
        if bot_msg:  chat.append({"role": "assistant", "content": bot_msg})
    chat.append({"role": "user", "content": message})

    prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt")
    input_len = inputs["input_ids"].shape[1]

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=300,
            do_sample=False,
            repetition_penalty=1.3,
            pad_token_id=tokenizer.eos_token_id,
        )

    reply = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()
    return reply + DISCLAIMERS.get(language, DISCLAIMERS["English"])

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
* { font-family: 'Inter', sans-serif !important; }
body, .gradio-container { background: #0a0f0d !important; }
footer { display: none !important; }
.gradio-container { max-width: 860px !important; margin: 0 auto !important; padding: 24px 16px !important; }
.mh {
    background: linear-gradient(135deg, #064e3b, #047857);
    border: 1px solid #10b981; border-radius: 20px; padding: 28px 32px; margin-bottom: 12px;
}
.hb { display: flex; align-items: center; gap: 14px; }
.hi {
    width: 52px; height: 52px; background: rgba(16,185,129,0.2);
    border: 1.5px solid rgba(16,185,129,0.5); border-radius: 14px;
    display: flex; align-items: center; justify-content: center; font-size: 24px; flex-shrink:0;
}
.ht { color: #ecfdf5; font-size: 26px; font-weight: 700; margin: 0; }
.hs { color: #6ee7b7; font-size: 13px; margin: 3px 0 0; }
.hbadges { display: flex; gap: 8px; flex-wrap: wrap; margin-top: 16px; }
.badge {
    background: rgba(16,185,129,0.15); border: 1px solid rgba(16,185,129,0.35);
    color: #6ee7b7; font-size: 12px; font-weight: 500; padding: 5px 12px; border-radius: 20px;
}
.lr {
    background: #111816; border: 1px solid #1f2e28; border-radius: 14px;
    padding: 14px 20px; margin-bottom: 8px; display: flex; align-items: center; gap: 12px;
}
.ll { color: #6ee7b7; font-size: 13px; font-weight: 500; white-space: nowrap; }
.db {
    background: #1a1500; border: 1px solid #78350f; border-left: 3px solid #f59e0b;
    border-radius: 10px; padding: 12px 16px; margin-top: 8px;
    font-size: 12px; color: #fbbf24; line-height: 1.6;
}
"""

with gr.Blocks(title="MediAssist") as demo:

    gr.HTML("""
    <div class="mh">
      <div class="hb">
        <div class="hi">🏥</div>
        <div>
          <p class="ht">MediAssist</p>
          <p class="hs">AI-powered health guidance for underserved communities</p>
        </div>
      </div>
      <div class="hbadges">
        <span class="badge">🧠 TinyLlama-1.1B</span>
        <span class="badge">🆓 100% Free</span>
        <span class="badge">🌍 3 Languages</span>
        <span class="badge">🇵🇰 Built for Rural Punjab</span>
        <span class="badge">✅ &lt;32B params</span>
      </div>
    </div>
    """)

    with gr.Row(elem_classes=["lr"]):
        gr.HTML('<span class="ll">🌐 Language / زبان / भाषा</span>')
        language = gr.Radio(
            choices=["English", "اردو", "हिन्दी"],
            value="English", show_label=False,
        )

    gr.ChatInterface(
        fn=respond,
        additional_inputs=[language],
        chatbot=gr.Chatbot(height=440, show_label=False),
        textbox=gr.Textbox(
            placeholder="Describe your symptoms or ask a health question…",
            show_label=False, lines=2,
        ),
        examples=[
            ["I have fever 38.5°C and headache for 2 days",   "English"],
            ["Paracetamol dose for a 5 year old child?",      "English"],
            ["Signs of dehydration in children?",             "English"],
            ["How to treat a minor burn at home?",            "English"],
            ["مجھے بخار اور سر درد ہے",                      "اردو"],
            ["بچے کو پیراسیٹامول کتنی دیں؟",                 "اردو"],
        ],
        title="", description="",
    )

    gr.HTML("""
    <div class="db">
      ⚠️ <strong>Medical Disclaimer:</strong> MediAssist provides general health information only.
      Always consult a qualified healthcare provider. In emergencies, call your local emergency number.
    </div>
    """)

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