"""
app.py — Hugging Face Space — Gradio 6 Compatible
Medical AI Diagnostic System
"""
import subprocess, sys

# Install LLaVA without overriding transformers
subprocess.run([
    sys.executable, "-m", "pip", "install", "--no-deps", "--quiet",
    "git+https://github.com/haotian-liu/LLaVA.git"
], check=True)

import io, os, gc, re, uuid, base64
import spaces
import torch
import numpy as np
import cv2
import gradio as gr
from PIL import Image
from peft import PeftModel
from llava.model.builder import load_pretrained_model
from llava.mm_utils import (get_model_name_from_path, process_images,
                             tokenizer_image_token)
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
from llava.conversation import conv_templates

BASE_MODEL_ID = "microsoft/llava-med-v1.5-mistral-7b"
LORA_MODEL_ID = "omaraboelmaaty/llava-med-lora-v2"
LORA_SUBFOLDER = "checkpoints/llava-med-finetuned-v2/final"

model = tokenizer = image_processor = None
_loaded = False

L = {
    "ar": {
        "title":   "نظام التشخيص الطبي بالذكاء الاصطناعي",
        "sub":     "رفع الأشعة  •  التقرير  •  خريطة الخطورة  •  محادثة",
        "upload":  "ارفع الأشعة", "analyze": "تحليل الأشعة",
        "report":  "التقرير الطبي", "overlay": "الأشعة المحللة",
        "heatmap": "خريطة الخطورة", "risk": "نسبة الخطورة",
        "tab1": "تحليل", "tab2": "محادثة",
        "q_title": "أسئلة مقترحة", "chat_ph": "اكتب سؤالك...",
        "send": "إرسال", "sid": "الجلسة",
        "hint": "اكتب شكراً أو thanks لإنهاء الجلسة",
        "no_img": "ارفع أشعة أولاً.", "no_rep": "حلّل الأشعة أولاً.",
        "welcome": "أهلاً دكتور — التحليل جاهز. يمكنك طرح أسئلتك.",
        "closed": "انتهت الجلسة. شكراً.", "rep_ph": "سيظهر التقرير هنا...",
        "lang_lbl": "اللغة",
    },
    "en": {
        "title":   "AI Medical Diagnostic System",
        "sub":     "Upload Scan  •  Report  •  Risk Heatmap  •  Chat",
        "upload":  "Upload Scan", "analyze": "Analyze Scan",
        "report":  "Medical Report", "overlay": "Annotated Scan",
        "heatmap": "Risk Heatmap", "risk": "Risk Score",
        "tab1": "Analysis", "tab2": "Chat",
        "q_title": "Quick Questions", "chat_ph": "Type your question...",
        "send": "Send", "sid": "Session",
        "hint": "Type thanks to end session",
        "no_img": "Upload a scan first.", "no_rep": "Analyze the scan first.",
        "welcome": "Welcome Doctor — Analysis done. Ask your questions.",
        "closed": "Session ended. Thank you.", "rep_ph": "Report will appear here...",
        "lang_lbl": "Language",
    },
}

AR_Q = ["ما أبرز الموجودات؟","هل توجد تغيّرات مرضية؟","ما الانطباع التشخيصي؟",
        "هل يلزم فحص إضافي؟","ما مدى الخطورة؟","ما التوصيات؟"]
EN_Q = ["What are the key findings?","Are there pathological changes?",
        "What is the likely diagnosis?","Is further imaging needed?",
        "How severe is this?","What do you recommend?"]

STOP_WORDS = {"شكرا","شكراً","thanks","thank you","bye","goodbye"}

