import gradio as gr
import torch
import numpy as np
import librosa
from transformers import (
    WhisperForConditionalGeneration, WhisperFeatureExtractor,
    WhisperProcessor, PreTrainedTokenizerFast,
)
from huggingface_hub import hf_hub_download

TARGET_SR = 16000

MODEL_IDS = {
    "NaijaVox-V1":  "Axiveri/NaijaVox-V1",
    "NaijaVox-2.0": "Axiveri/NaijaVox-2.0",
}

LANGUAGES = {
    "🇳🇬 Nigerian English": "<|en|>",
    "🇳🇬 Nigerian Pidgin":  "<|pcm|>",
    "🇳🇬 Yoruba":           "<|yo|>",
    "🇳🇬 Hausa":            "<|ha|>",
    "🇳🇬 Igbo":             "<|ig|>",
}

MODEL_CACHE = {}


def load_model(model_key):
    if model_key in MODEL_CACHE:
        return MODEL_CACHE[model_key]

    model_id = MODEL_IDS[model_key]
    print(f"Loading {model_key} from {model_id}...")
    model = WhisperForConditionalGeneration.from_pretrained(
        model_id, torch_dtype=torch.float32
    )

    # Try standard load first; fall back to manual tokenizer (same approach as V1 deploy)
    try:
        processor = WhisperProcessor.from_pretrained(model_id)
        vocab = processor.tokenizer.get_vocab()
        assert "<|pcm|>" in vocab and "<|ig|>" in vocab
    except Exception as e:
        print(f"  Standard load failed ({e}), using manual tokenizer...")
        fe  = WhisperFeatureExtractor.from_pretrained(model_id)
        tok = hf_hub_download(repo_id=model_id, filename="tokenizer.json")
        tokenizer = PreTrainedTokenizerFast(tokenizer_file=tok)
        tokenizer.add_special_tokens({
            "additional_special_tokens": [
                t for t in [
                    "<|startoftranscript|>", "<|endoftext|>", "<|transcribe|>",
                    "<|notimestamps|>", "<|en|>", "<|yo|>", "<|ha|>", "<|ig|>", "<|pcm|>",
                ]
                if t not in tokenizer.get_vocab()
            ]
        })
        processor = WhisperProcessor(feature_extractor=fe, tokenizer=tokenizer)
        vocab = processor.tokenizer.get_vocab()

    model.eval()
    MODEL_CACHE[model_key] = (model, processor, vocab)
    print(f"  {model_key} ready.")
    return model, processor, vocab


# Pre-load V1 on startup so the demo is immediately responsive
print("Pre-loading NaijaVox-V1...")
load_model("NaijaVox-V1")
print("Startup complete.")


def transcribe(audio, language, model_key):
    if audio is None:
        return "⚠️ Record or upload audio first, then tap Transcribe."

    try:
        model, processor, vocab = load_model(model_key)
    except Exception as e:
        return f"❌ Error loading {model_key}: {e}"

    sr, arr = audio
    arr = np.array(arr, dtype=np.float32)
    if arr.ndim > 1:
        arr = arr.mean(axis=1)
    if np.abs(arr).max() > 1.0:
        arr /= 32768.0
    if sr != TARGET_SR:
        arr = librosa.resample(arr, orig_sr=sr, target_sr=TARGET_SR)

    lang_id = vocab[LANGUAGES[language]]
    start   = vocab["<|startoftranscript|>"]
    trans   = vocab["<|transcribe|>"]
    nots    = vocab["<|notimestamps|>"]

    dec_ids = torch.tensor([[start, lang_id, trans, nots]])
    inputs  = processor.feature_extractor(
        arr, sampling_rate=TARGET_SR, return_tensors="pt"
    ).input_features

    with torch.no_grad():
        gen = model.generate(
            input_features=inputs,
            decoder_input_ids=dec_ids,
            max_new_tokens=200,
        )

    return processor.tokenizer.decode(gen[0], skip_special_tokens=True).strip()


# ── Theme ──────────────────────────────────────────────────────────────────────
GREEN = gr.themes.Color(
    c50="#f0fdf4", c100="#dcfce7", c200="#bbf7d0",
    c300="#86efac", c400="#4ade80", c500="#22c55e",
    c600="#16a34a", c700="#15803d", c800="#166534",
    c900="#14532d", c950="#052e16",
)

