import os, json, sys, html, numpy as np, gradio as gr, torch
from huggingface_hub import snapshot_download

PRIVATE = os.environ.get("MODEL_REPO", "")
d = snapshot_download(PRIVATE, token=os.environ.get("HF_TOKEN")) if PRIVATE else os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, d)
import keep_audio_core as K
cfg = json.load(open(os.path.join(d, "config.json"))) if os.path.exists(os.path.join(d, "config.json")) else {}
R = K.Recognizer(os.path.join(d, "audio2b_KEEP_graine0.pt"), os.path.join(d, "audio2b_normalisation.pt"), gain=cfg.get("gain", "peak"))

# Voice-first shell, same language as KEEP Listen: pure black, one narrow column,
# numbered steps instead of cards, huge type, pill actions.
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@400;500;700&display=swap');
html, body { background: #000000 !important; }
.gradio-container {
  --body-background-fill: #000000;
  --background-fill-primary: #000000;
  --background-fill-secondary: #000000;
  --block-background-fill: transparent;
  --block-border-color: transparent;
  --block-label-background-fill: transparent;
  --block-label-text-color: #8A8A93;
  --block-title-text-color: #F5F5F7;
  --body-text-color: #F5F5F7;
  --body-text-color-subdued: #8A8A93;
  --button-primary-background-fill: #F5F5F7;
  --button-primary-background-fill-hover: #FFFFFF;
  --button-primary-text-color: #000000;
  --button-secondary-background-fill: transparent;
  --button-secondary-background-fill-hover: #1A1A1E;
  --button-secondary-text-color: #F5F5F7;
  --button-secondary-border-color: rgba(255,255,255,.22);
  --input-background-fill: transparent;
  --input-border-color: rgba(255,255,255,.22);
  --input-text-color: #F5F5F7;
  --input-placeholder-color: #6E6E73;
  max-width: 560px !important; margin: 0 auto !important; padding: 26px 22px 44px !important;
  background: #000000 !important; color: #F5F5F7 !important;
  font-family: 'Space Grotesk', system-ui, sans-serif !important;
}
.gradio-container .block, .gradio-container .form, .gradio-container .wrap,
.gradio-container .panel, .gradio-container .component, .gradio-container .audio,
.gradio-container .plot, .gradio-container .input, .gradio-container .output {
  background: transparent !important; border: none !important;
  box-shadow: none !important; border-radius: 0 !important;
}
.gradio-container label span {
  color: #8A8A93 !important; font-size: 11px !important; font-weight: 500 !important;
  letter-spacing: 2.5px !important; text-transform: uppercase !important;
}
.gradio-container .prose { color: #C7C7CC !important; }
.gradio-container .prose strong { color: #F5F5F7 !important; }
.topbar { display: flex; justify-content: space-between; align-items: center; margin-bottom: 8px; }
.brand { font-weight: 700; font-size: 19px; letter-spacing: -.5px; }
.brand i { color: #E5C15D; font-style: normal; }
.ondevice { font-size: 12px; color: #8A8A93; border: 1px solid rgba(255,255,255,.18); border-radius: 999px; padding: 5px 12px; }
.intro h1 { font-size: 46px; font-weight: 700; letter-spacing: -2px; line-height: 1.04; margin: 30px 0 14px; }
.intro .sub { color: #8A8A93; font-size: 16px; line-height: 1.6; margin: 0; }
.intro .sub b { color: #C7C7CC; }
.step { display: flex; align-items: baseline; gap: 14px; margin: 36px 0 4px; }
.step .num { color: #E5C15D; font-weight: 700; font-size: 14px; letter-spacing: 1px; }
.step h2 { font-size: 27px; font-weight: 700; letter-spacing: -.5px; margin: 0; }
.hairline { height: 1px; background: rgba(255,255,255,.10); margin: 32px 0; }
.chips { display: flex; flex-wrap: wrap; gap: 8px; }
.chip { font-size: 14px; color: #F5F5F7; background: #141416; border-radius: 999px; padding: 9px 15px; }
.chip span { color: #8A8A93; }
.dots { display: flex; gap: 9px; align-items: center; margin: 14px 0 4px; }
.dots i { width: 11px; height: 11px; border-radius: 50%; background: #26262A; }
.dots i.on { background: #E5C15D; }
.dots span { color: #8A8A93; font-size: 13px; margin-left: 8px; }
.qstage { text-align: center; padding: 10px 0 4px; }
.qstage .eyebrow { font-size: 11px; letter-spacing: 3px; text-transform: uppercase; color: #8A8A93; margin-bottom: 12px; }
.qstage .word { font-size: 64px; font-weight: 700; letter-spacing: -2.5px; line-height: 1; margin: 0; }
.qstage .word.dim { color: #333338; }
.qstage .bars { margin-top: 20px; display: flex; flex-direction: column; gap: 8px; }
.qstage .bars .row { display: flex; align-items: center; gap: 12px; font-size: 14px; color: #C7C7CC; }
.qstage .bars .row b { width: 110px; text-align: right; font-weight: 500; color: #F5F5F7; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
.qstage .bars .bar { flex: 1; height: 5px; background: #1C1C1F; border-radius: 99px; overflow: hidden; }
.qstage .bars .bar i { display: block; height: 100%; background: linear-gradient(90deg, #C9A227, #F5D98B); border-radius: 99px; }
.qstage .bars .pct { width: 46px; text-align: left; color: #8A8A93; }
.quiet { color: #8A8A93 !important; font-size: 13.5px !important; line-height: 1.65 !important; }
.quiet strong { color: #C7C7CC !important; }
.gradio-container button.primary { border-radius: 999px !important; height: 56px !important; font-size: 17px !important; font-weight: 700 !important; }
.gradio-container button:not(.primary) { border-radius: 999px !important; }
.gradio-container .audio button { border-radius: 999px !important; background: #141416 !important; border: 1px solid rgba(255,255,255,.14) !important; }
.gradio-container input[type="text"] {
  background: transparent !important; border: none !important;
  border-bottom: 2px solid rgba(255,255,255,.25) !important; border-radius: 0 !important;
  font-size: 30px !important; font-weight: 700 !important; padding: 10px 0 !important; letter-spacing: -.5px !important;
}
.gradio-container input[type="text"]:focus { border-bottom-color: #E5C15D !important; box-shadow: none !important; outline: none !important; }
/* Gradio 6 renders Textbox as a textarea inside label.container.show_textbox_border */
.gradio-container label.container.show_textbox_border { border: none !important; background: transparent !important; box-shadow: none !important; }
.gradio-container .input-container { background: transparent !important; border: none !important; }
.gradio-container textarea {
  background: transparent !important; border: none !important;
  border-bottom: 2px solid rgba(255,255,255,.25) !important; border-radius: 0 !important;
  font-size: 30px !important; font-weight: 700 !important; padding: 10px 0 !important; letter-spacing: -.5px !important;
  color: #F5F5F7 !important;
}
.gradio-container textarea:focus { border-bottom-color: #E5C15D !important; box-shadow: none !important; outline: none !important; }
.gradio-container label[data-testid$="-radio-label"] {
  border-radius: 999px !important; border: 1px solid rgba(255,255,255,.14) !important;
  padding: 7px 15px !important; background: transparent !important; color: #8A8A93 !important;
}
.gradio-container label[data-testid$="-radio-label"].selected { border-color: #E5C15D !important; color: #F5F5F7 !important; }
.gradio-container label[data-testid$="-radio-label"] span { color: inherit !important; background: transparent !important; letter-spacing: 0 !important; text-transform: none !important; font-size: 14px !important; }
.footer { opacity: .45; }
"""

TOPBAR = """<div class="topbar"><div class="brand">KEEP <i>•</i> Learn</div><div class="ondevice">On-device · CPU</div></div>"""
INTRO = """<div class="intro"><h1>Teach KEEP<br>a new word.</h1><p class="sub">No retraining. It keeps the <b>fingerprint</b> of your examples, then recognizes by resemblance. What you teach here lives only for this session.</p></div>"""

def learned_html(mem):
    if not mem:
        return '<p class="quiet" style="margin:8px 0 0">No word learned yet.</p>'
    chips = "".join(
        f'<span class="chip">{html.escape(w)} <span>×{len(v)}</span></span>' for w, v in mem.items())
    return f'<div class="chips">{chips}</div>'

def dots_html(mem, name):
    n = len((mem or {}).get((name or "").strip().lower(), []))
    dots = "".join('<i class="on"></i>' if i < min(n, 3) else '<i></i>' for i in range(3))
    label = "example" if n == 1 else "examples"
    return f'<div class="dots">{dots}<span>{n} {label} for this word</span></div>'

def qstage_empty():
    return """<div class="qstage"><div class="eyebrow">Test it</div><p class="word dim">Say it back</p></div>"""

def qstage_result(names, probs):
    rows = "".join(
        f'<div class="row"><b>{html.escape(w)}</b><span class="bar"><i style="width:{p * 100:.0f}%"></i></span><span class="pct">{p:.0%}</span></div>'
        for w, p in zip(names, probs))
    return (f"""<div class="qstage"><div class="eyebrow">KEEP thinks you said</div><p class="word">{html.escape(names[int(np.argmax(probs))])}</p>"""
            f"""<div class="bars">{rows}</div></div>""")

def add_example(name, audio, gain_choice, mem):
    mem = dict(mem or {}); name = (name or "").strip().lower()
    if not name:
        return mem, learned_html(mem), dots_html(mem, name), "Type the word name first.", gr.update()
    if audio is None:
        return mem, learned_html(mem), dots_html(mem, name), "Record the word first.", gr.update()
    R.gain = gain_choice; sr, wav = audio
    mem[name] = list(mem.get(name, [])) + [R.embedding(wav, sr).tolist()]
    n = len(mem[name]); tip = " Add 1 or 2 more examples to make it more sure." if n < 3 else ""
    return mem, learned_html(mem), dots_html(mem, name), f"Example no. {n} of “{name}” memorized.{tip}", None

def forget_all():
    return {}, learned_html({}), dots_html({}, ""), "Everything forgotten.", None

def recognize(audio, gain_choice, mem):
    if not mem or len(mem) < 2:
        return qstage_empty(), "Teach it at least **2 words** first (steps above)."
    if audio is None:
        return qstage_empty(), "Say one of the learned words."
    R.gain = gain_choice; sr, wav = audio; q = R.embedding(wav, sr)
    names = list(mem); P = torch.nn.functional.normalize(torch.stack([torch.tensor(mem[w]).mean(0) for w in names]), dim=-1)
    sim = (P @ q).numpy(); p = np.exp(20 * (sim - sim.max())); p = p / p.sum()
    i = int(sim.argmax()); j = np.argsort(-sim)[1] if len(sim) > 1 else i
    gap = sim[i] - sim[j]
    msg = (f"I heard: **{names[i]}**." if gap > 0.05 else f"I'm hesitating… rather **{names[i]}**, but it also sounds like “{names[j]}”. Add examples.")
    return qstage_result(names, p), msg

with gr.Blocks(title="KEEP Learn", css=CSS) as demo:
    gr.HTML(TOPBAR)
    gr.HTML(INTRO)
    mem = gr.State({})
    gr.HTML('<div class="step"><span class="num">01</span><h2>Name it</h2></div>')
    name = gr.Textbox(label="Word name", placeholder="e.g. bonjour", show_label=False)
    gr.HTML('<div class="step"><span class="num">02</span><h2>Say it</h2></div>')
    ex_audio = gr.Audio(sources=["microphone", "upload"], type="numpy", label="Say this word, 1 to 3 times")
    dots = gr.HTML(dots_html({}, ""))
    with gr.Row():
        b_add = gr.Button("Memorize", variant="primary"); b_forget = gr.Button("Forget everything")
    info = gr.Markdown(elem_classes=["quiet"])
    learned = gr.HTML(learned_html({}))
    gr.HTML('<div class="hairline"></div>')
    gr.HTML('<div class="step"><span class="num">03</span><h2>Test it</h2></div>')
    q_audio = gr.Audio(sources=["microphone", "upload"], type="numpy", label="Say one of the learned words again")
    q_stage = gr.HTML(qstage_empty())
    q_msg = gr.Markdown(elem_classes=["quiet"])
    q_btn = gr.Button("What did I say?", variant="primary")
    gain_radio = gr.Radio(["peak", "none"], value=cfg.get("gain", "peak"), label="Volume setting (peak = automatic, none = raw)")
    b_add.click(add_example, [name, ex_audio, gain_radio, mem], [mem, learned, dots, info, ex_audio])
    b_forget.click(forget_all, None, [mem, learned, dots, info, ex_audio])
    q_btn.click(recognize, [q_audio, gain_radio, mem], [q_stage, q_msg])
    gr.HTML('<div class="hairline"></div>')
    gr.Markdown("_Experimental: what it learns here is forgotten when you close the page. "
                "It works best with the same voice as the examples, and 2 to 3 examples per word._",
                elem_classes=["quiet"])
    gr.Markdown("**Limits.** It recognizes word sounds, it does not understand speech. A noisy microphone or words that sound very alike can lower accuracy.",
                elem_classes=["quiet"])
demo.launch()