# ── CSS ───────────────────────────────────────────────────────────────────────
CSS = """
:root {
  --bg:#0b1120; --surface:#151f32; --surface2:#1c2a42; --border:#263450;
  --text:#dce8f5; --text2:#7b9bbf; --text3:#4a6484;
  --accent:#4d8ef0; --accent-h:#3b7de0; --lbl-bg:#1a3a7c; --lbl-text:#c8deff;
  --usr-msg:#162447; --usr-bdr:#2d5499; --bot-msg:#1c2a42; --bot-bdr:#263450;
  --shadow:0 1px 4px rgba(0,0,0,.35); --r:8px;
}
body[data-dark="false"] {
  --bg:#eef2f7; --surface:#ffffff; --surface2:#f4f7fb; --border:#d0d9e8;
  --text:#0f1e35; --text2:#3d5478; --text3:#8fa3be;
  --accent:#2563eb; --accent-h:#1a4fd6; --lbl-bg:#2563eb; --lbl-text:#ffffff;
  --usr-msg:#e8f0fe; --usr-bdr:#aac4f6; --bot-msg:#f4f7fb; --bot-bdr:#d0d9e8;
  --shadow:0 1px 4px rgba(15,30,53,.10);
}
footer { display:none!important }
body, .gradio-container { background:var(--bg)!important; color:var(--text)!important; font-family:Inter,sans-serif!important; font-size:14px!important }
.block,.panel,.form { background:var(--surface)!important; border-color:var(--border)!important; border-radius:var(--r)!important; box-shadow:var(--shadow)!important }
.block .label-wrap { background:var(--lbl-bg)!important; border-radius:5px 5px 0 0!important; padding:3px 8px!important }
.block .label-wrap span { color:var(--lbl-text)!important; font-size:11px!important; font-weight:600!important; text-transform:uppercase!important }
textarea, input[type=text] { background:var(--surface2)!important; border:1px solid var(--border)!important; color:var(--text)!important }
.tab-nav { background:var(--surface)!important; border-bottom:2px solid var(--border)!important }
.tab-nav button { background:transparent!important; border:none!important; border-bottom:3px solid transparent!important; color:var(--text2)!important; box-shadow:none!important }
.tab-nav button.selected { color:var(--accent)!important; border-bottom-color:var(--accent)!important; font-weight:700!important }
button.primary { background:var(--accent)!important; color:#fff!important; border:none!important; border-radius:var(--r)!important; font-weight:600!important; box-shadow:none!important }
button.primary:hover { background:var(--accent-h)!important }
button.secondary { background:var(--surface2)!important; color:var(--text)!important; border:1px solid var(--border)!important; border-radius:var(--r)!important; box-shadow:none!important }
.analyze-btn button { background:var(--accent)!important; color:#fff!important; border:none!important; border-radius:var(--r)!important; font-weight:700!important; width:100%!important; box-shadow:none!important }
.q-btn button { background:var(--surface2)!important; color:var(--text2)!important; border:1px solid var(--border)!important; border-radius:6px!important; font-size:.8rem!important; width:100%!important; margin-bottom:4px!important; text-align:right!important; box-shadow:none!important }
.med-header { background:var(--accent); border-radius:var(--r); padding:14px 20px; margin-bottom:10px; display:flex; align-items:center; gap:12px }
.med-title { font-size:1rem; font-weight:700; color:#fff; margin:0 }
.med-sub { font-size:.75rem; color:rgba(255,255,255,.78); margin:2px 0 0 }
.rtl-box textarea { direction:rtl; text-align:right }
.ltr-box textarea { direction:ltr; text-align:left }
.toggle-row { display:flex; align-items:center; gap:6px; padding:4px 2px }
.tog-icon { font-size:1rem; user-select:none }
.tog-switch { position:relative; width:38px; height:21px; flex-shrink:0 }
.tog-switch input { display:none }
.tog-track { position:absolute; inset:0; background:var(--border); border-radius:21px; transition:background .2s }
.tog-thumb { position:absolute; width:15px; height:15px; background:#fff; border-radius:50%; top:3px; left:3px; transition:left .2s; box-shadow:0 1px 3px rgba(0,0,0,.25) }
.tog-switch input:checked ~ .tog-track { background:var(--accent) }
.tog-switch input:checked ~ .tog-thumb { left:20px }
.toggle-col, .toggle-col > *, .toggle-col .block { background:transparent!important; border:none!important; box-shadow:none!important; padding:0!important }
.svelte-12ioyct, .wrap.svelte-12ioyct { border:none!important; box-shadow:none!important; background:transparent!important }
.wrap.svelte-1kzox3m { background:transparent!important; border:none!important; box-shadow:none!important; padding:0!important; display:flex!important; gap:6px!important }
label.svelte-1mhtq7j { background:var(--surface2)!important; border:1px solid var(--border)!important; border-radius:6px!important; padding:6px 14px!important; box-shadow:none!important }
label.svelte-1mhtq7j.selected { background:var(--accent)!important; border-color:var(--accent)!important }
label.svelte-1mhtq7j.selected span { color:#fff!important }
.hint-txt { color:var(--text3)!important; font-size:.78rem!important; text-align:center }
@media(max-width:768px) { .main-row > div { min-width:100%!important } }
"""