theme = gr.themes.Base(
    primary_hue=GREEN,
    secondary_hue=GREEN,
    neutral_hue=GREEN,
    font=gr.themes.GoogleFont("Space Grotesk"),
    font_mono=gr.themes.GoogleFont("JetBrains Mono"),
).set(
    body_background_fill="#060d07",
    body_background_fill_dark="#060d07",
    block_background_fill="#0d1a0f",
    block_background_fill_dark="#0d1a0f",
    block_border_color="#1c3824",
    block_border_color_dark="#1c3824",
    block_border_width="1px",
    block_radius="14px",
    button_primary_background_fill="linear-gradient(135deg,#16a34a 0%,#22c55e 100%)",
    button_primary_background_fill_hover="linear-gradient(135deg,#22c55e 0%,#4ade80 100%)",
    button_primary_text_color="#ffffff",
    button_primary_border_color="transparent",
    button_primary_border_color_hover="transparent",
    input_background_fill="#112015",
    input_background_fill_dark="#112015",
    input_border_color="#1c3824",
    input_border_color_focus="#22c55e",
    body_text_color="#dcfce7",
    body_text_color_dark="#dcfce7",
    body_text_color_subdued="#6ee7a0",
    block_label_text_color="#4ade80",
    block_label_text_color_dark="#4ade80",
    block_label_text_weight="600",
    block_title_text_color="#4ade80",
    block_title_text_color_dark="#4ade80",
    checkbox_label_background_fill="#112015",
    checkbox_label_background_fill_hover="#1a3020",
    checkbox_label_background_fill_selected="#0d2e1a",
    checkbox_label_border_color="#1c3824",
    checkbox_label_border_color_hover="#22c55e",
    checkbox_label_text_color="#dcfce7",
    table_even_background_fill="#0d1a0f",
    table_odd_background_fill="#112015",
    slider_color="#22c55e",
)

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;600&display=swap');

*, *::before, *::after {
    font-family: 'Space Grotesk', system-ui, sans-serif !important;
    box-sizing: border-box;
}

body, .gradio-container { background: #060d07 !important; }

.gradio-container {
    max-width: 820px !important;
    margin: 0 auto !important;
    padding: 0 1rem !important;
}

/* ── Header ── */
.nv-header { text-align: center; padding: 0.75rem 0 0.25rem; }

.nv-logo {
    display: block;
    margin: 0 auto 1.2rem;
    max-width: 100%;
    border-radius: 12px;
}

.nv-title {
    font-size: clamp(2.4rem, 6vw, 3.6rem);
    font-weight: 700;
    letter-spacing: -0.03em;
    background: linear-gradient(135deg, #22c55e 0%, #4ade80 55%, #86efac 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    line-height: 1.1;
    margin: 0 0 0.5rem;
}

.nv-sub {
    color: #6ee7a0;
    font-size: 0.95rem;
    margin: 0 0 1.8rem;
    letter-spacing: 0.01em;
}


/* ── Section labels ── */
.sl {
    color: #4ade80;
    font-size: 0.75rem;
    font-weight: 700;
    letter-spacing: 0.09em;
    text-transform: uppercase;
    margin: 1rem 0 0.35rem;
}

/* ── Model radio styled as pill tabs ── */
#model-radio .wrap { gap: 0.6rem !important; }

#model-radio label {
    background: #0d1a0f !important;
    border: 1.5px solid #1c3824 !important;
    border-radius: 10px !important;
    padding: 0.65rem 1rem !important;
    color: #6ee7a0 !important;
    font-weight: 500 !important;
    cursor: pointer !important;
    transition: all 0.16s !important;
    flex: 1 !important;
    text-align: center !important;
}

#model-radio label:hover {
    border-color: #22c55e !important;
    color: #dcfce7 !important;
    background: #112015 !important;
}

#model-radio label:has(input:checked) {
    border-color: #22c55e !important;
    background: linear-gradient(135deg,#0d2e1a,#112015) !important;
    color: #4ade80 !important;
    box-shadow: 0 0 0 1px #22c55e, 0 4px 16px rgba(34,197,94,0.18) !important;
}