# JS for dark mode default (Gradio 6 uses js parameter in launch())
DARK_JS = """
() => {
    function applyDark() {
        document.body.setAttribute('data-dark', 'true');
    }
    applyDark();
    setTimeout(applyDark, 100);
    setTimeout(applyDark, 500);
    setTimeout(applyDark, 1500);
}
"""

def _toggle_html(checked):
    chk = "checked" if checked else ""
    js = (
        "var d=this.checked;"
        "document.body.setAttribute('data-dark',d?'true':'false');"
        "document.querySelectorAll('label span,.label-wrap span').forEach(function(el){"
        "el.style.color=d?'#c8deff':'#ffffff';});"
        "document.querySelectorAll('textarea,input[type=text]').forEach(function(el){"
        "el.style.background=d?'#1c2a42':'#f4f7fb';"
        "el.style.color=d?'#dce8f5':'#0f1e35';});"
        "document.getElementById('__mode_btn__').click();"
    )
    return (
        f'<div class="toggle-row">'
        f'<span class="tog-icon">☀️</span>'
        f'<label class="tog-switch">'
        f'<input type="checkbox" {chk} onchange="{js}">'
        f'<span class="tog-track"></span>'
        f'<span class="tog-thumb"></span>'
        f'</label>'
        f'<span class="tog-icon">🌙</span>'
        f'</div>'
    )

def _header_html(lang):
    t = L[lang]
    return (
        f'<div class="med-header">'
        f'<span style="font-size:1.5rem">🏥</span>'
        f'<div><p class="med-title">{t["title"]}</p>'
        f'<p class="med-sub">{t["sub"]}</p></div></div>'
    )

def init_state():
    return {"session_id": str(uuid.uuid4())[:8],
            "report": "", "ended": False,
            "lang": "ar", "dark": True, "image": None}

# ── Model loading ─────────────────────────────────────────────────────────────
def load_model():
    global model, tokenizer, image_processor, _loaded
    if _loaded: return
    print("🔄 Loading LLaVA-Med + LoRA...")
    from transformers import BitsAndBytesConfig
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=False,
        bnb_4bit_quant_type="nf4",
    )
    name = get_model_name_from_path(BASE_MODEL_ID)
    tokenizer, base, image_processor, _ = load_pretrained_model(
        model_path=BASE_MODEL_ID, model_base=None,
        model_name=name, device_map="cuda",
        quantization_config=bnb_config,
    )
    model = PeftModel.from_pretrained(base, LORA_MODEL_ID, subfolder=LORA_SUBFOLDER, autocast_adapter_dtype=False)
    model.eval()
    _loaded = True
    print("✅ Model ready!")

# ── Generation ────────────────────────────────────────────────────────────────
@spaces.GPU(duration=120)
def _generate(image, prompt, max_tokens=512, greedy=False):
    load_model()
    img_t = process_images([image], image_processor, model.config).to(model.device, dtype=torch.float16)
    conv = conv_templates["mistral_instruct"].copy()
    conv.append_message(conv.roles[0], DEFAULT_IMAGE_TOKEN + "\n" + prompt)
    conv.append_message(conv.roles[1], None)
    ids = tokenizer_image_token(conv.get_prompt(), tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(model.device)
    gen_kwargs = dict(images=img_t, image_sizes=[image.size], max_new_tokens=max_tokens, pad_token_id=tokenizer.eos_token_id)
    if greedy:
        gen_kwargs["do_sample"] = False
    else:
        gen_kwargs.update({"do_sample": True, "temperature": 0.7, "top_p": 0.9, "repetition_penalty": 1.1})
    with torch.inference_mode():
        out = model.generate(ids, **gen_kwargs)
    full = tokenizer.decode(out[0], skip_special_tokens=True)
    result = full.split("[/INST]")[-1].strip() if "[/INST]" in full else full.strip()
    for pat in [r'\[/?INST\]', r'\[/?s\]', r'</?s>', r'End now with[^\n]*\n?',
                r'Radiologist:[^\n]*\n?', r'You MUST[^\n]*\n?', r'Do NOT[^\n]*\n?']:
        result = re.sub(pat, '', result, flags=re.IGNORECASE)
    return result.strip()

REPORT_PROMPT = (
    "You are a radiologist. Analyze this medical image.\n"
    "Respond in ENGLISH ONLY.\n\n"
    "FINDINGS: [describe what you see]\n"
    "IMPRESSION: [your diagnosis]\n"
    "RECOMMENDATIONS: [next steps]\n\n"
    "FINDINGS:"
)

def _format_report(text, lang):
    text = re.sub(r'^\s*Radiology\s*', '', text, flags=re.IGNORECASE).strip()
    text = re.sub(r'(?im)^(english|arabic|french)\s*:\s*', '', text)
    text = re.sub(r'(?m)^\s*None\.?\s*$', '', text)
    text = re.sub(r'None provided[^\n]*', '', text, flags=re.IGNORECASE)
    text = re.sub(r'\n{3,}', '\n\n', text).strip()
    if not re.search(r'(?i)FINDINGS\s*:', text):
        text = "FINDINGS:\n" + text
    if not re.search(r'(?i)IMPRESSION\s*:', text):
        text += "\n\nIMPRESSION: Normal study unless otherwise noted."
    if not re.search(r'(?i)RECOMMENDATIONS?\s*:', text):
        text += "\n\nRECOMMENDATIONS: Clinical correlation recommended."
    if lang == "ar":
        for pat, rep in [
            (r'(?i)FINDINGS\s*:',        '\n📋 الموجودات (Findings):'),
            (r'(?i)IMPRESSION\s*:',      '\n\n🔍 الانطباع (Impression):'),
            (r'(?i)RECOMMENDATIONS?\s*:', '\n\n💡 التوصيات (Recommendations):'),
        ]:
            text = re.sub(pat, rep, text)
    else:
        for pat, rep in [
            (r'(?i)IMPRESSION\s*:',      '\n\nIMPRESSION:'),
            (r'(?i)RECOMMENDATIONS?\s*:', '\n\nRECOMMENDATIONS:'),
        ]:
            text = re.sub(pat, rep, text)
    return text.strip()

@spaces.GPU(duration=120)
def generate_report(image, lang="ar"):
    load_model()
    result = _generate(image, REPORT_PROMPT, max_tokens=600, greedy=False)
    see_above = re.compile(r'(IMPRESSION|RECOMMENDATIONS?)\s*:\s*(see findings[^\n]*|see above[^\n]*)', re.IGNORECASE)
    if see_above.search(result):
        fm = re.search(r'FINDINGS\s*:(.*?)(?=\nIMPRESSION|\nRECOMMENDATION|$)', result, re.IGNORECASE | re.DOTALL)
        ft = fm.group(1).strip()[:200] if fm else ""
        new_imp = _generate(image, f"Based on: {ft} State the diagnosis in one sentence. English only.", max_tokens=80, greedy=True)
        new_imp = new_imp.split('\n')[0].strip()
        result = see_above.sub(lambda m: m.group(1).upper() + ": " + new_imp, result)
    return _format_report(result, lang)

@spaces.GPU(duration=60)
def answer_question(image, question, lang="ar"):
    load_model()
    prompt = (
        "You are a radiologist. Look at this medical image carefully. "
        "Answer the following question briefly and accurately. English only.\n"
        f"Question: {question}\nAnswer:"
    )
    return _generate(image, prompt, max_tokens=300, greedy=True)

# ── Heatmap ───────────────────────────────────────────────────────────────────
def _add_legend(img_np):
    h, w = img_np.shape[:2]
    bar = np.full((30, w, 3), (20, 20, 20), dtype=np.uint8)
    step = w // 3
    for i, (color, lbl) in enumerate([((60,50,220),"High Risk"),((50,140,220),"Medium"),((50,220,220),"Low")]):
        x = i*step+6
        cv2.rectangle(bar,(x,8),(x+12,22),color[::-1],-1)
        cv2.putText(bar,lbl,(x+16,20),cv2.FONT_HERSHEY_SIMPLEX,0.45,(200,200,200),1)
    return np.vstack([img_np, bar])

def _find_clip_vision_model():
    """Find the actual transformers CLIPVisionModel instance by CLASS NAME,
    walking the real module graph via named_modules(). This is robust to
    transformers version changes that alter attribute nesting/names (e.g.
    `.vision_tower.vision_model` no longer resolving the same way) — we let
    the module's own forward() do its own internal attribute access instead
    of us guessing the attribute chain from outside."""
    for name, module in model.named_modules():
        if module.__class__.__name__ == "CLIPVisionModel":
            return name, module
    return None, None

def _find_last_vision_encoder_layer(clip_vision_model):
    """Find the LAST CLIP transformer layer by CLASS NAME (CLIPEncoderLayer),
    searched only within the already-resolved CLIPVisionModel's own subtree.
    Deliberately avoids any attribute-name check (self_attn, vision_model,
    encoder.layers substring, numeric suffix parsing) — those attribute names
    have already changed between transformers versions in this environment
    and broken previous fixes. Class names for CLIP have been stable across
    versions. nn.ModuleList preserves registration order, so the last match
    encountered during named_modules() traversal IS the last layer."""
    last_name, last_module = None, None
    for name, module in clip_vision_model.named_modules():
        if module.__class__.__name__ == "CLIPEncoderLayer":
            last_name, last_module = name, module
    return last_name, last_module

@spaces.GPU(duration=60)
def generate_heatmap(image):
    load_model()
    feat = [None]

    def _hook(module, inp, out):
        o = out[0] if isinstance(out, tuple) else out
        if hasattr(o, 'dim') and o.dim() == 3 and o.shape[1] > 100:
            feat[0] = o[0, 1:].norm(dim=-1).detach().cpu().float().numpy()

    handles = []
    try:
        cv_name, clip_vision_model = _find_clip_vision_model()
        if clip_vision_model is None:
            raise RuntimeError("CLIPVisionModel instance not found")
        print(f"Found CLIPVisionModel: {cv_name}")

        layer_name, layer_module = _find_last_vision_encoder_layer(clip_vision_model)
        if layer_module is None:
            raise RuntimeError("vision encoder layer not found")
        handles.append(layer_module.register_forward_hook(_hook))
        print(f"Hooked last CLIP encoder layer: {cv_name}.{layer_name}")

        img_t = process_images([image], image_processor, model.config).to(model.device, dtype=torch.float16)

        # Call the CLIPVisionModel DIRECTLY (standard transformers API takes
        # pixel_values=) instead of the full LLaVA causal LM. The full model
        # only routes pixels through the vision tower when input_ids contains
        # the special IMAGE_TOKEN_INDEX (-200) placeholder; a dummy zero
        # input_ids never triggers that path, so the hook never fires.
        with torch.no_grad():
            clip_vision_model(pixel_values=img_t)
    except Exception as e:
        print(f"Heatmap forward error: {e}")
    finally:
        for h in handles: h.remove()

    if feat[0] is None:
        print("Heatmap: no feature captured, returning original")
        return image, image, 0.0

    cam = feat[0]
    n = int(round(len(cam)**0.5))
    if n*n != len(cam):
        print(f"Heatmap: non-square feature {len(cam)}")
        return image, image, 0.0

    cam = cam.reshape(n,n)
    cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)
    W,H = image.size
    cam_r = cv2.GaussianBlur(cv2.resize(cam,(W,H),interpolation=cv2.INTER_CUBIC),(15,15),0)
    cam_r = (cam_r - cam_r.min()) / (cam_r.max() - cam_r.min() + 1e-8)
    top_px = cam_r[cam_r >= float(np.percentile(cam_r,80))]
    risk = round(min(max(float(np.mean(top_px))*80 + float(np.std(cam_r))*20, 0), 100), 1)
    hm_c = cv2.applyColorMap(np.uint8(255*cam_r), cv2.COLORMAP_JET)
    hm_rgb = cv2.cvtColor(hm_c, cv2.COLOR_BGR2RGB)
    orig = np.array(image.convert("RGB"))
    overlay = cv2.addWeighted(orig, 0.55, hm_rgb, 0.45, 0)
    print(f"✅ Heatmap generated, risk={risk}%")
    return Image.fromarray(_add_legend(overlay)), Image.fromarray(_add_legend(hm_rgb)), risk