/* ── Transcribe button ── */
#transcribe-btn {
    margin-top: 0.5rem !important;
}

#transcribe-btn button {
    font-size: 1.05rem !important;
    font-weight: 700 !important;
    letter-spacing: 0.04em !important;
    padding: 0.9rem !important;
    border-radius: 14px !important;
    box-shadow: 0 4px 20px rgba(34,197,94,0.28) !important;
    transition: all 0.18s ease !important;
}

#transcribe-btn button:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 28px rgba(34,197,94,0.4) !important;
}

#transcribe-btn button:active { transform: translateY(0) !important; }

/* ── Output ── */
#output-box textarea {
    font-size: 1.1rem !important;
    line-height: 1.75 !important;
    min-height: 96px !important;
    color: #dcfce7 !important;
}

/* ── Footer ── */
.nv-footer {
    text-align: center;
    padding: 1.5rem 0 2rem;
    color: #6ee7a0;
    font-size: 0.82rem;
    border-top: 1px solid #1c3824;
    margin-top: 1.5rem;
}

.nv-footer a { color: #4ade80; text-decoration: none; }
.nv-footer a:hover { text-decoration: underline; }


/* ── Prevent horizontal reflow from audio waveform ── */
.gradio-container { overflow-x: hidden !important; max-width: 100% !important; }
#audio-input { max-width: 100% !important; overflow-x: hidden !important; }
#audio-input canvas, #audio-input svg { max-width: 100% !important; }

/* ── Hide default Gradio component labels we replaced with .sl ── */
#audio-input > .block > span,
#audio-input label { display: none !important; }

/* ── Scrollbar ── */
::-webkit-scrollbar { width: 5px; height: 5px; }
::-webkit-scrollbar-track { background: #060d07; }
::-webkit-scrollbar-thumb { background: #1c3824; border-radius: 3px; }
::-webkit-scrollbar-thumb:hover { background: #22c55e; }
"""

# ── UI ─────────────────────────────────────────────────────────────────────────
with gr.Blocks(title="NaijaVox Demo — Nigerian ASR") as demo:

    # Header
    gr.HTML("""
    <div class="nv-header">
      <h1 class="nv-title">NaijaVox</h1>
      <p class="nv-sub">Yoruba · Hausa · Igbo · Pidgin · Nigerian English</p>
    </div>
    """)

    # Model selector
    gr.HTML('<p class="sl">Model</p>')
    model_radio = gr.Radio(
        choices=["NaijaVox-V1", "NaijaVox-2.0"],
        value="NaijaVox-2.0",
        label="",
        elem_id="model-radio",
    )

    # Language + Audio
    gr.HTML('<p class="sl">Language</p>')
    lang_dropdown = gr.Dropdown(
        choices=list(LANGUAGES.keys()),
        value="🇳🇬 Nigerian English",
        label="",
    )

    gr.HTML('<p class="sl">Audio</p>')
    audio_input = gr.Audio(
        sources=["microphone", "upload"],
        type="numpy",
        label="",
        elem_id="audio-input",
    )

    # Transcribe button
    transcribe_btn = gr.Button(
        "Transcribe",
        variant="primary",
        elem_id="transcribe-btn",
    )

    # Output
    gr.HTML('<p class="sl">Transcription</p>')
    output_box = gr.Textbox(
        label="",
        placeholder="Transcription will appear here...",
        lines=4,
        elem_id="output-box",
    )

    # Wire up
    transcribe_btn.click(
        fn=transcribe,
        inputs=[audio_input, lang_dropdown, model_radio],
        outputs=output_box,
    )

    # Footer
    gr.HTML("""
    <div class="nv-footer">
      <a href="https://huggingface.co/Axiveri/NaijaVox-V1">NaijaVox-V1</a>
      &nbsp;·&nbsp;
      <a href="https://huggingface.co/Axiveri/NaijaVox-2.0">NaijaVox-2.0</a>
      &nbsp;·&nbsp;
      <a href="https://huggingface.co/collections/Axiveri/naijavox-nigerian-speech-recognition">NaijaVox Collection</a>
      &nbsp;·&nbsp;
      Built by <a href="https://huggingface.co/ememzyvisuals">Emmanuel Ariyo</a> · Axiveri
    </div>
    """)

demo.launch(theme=theme, css=CSS)