# ── Handlers ──────────────────────────────────────────────────────────────────
def handle_analyze(image, lang, state):
    t = L[lang]
    if image is None:
        return ("", image, image, "—", state, [], state["session_id"])
    state["lang"] = lang; state["image"] = image
    try:
        report = generate_report(image, lang)
        state["report"] = report
        ov, hm, risk = generate_heatmap(image)
        return (report, ov, hm, f"{risk}%", state, [{"role":"assistant","content":t["welcome"]}], state["session_id"])
    except Exception as e:
        return (f"خطأ: {e}", image, image, "—", state, [], state["session_id"])

def handle_chat(question, history, state):
    if not isinstance(question, str) or not question.strip():
        return history, state, ""
    lang = state.get("lang","ar"); t = L[lang]
    if state.get("ended"):
        return history+[{"role":"user","content":question},{"role":"assistant","content":t["closed"]}], state, ""
    if not state.get("report"):
        return history+[{"role":"user","content":question},{"role":"assistant","content":t["no_rep"]}], state, ""
    if question.strip().lower() in STOP_WORDS:
        state["ended"] = True
        return history+[{"role":"user","content":question},{"role":"assistant","content":t["closed"]}], state, ""
    try:
        image = state.get("image")
        answer = answer_question(image, question, lang) if image else t["no_img"]
        return history+[{"role":"user","content":question},{"role":"assistant","content":answer}], state, ""
    except Exception as e:
        return history+[{"role":"user","content":question},{"role":"assistant","content":f"خطأ: {e}"}], state, ""

def quick_q(q_text, history, state):
    return handle_chat(q_text, history, state)



def toggle_mode(state):
    state["dark"] = not state.get("dark", True)
    return state, gr.update(value=_toggle_html(state["dark"]))

def switch_lang(lang, state):
    state["lang"] = lang; t = L[lang]; rtl = lang == "ar"; dark = state.get("dark", True)
    return (state, gr.update(value=_header_html(lang)),
            gr.update(label=t["upload"]), gr.update(value=t["analyze"]),
            gr.update(label=t["report"], placeholder=t["rep_ph"], elem_classes="rtl-box" if rtl else "ltr-box"),
            gr.update(label=t["overlay"]), gr.update(label=t["heatmap"]),
            gr.update(label=t["risk"]), gr.update(label=t["sid"]),
            gr.update(label=t["tab1"]), gr.update(label=t["tab2"]),
            gr.update(label=t["chat_ph"]), gr.update(value=t["send"]),
            gr.update(value=f"**{t['q_title']}**"),
            gr.update(visible=rtl), gr.update(visible=not rtl),
            gr.update(value=f'<p class="hint-txt">{t["hint"]}</p>'),
            gr.update(label=t["lang_lbl"]), gr.update(value=_toggle_html(dark)))

# ── Build UI ──────────────────────────────────────────────────────────────────
def build_ui():
    with gr.Blocks(title="Medical AI") as demo:

        gr.HTML(f"<style>{CSS}</style>")
        gr.HTML("""<style>
.svelte-12ioyct,.wrap.svelte-12ioyct{border:none!important;box-shadow:none!important;background:transparent!important}
.wrap.svelte-1kzox3m{background:transparent!important;border:none!important;box-shadow:none!important;padding:0!important;display:flex!important;gap:6px!important}
label.svelte-1mhtq7j{background:var(--surface2)!important;border:1px solid var(--border)!important;border-radius:6px!important;padding:6px 14px!important;box-shadow:none!important}
label.svelte-1mhtq7j.selected{background:var(--accent)!important;border-color:var(--accent)!important}
label.svelte-1mhtq7j.selected span{color:#fff!important}
</style>""")

        state = gr.State(init_state())
        header = gr.HTML(_header_html("ar"))

        with gr.Row(equal_height=True):
            lang_r  = gr.Radio(choices=[("العربية","ar"),("English","en")], value="ar", label=L["ar"]["lang_lbl"], scale=2)
            sid_box = gr.Textbox(label=L["ar"]["sid"], value="—", interactive=False, scale=3)
            with gr.Column(scale=1, min_width=100, elem_classes="toggle-col"):
                mode_html = gr.HTML(_toggle_html(True))
                mode_btn  = gr.Button("m", visible=False, elem_id="__mode_btn__")

        with gr.Tabs():
            with gr.Tab(L["ar"]["tab1"]) as tab1:
                with gr.Row(elem_classes="main-row"):
                    with gr.Column(scale=1, min_width=200):
                        img_in = gr.Image(label=L["ar"]["upload"], type="pil", height=240)
                        an_btn = gr.Button(L["ar"]["analyze"], variant="primary", elem_classes="analyze-btn")
                        risk   = gr.Textbox(label=L["ar"]["risk"], value="—", interactive=False)
                    with gr.Column(scale=1, min_width=200):
                        ov_out = gr.Image(label=L["ar"]["overlay"], height=185, interactive=False)
                        hm_out = gr.Image(label=L["ar"]["heatmap"], height=185, interactive=False)
                    with gr.Column(scale=1, min_width=240):
                        rep_out = gr.Textbox(label=L["ar"]["report"], lines=17, interactive=False, elem_classes="rtl-box", placeholder=L["ar"]["rep_ph"])

            with gr.Tab(L["ar"]["tab2"]) as tab2:
                with gr.Row():
                    with gr.Column(scale=1, min_width=175):
                        q_title = gr.Markdown(f"**{L['ar']['q_title']}**")
                        with gr.Group(visible=True) as ar_grp:
                            ar_btns = [gr.Button(q, elem_classes="q-btn", size="sm") for q in AR_Q]
                        with gr.Group(visible=False) as en_grp:
                            en_btns = [gr.Button(q, elem_classes="q-btn", size="sm") for q in EN_Q]
                    with gr.Column(scale=3, min_width=280):
                        chatbot = gr.Chatbot(height=380,
                                             avatar_images=("https://cdn-icons-png.flaticon.com/512/3774/3774299.png",
                                                            "https://cdn-icons-png.flaticon.com/512/4712/4712035.png"))
                        with gr.Row():
                            chat_in = gr.Textbox(label=L["ar"]["chat_ph"], lines=2, scale=5, placeholder="اكتب سؤالك...")
                            with gr.Column(scale=1, min_width=100):
                                send_btn = gr.Button(L["ar"]["send"], variant="primary", size="lg")
                        hint_html = gr.HTML(f'<p class="hint-txt">{L["ar"]["hint"]}</p>')

        an_btn.click(fn=handle_analyze, inputs=[img_in, lang_r, state],
                     outputs=[rep_out, ov_out, hm_out, risk, state, chatbot, sid_box])
        def _cw(q,h,s): return handle_chat(q,h,s)
        send_btn.click(_cw, [chat_in,chatbot,state], [chatbot,state,chat_in])
        chat_in.submit(_cw, [chat_in,chatbot,state], [chatbot,state,chat_in])
        for btn in ar_btns+en_btns:
            btn.click(fn=quick_q, inputs=[gr.State(btn.value),chatbot,state], outputs=[chatbot,state,chat_in])
        mode_btn.click(fn=toggle_mode, inputs=[state], outputs=[state,mode_html])
        lang_r.change(fn=switch_lang, inputs=[lang_r,state],
                      outputs=[state,header,img_in,an_btn,rep_out,ov_out,hm_out,risk,sid_box,
                               tab1,tab2,chat_in,send_btn,q_title,ar_grp,en_grp,hint_html,lang_r,mode_html])
    return demo

demo = build_ui()

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