app.py

"""
Sahel-Agri Voice AI — HuggingFace Spaces (ZeroGPU)
Two-way voice assistant: Bambara / Fula / French / English → voice response

Environment variables (set in Space Settings → Secrets):
  HF_TOKEN          — HF write-access token
  FEEDBACK_REPO_ID  — e.g. ous-sow/sahel-agri-feedback  (dataset, private)
  ADAPTER_REPO_ID   — e.g. ous-sow/sahel-agri-adapters   (model, private)
  WHISPER_MODEL_ID  — default: openai/whisper-large-v3-turbo
                      (use openai/whisper-base for local CPU testing)
"""

from __future__ import annotations

import io
import json
import os
import sys
import tempfile
import threading
from datetime import datetime, timezone
from pathlib import Path

import gradio as gr
import numpy as np

ROOT = Path(__file__).parent
sys.path.insert(0, str(ROOT))

# ── env ───────────────────────────────────────────────────────────────────────
HF_TOKEN         = os.environ.get("HF_TOKEN")
FEEDBACK_REPO_ID = os.environ.get("FEEDBACK_REPO_ID", "ous-sow/sahel-agri-feedback")
ADAPTER_REPO_ID  = os.environ.get("ADAPTER_REPO_ID",  "ous-sow/sahel-agri-adapters")
# whisper-small: ~10s on cpu-basic, good multilingual quality.
# Override via WHISPER_MODEL_ID env var if you upgrade to a GPU Space later.
WHISPER_MODEL_ID = os.environ.get("WHISPER_MODEL_ID",  "openai/whisper-small")

# On local CPU (no HF_TOKEN / no spaces package) fall back gracefully
_ON_SPACES = os.environ.get("SPACE_ID") is not None

SUPPORTED_LANGUAGES = {
    "Bambara (bam)":    "bam",
    "Fula (ful)":       "ful",
    "French / Français": "fr",
    "English":          "en",
}

# ── ZeroGPU decorator (no-op locally) ────────────────────────────────────────
try:
    import spaces  # type: ignore
    _gpu = spaces.GPU(duration=55)
except ImportError:
    def _gpu(fn):           # local fallback: plain function
        return fn

# ── Module-level model state (CPU-resident between requests) ─────────────────
_whisper_model     = None   # WhisperForConditionalGeneration (base)
_whisper_processor = None
_adapter_manager   = None   # AdapterManager (wraps base model with PEFT if adapters loaded)
_model_lock        = threading.Lock()
_model_status      = "not loaded"
_adapters_loaded   = set()  # set of language codes with loaded adapters, e.g. {"bam", "ful"}

from src.tts.mms_tts                  import MMSTTSEngine
from src.iot.intent_parser             import IntentParser
from src.iot.sensor_bridge             import SensorBridge
from src.iot.voice_responder           import VoiceResponder
from src.conversation.phrase_matcher   import PhraseMatcher

_tts             = MMSTTSEngine()
_intent_parser   = IntentParser()
_sensor_bridge   = SensorBridge()
_phrase_matcher  = PhraseMatcher()

# HF API — only instantiate when token present
_hf_api = None
if HF_TOKEN:
    from huggingface_hub import HfApi
    _hf_api = HfApi(token=HF_TOKEN)


# ── Model loading ─────────────────────────────────────────────────────────────

def _do_load_whisper():
    global _whisper_model, _whisper_processor, _adapter_manager, _model_status
    import torch
    from src.engine.adapter_manager import AdapterManager

    # Import concrete Whisper classes directly — bypasses transformers __init__.py
    # Auto-class exports differ between transformers 4.x and 5.x; direct paths are stable.
    try:
        from transformers.models.whisper import WhisperProcessor, WhisperForConditionalGeneration
    except ImportError:
        from transformers.models.whisper.processing_whisper import WhisperProcessor
        from transformers.models.whisper.modeling_whisper import WhisperForConditionalGeneration

    _model_status = "loading…"
    try:
        _whisper_processor = WhisperProcessor.from_pretrained(
            WHISPER_MODEL_ID, token=HF_TOKEN
        )
        try:
            _whisper_model = WhisperForConditionalGeneration.from_pretrained(
                WHISPER_MODEL_ID,
                torch_dtype=torch.float32,
                token=HF_TOKEN,
            )
        except TypeError:
            _whisper_model = WhisperForConditionalGeneration.from_pretrained(
                WHISPER_MODEL_ID,
                token=HF_TOKEN,
            )
        _whisper_model.eval()

        # Create the AdapterManager wrapping the base model
        _adapter_manager = AdapterManager(base_model=_whisper_model, config={})

        # Try to load adapters from the local adapter repo snapshot (if already downloaded)
        _try_load_local_adapters()

        _model_status = f"ready ({WHISPER_MODEL_ID})"
    except Exception as e:
        _model_status = f"error: {e}"


def _try_load_local_adapters() -> None:
    """Load any adapter snapshots that are already on disk (downloaded previously)."""
    global _adapters_loaded
    if _adapter_manager is None:
        return
    if not ADAPTER_REPO_ID:
        return
    try:
        from huggingface_hub import try_to_load_from_cache
        lang_dirs = {"bam": "adapters/bambara", "ful": "adapters/fula"}
        for lang, subdir in lang_dirs.items():
            cached = try_to_load_from_cache(
                repo_id=ADAPTER_REPO_ID,
                filename=f"{subdir}/adapter_config.json",
                repo_type="model",
                token=HF_TOKEN,
            )
            if cached:
                import os
                adapter_path = str(os.path.dirname(cached))
                _adapter_manager.register(lang, adapter_path)
                try:
                    _adapter_manager.load_adapter(lang)
                    _adapters_loaded.add(lang)
                except Exception:
                    pass
    except Exception:
        pass  # Adapters not cached yet — will load after first Hub download


def _ensure_whisper_loaded():
    """Load Whisper to CPU in a background thread on first call. Non-blocking."""
    global _model_status
    with _model_lock:
        # Retry if previous attempt errored (e.g. import failed on first try)
        if _whisper_model is None and "loading" not in _model_status:
            _model_status = "loading…"
            t = threading.Thread(target=_do_load_whisper, daemon=True)
            t.start()
    return _model_status


def get_model_status() -> str:
    s = _ensure_whisper_loaded()
    if "ready" in s:
        return f"🟢 {s}"
    if "loading" in s:
        return f"🟡 {s}"
    if "error" in s:
        return f"🔴 {s}"
    return f"⚪ {s}"


# ── Core GPU pipeline ─────────────────────────────────────────────────────────

@_gpu
def _run_pipeline(audio_path: str, language_code: str):
    """
    Full STT → Intent → Sensor → TTS pipeline.
    Decorated with @spaces.GPU(duration=55) on HF Spaces; plain function locally.
    Returns: (transcript, response_text, (sample_rate, wav_np))
    """
    import asyncio
    import torch

    device = "cuda" if torch.cuda.is_available() else "cpu"

    # ── 1. Whisper STT ────────────────────────────────────────────────────────
    if _whisper_model is None:
        return "⏳ Model still loading…", "", None

    import librosa

    audio_np, _ = librosa.load(audio_path, sr=16000, mono=True)

    # Use adapter-wrapped model if an adapter for this language is loaded;
    # otherwise fall back to base Whisper.
    if _adapter_manager is not None and language_code in _adapters_loaded:
        _adapter_manager.activate(language_code)
        active_model = _adapter_manager.get_model()
    else:
        active_model = _whisper_model

    active_model.to(device)
    with _model_lock:
        inputs = _whisper_processor.feature_extractor(
            audio_np, sampling_rate=16000, return_tensors="pt"
        )
        input_features = inputs.input_features.to(device)

        # Bambara and Fula have no Whisper language token — pass None so the model
        # auto-detects or falls back to multilingual decoding.
        if language_code in ("bam", "ful"):
            forced_ids = None
        else:
            forced_ids = _whisper_processor.get_decoder_prompt_ids(
                language=language_code, task="transcribe"
            )

        with torch.no_grad():
            predicted_ids = active_model.generate(
                input_features,
                forced_decoder_ids=forced_ids if forced_ids else None,
                max_new_tokens=256,
            )

    transcript = _whisper_processor.batch_decode(
        predicted_ids, skip_special_tokens=True
    )[0].strip()

    # Free GPU VRAM before TTS
    active_model.to("cpu")
    if device == "cuda":
        torch.cuda.empty_cache()

    # ── 2. Phrase library (general conversation — no sensors needed) ─────────
    phrase_match = _phrase_matcher.match(transcript, language_code)
    if phrase_match:
        response_text      = phrase_match["response"]
        english_translation = phrase_match["english"]
    else:
        # ── 3. Intent + sensor data (agricultural queries) ────────────────────
        intent = _intent_parser.parse(transcript, language=language_code)

        try:
            loop = asyncio.new_event_loop()
            sensor_data = loop.run_until_complete(_sensor_bridge.fetch(intent))
            loop.close()
        except Exception:
            from src.iot.sensor_bridge import SensorData
            sensor_data = SensorData(sensor_type="soil", values={
                "moisture_pct": 45.0, "ph": 6.5, "temperature_c": 28.0
            })

        responder = VoiceResponder(language=language_code)
        response_text, english_translation = responder.generate_response(intent, sensor_data)

        # Low-confidence fallback: say "I didn't understand" in the native language
        if intent.action == "unknown" and intent.confidence < 0.15:
            from src.iot.voice_responder import BAMBARA_TEMPLATES, FULA_TEMPLATES
            if language_code == "bam":
                response_text, english_translation = BAMBARA_TEMPLATES["not_understood"]
            elif language_code == "ful":
                response_text, english_translation = FULA_TEMPLATES["not_understood"]

    # ── 3. MMS-TTS (GPU) ──────────────────────────────────────────────────────
    wav_np, sample_rate = _tts.synthesize(response_text, language_code, device=device)

    return transcript, english_translation, response_text, (sample_rate, wav_np)


# ── HF Hub feedback persistence ───────────────────────────────────────────────

def _save_feedback_to_hub(
    audio_path: str | None,
    transcript: str,
    corrected_text: str,
    english_translation: str,
    corrected_english: str,
    response_text: str,
    corrected_response: str,
    rating: int,
    notes: str,
    language_label: str,
) -> str:
    language_code = SUPPORTED_LANGUAGES.get(language_label, "bam")

    if not corrected_text.strip():
        return "⚠️ Corrected transcription is empty — please fill in what was actually said."

    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S_%f")

    record = {
        "id": timestamp,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "language": language_code,
        "audio_file": f"audio/{language_code}_{timestamp}.wav",
        "whisper_output": transcript,
        "corrected_text": corrected_text.strip(),
        "english_translation": english_translation.strip(),
        "corrected_english": corrected_english.strip() or english_translation.strip(),
        "response_text": response_text,
        "corrected_response": corrected_response.strip() or response_text.strip(),
        "rating": rating,
        "notes": notes.strip(),
        "is_correction": transcript.strip() != corrected_text.strip(),
        "model": WHISPER_MODEL_ID,
    }

    if _hf_api is None:
        # Local: save to disk instead
        fb_dir = ROOT / "feedback"
        fb_dir.mkdir(exist_ok=True)
        (fb_dir / "audio").mkdir(exist_ok=True)
        corrections_path = fb_dir / "corrections.jsonl"
        if audio_path:
            import shutil
            shutil.copy2(audio_path, fb_dir / "audio" / f"{language_code}_{timestamp}.wav")
        with open(corrections_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
        total = sum(1 for _ in open(corrections_path, encoding="utf-8"))
        return f"✅ Saved locally (#{total}) — HF_TOKEN not set, Hub upload skipped."

    try:
        # Upload audio
        if audio_path:
            _hf_api.upload_file(
                path_or_fileobj=audio_path,
                path_in_repo=f"audio/{language_code}_{timestamp}.wav",
                repo_id=FEEDBACK_REPO_ID,
                repo_type="dataset",
            )

        # Download → append → re-upload corrections.jsonl (with retry on conflict)
        from huggingface_hub import hf_hub_download
        for attempt in range(2):
            try:
                local_jsonl = hf_hub_download(
                    repo_id=FEEDBACK_REPO_ID,
                    filename="corrections.jsonl",
                    repo_type="dataset",
                    token=HF_TOKEN,
                )
                with open(local_jsonl, encoding="utf-8") as f:
                    existing = f.read()
            except Exception:
                existing = ""

            updated = existing + json.dumps(record, ensure_ascii=False) + "\n"
            buf = io.BytesIO(updated.encode("utf-8"))

            try:
                _hf_api.upload_file(
                    path_or_fileobj=buf,
                    path_in_repo="corrections.jsonl",
                    repo_id=FEEDBACK_REPO_ID,
                    repo_type="dataset",
                )
                break
            except Exception as e:
                if attempt == 1:
                    return f"⚠️ Audio uploaded but corrections.jsonl update failed: {e}"

        total = updated.count("\n")
        return f"✅ Saved to Hub (#{total}) — {FEEDBACK_REPO_ID}"

    except Exception as e:
        return f"❌ Hub upload error: {e}"


# ── Adapter reload ────────────────────────────────────────────────────────────

def _reload_adapters_from_hub() -> str:
    global _adapters_loaded
    if _hf_api is None:
        return "⚠️ HF_TOKEN not set — cannot download adapters."
    if _adapter_manager is None:
        return "⏳ Base model not loaded yet — wait for model to finish loading and try again."
    try:
        from huggingface_hub import snapshot_download
        local_dir = snapshot_download(
            repo_id=ADAPTER_REPO_ID, repo_type="model", token=HF_TOKEN
        )
        results = []
        for lang, subdir in (("bam", "adapters/bambara"), ("ful", "adapters/fula")):
            adapter_path = Path(local_dir) / subdir
            if not adapter_path.exists():
                results.append(f"⚠️ {lang}: `{subdir}` not found in repo")
                continue
            # Check that this looks like a valid PEFT adapter
            if not (adapter_path / "adapter_config.json").exists():
                results.append(f"⚠️ {lang}: `{subdir}` missing adapter_config.json — run training first")
                continue
            try:
                _adapter_manager.register(lang, str(adapter_path))
                _adapter_manager.load_adapter(lang)
                _adapters_loaded.add(lang)
                results.append(f"✅ {lang}: adapter loaded from `{subdir}`")
            except Exception as e:
                results.append(f"❌ {lang}: load failed — {e}")

        summary = "\n".join(results)
        active = ", ".join(_adapters_loaded) if _adapters_loaded else "none"
        return f"{summary}\n\n**Active adapters:** {active}\n**Repo:** `{ADAPTER_REPO_ID}`"
    except Exception as e:
        return f"❌ Adapter reload failed: {e}"


def _get_adapter_status() -> str:
    lines = []

    # Show which adapters are currently active in memory
    if _adapters_loaded:
        lines.append(f"**Active adapters (in memory):** {', '.join(sorted(_adapters_loaded))}")
    else:
        lines.append("**Active adapters:** none — using base Whisper")

    if _hf_api is None:
        lines.append("_HF_TOKEN not set — Hub check skipped._")
        return "\n".join(lines)

    try:
        from huggingface_hub import list_repo_files
        files = list(list_repo_files(ADAPTER_REPO_ID, repo_type="model", token=HF_TOKEN))
        bam_ok = any("bambara" in f and "adapter_config" in f for f in files)
        ful_ok = any("fula" in f and "adapter_config" in f for f in files)
        lines += [
            f"\n**Hub repo:** `{ADAPTER_REPO_ID}`",
            f"- Bambara (bam): {'✅ trained adapter present' if bam_ok else '⚠️ not yet trained — run bootstrap notebook'}",
            f"- Fula (ful): {'✅ trained adapter present' if ful_ok else '⚠️ not yet trained — run bootstrap notebook'}",
        ]
        if bam_ok or ful_ok:
            lines.append("\n_Click **Reload Adapters** to activate them._")
    except Exception as e:
        lines.append(f"_Could not read Hub repo: {e}_")

    return "\n".join(lines)


# ── Knowledge Base handlers ───────────────────────────────────────────────────

def _import_phrase_pairs(lang_label: str, pairs_text: str) -> str:
    """Import pasted phrase pairs into the phrase library."""
    if not pairs_text.strip():
        return "⚠️ Nothing entered. Use the format:  native phrase | english translation"
    lang = SUPPORTED_LANGUAGES.get(lang_label, "bam")
    count = _phrase_matcher.import_pairs(lang, pairs_text)
    if count == 0:
        return "⚠️ No valid phrases found. Each line must contain a | separator.\nExample:  I ni ce | Hello, good day"
    _upload_phrase_additions_to_hub(lang)
    total = _phrase_matcher.phrase_count(lang)
    return f"✅ Added {count} phrase(s) for {lang_label}. Library now has {total} phrases. Available immediately."


def _upload_phrase_additions_to_hub(lang: str) -> None:
    """Persist user phrase additions to HF Hub so they survive Space restarts."""
    if _hf_api is None or not FEEDBACK_REPO_ID:
        return
    try:
        import io
        data = _phrase_matcher.get_additions_json(lang)
        buf = io.BytesIO(data.encode("utf-8"))
        _hf_api.upload_file(
            path_or_fileobj=buf,
            path_in_repo=f"phrase_additions/{lang}.json",
            repo_id=FEEDBACK_REPO_ID,
            repo_type="dataset",
        )
    except Exception as exc:
        import logging
        logging.getLogger(__name__).warning("Could not upload phrase additions: %s", exc)


def _load_phrase_additions_from_hub() -> None:
    """Download and merge user phrase additions from HF Hub at startup."""
    if _hf_api is None or not FEEDBACK_REPO_ID:
        return
    for lang in ("bam", "ful"):
        try:
            from huggingface_hub import hf_hub_download
            local = hf_hub_download(
                repo_id=FEEDBACK_REPO_ID,
                filename=f"phrase_additions/{lang}.json",
                repo_type="dataset",
                token=HF_TOKEN,
            )
            with open(local, encoding="utf-8") as f:
                data = f.read()
            _phrase_matcher.reload_from_hub_data(lang, data)
        except Exception:
            pass  # No additions saved yet — fine

# Load user phrase additions in background at module import time
threading.Thread(target=_load_phrase_additions_from_hub, daemon=True).start()


def _save_audio_for_training(lang_label: str, audio_path: str | None, transcript: str, source_note: str) -> str:
    """Save an uploaded audio file + transcription as a training sample to HF Hub."""
    transcript = transcript.strip()
    if audio_path is None:
        return "⚠️ Please upload an audio file first."
    if not transcript:
        return "⚠️ Please type the transcription — what is said in this audio."

    lang = SUPPORTED_LANGUAGES.get(lang_label, "bam")
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    audio_repo_path = f"training_audio/{lang}/{timestamp}.wav"
    meta_repo_path  = f"training_audio/{lang}/{timestamp}.txt"

    if _hf_api is None or not FEEDBACK_REPO_ID:
        return "⚠️ HF_TOKEN not set — file saved locally only, not uploaded to Hub."

    try:
        import io
        _hf_api.upload_file(
            path_or_fileobj=audio_path,
            path_in_repo=audio_repo_path,
            repo_id=FEEDBACK_REPO_ID,
            repo_type="dataset",
        )
        meta = (
            f"language: {lang}\n"
            f"transcription: {transcript}\n"
            f"source: {source_note.strip() or 'uploaded'}\n"
            f"timestamp: {timestamp}\n"
        )
        _hf_api.upload_file(
            path_or_fileobj=io.BytesIO(meta.encode()),
            path_in_repo=meta_repo_path,
            repo_id=FEEDBACK_REPO_ID,
            repo_type="dataset",
        )
        return (
            f"✅ Saved to training dataset!\n"
            f"Audio: {audio_repo_path}\n"
            f"Transcription: {transcript[:80]}{'…' if len(transcript) > 80 else ''}\n"
            f"Run the training notebook on Kaggle to include this in the next model update."
        )
    except Exception as exc:
        return f"❌ Upload failed: {exc}"


# ── Main ask handler ──────────────────────────────────────────────────────────

def handle_ask(audio_path, language_label):
    if audio_path is None:
        return "⚠️ No audio — press Record or upload a file.", "", "", None

    language_code = SUPPORTED_LANGUAGES.get(language_label, "bam")
    status = _ensure_whisper_loaded()

    if _whisper_model is None:
        return f"⏳ Model loading ({status}). Wait a moment and try again.", "", "", None

    try:
        transcript, english_translation, response_text, audio_out = _run_pipeline(audio_path, language_code)
        return transcript, english_translation, response_text, audio_out
    except Exception as e:
        return f"❌ {e}", "", "", None


# ── Gradio UI ─────────────────────────────────────────────────────────────────

def build_ui() -> gr.Blocks:
    with gr.Blocks(title="Sahel-Agri Voice AI") as demo:
        gr.Markdown("# 🌾 Sahel-Agri Voice AI")
        gr.Markdown(
            "Speak in **Bambara** or **Fula** — get agricultural insights spoken back "
            "in your language. Also supports French and English."
        )

        model_status_box = gr.Textbox(
            value=get_model_status(),
            label="Model status",
            interactive=False,
        )
        # gr.Timer polls get_model_status every 3s and updates the box (Gradio 5)
        status_timer = gr.Timer(value=3)
        status_timer.tick(fn=get_model_status, outputs=model_status_box)

        with gr.Tabs() as tabs:

            # ── Tab 1: Voice Assistant ────────────────────────────────────────
            with gr.TabItem("🎙️ Voice Assistant", id="tab_voice"):
                with gr.Row():
                    with gr.Column(scale=1):
                        language_dd = gr.Dropdown(
                            choices=list(SUPPORTED_LANGUAGES.keys()),
                            value="Bambara (bam)",
                            label="Language / Kan",
                        )
                        audio_input = gr.Audio(
                            sources=["microphone", "upload"],
                            type="filepath",
                            label="Record or upload audio",
                        )
                        ask_btn = gr.Button("▶ Ask / Ɲinɛ", variant="primary")

                    with gr.Column(scale=1):
                        transcript_box = gr.Textbox(
                            label="Whisper heard (transcription)",
                            lines=2,
                            placeholder="Your words will appear here…",
                            interactive=False,
                        )
                        translation_box = gr.Textbox(
                            label="English translation",
                            lines=2,
                            placeholder="English meaning will appear here…",
                            interactive=False,
                        )
                        response_box = gr.Textbox(
                            label="Response in your language",
                            lines=2,
                            placeholder="Agricultural advice will appear here…",
                            interactive=False,
                        )
                        audio_output = gr.Audio(
                            label="Voice response",
                            autoplay=True,
                            interactive=False,
                        )
                        correct_btn = gr.Button(
                            "✏️ Something wrong? Send to Correction tab",
                            variant="secondary",
                            size="sm",
                        )

                ask_btn.click(
                    fn=handle_ask,
                    inputs=[audio_input, language_dd],
                    outputs=[transcript_box, translation_box, response_box, audio_output],
                )

            # ── Tab 2: Feedback & Correction ─────────────────────────────────
            with gr.TabItem("📝 Feedback & Correction", id="tab_feedback"):
                gr.Markdown(
                    "Correct what Whisper heard, the English translation, and the response. "
                    "All corrections are saved to the training dataset to improve future accuracy."
                )
                with gr.Row():
                    with gr.Column():
                        fb_lang = gr.Dropdown(
                            choices=list(SUPPORTED_LANGUAGES.keys()),
                            value="Bambara (bam)",
                            label="Language",
                        )
                        fb_audio = gr.Audio(
                            sources=["microphone", "upload"],
                            type="filepath",
                            label="Audio",
                        )
                        gr.Markdown("**Step 1 — Fix the transcription**")
                        fb_transcript = gr.Textbox(
                            label="What Whisper heard",
                            lines=2,
                            placeholder="Auto-filled from Tab 1…",
                        )
                        fb_corrected = gr.Textbox(
                            label="✏️ What was actually said (in Bambara/Fula)",
                            lines=2,
                            placeholder="Type the correct transcription here…",
                        )

                    with gr.Column():
                        gr.Markdown("**Step 2 — Fix the English translation**")
                        fb_english = gr.Textbox(
                            label="Auto-generated English translation",
                            lines=2,
                            placeholder="Auto-filled from Tab 1…",
                        )
                        fb_corrected_english = gr.Textbox(
                            label="✏️ Correct English translation",
                            lines=2,
                            placeholder="Type the correct English meaning here…",
                        )
                        gr.Markdown("**Step 3 — Fix the response**")
                        fb_response = gr.Textbox(
                            label="Auto-generated response",
                            lines=2,
                            placeholder="Auto-filled from Tab 1…",
                        )
                        fb_corrected_response = gr.Textbox(
                            label="✏️ Better response (in farmer's language)",
                            lines=2,
                            placeholder="Type a better response here…",
                        )
                        fb_rating = gr.Slider(
                            minimum=1, maximum=5, step=1, value=3,
                            label="Overall quality (1 = poor, 5 = excellent)",
                        )
                        fb_notes = gr.Textbox(
                            label="Notes (optional)",
                            lines=2,
                            placeholder="e.g. noisy background, strong accent…",
                        )
                        save_btn = gr.Button("💾 Save to Dataset", variant="primary")
                        save_status = gr.Textbox(
                            label="Save status", interactive=False, lines=2
                        )

                save_btn.click(
                    fn=_save_feedback_to_hub,
                    inputs=[
                        fb_audio, fb_transcript, fb_corrected,
                        fb_english, fb_corrected_english,
                        fb_response, fb_corrected_response,
                        fb_rating, fb_notes, fb_lang,
                    ],
                    outputs=[save_status],
                )

            # Wire "Send to Correction" button — populates Tab 2 fields from Tab 1
            correct_btn.click(
                fn=lambda t, tr, r, lang: (t, t, tr, tr, r, r, lang),
                inputs=[transcript_box, translation_box, response_box, language_dd],
                outputs=[fb_transcript, fb_corrected, fb_english, fb_corrected_english,
                         fb_response, fb_corrected_response, fb_lang],
            )

            # ── Tab 3: Knowledge Base ─────────────────────────────────────────
            with gr.TabItem("📚 Knowledge Base"):
                gr.Markdown(
                    "## Teach the assistant new phrases — no technical knowledge required\n\n"
                    "Add phrases the assistant should recognise and respond to. "
                    "Changes take effect **immediately** and are saved to the Hub so they survive restarts."
                )
                with gr.Row():
                    # ── Left: phrase pair import ──────────────────────────────
                    with gr.Column():
                        gr.Markdown(
                            "### ➕ Add phrases manually\n"
                            "One phrase per line in the format:\n"
                            "```\nnative phrase | English translation\n```\n"
                            "**Examples (Bambara):**\n"
                            "```\nI ni ce | Hello, good day\n"
                            "Sanji bɛ na | Rain is coming\n"
                            "N bɛ i dɛmɛ | I will help you\n```\n"
                            "**Examples (Fula):**\n"
                            "```\nJam waali | Hello, peace be with you\n"
                            "Ndiyam wadata | Rain is coming\n"
                            "Mi woni ɗoo | I am here\n```"
                        )
                        kb_lang = gr.Dropdown(
                            choices=["Bambara (bam)", "Fula (ful)"],
                            value="Bambara (bam)",
                            label="Language",
                        )
                        kb_pairs = gr.Textbox(
                            lines=10,
                            placeholder="I ni ce | Hello, good day\nI ni sogoma | Good morning\nSanji bɛ na | Rain is coming",
                            label="Phrase pairs (native | english) — one per line",
                        )
                        kb_import_btn = gr.Button("➕ Add to Knowledge Base", variant="primary")
                        kb_status = gr.Textbox(label="Status", interactive=False, lines=3)

                    # ── Right: audio upload for training ─────────────────────
                    with gr.Column():
                        gr.Markdown(
                            "### 🎬 Add audio from YouTube (or anywhere)\n"
                            "HuggingFace Spaces cannot download YouTube directly, "
                            "so convert the video to audio first on your computer:\n\n"
                            "**Free online converters:**\n"
                            "- [ytmp3.cc](https://ytmp3.cc) — paste YouTube URL → download MP3\n"
                            "- [cobalt.tools](https://cobalt.tools) — paste any video URL → download audio\n"
                            "- [y2mate.com](https://y2mate.com) — paste YouTube URL → download MP3\n\n"
                            "**Good YouTube search terms:**\n"
                            "- Bambara: *'Bamanankan conversation'*, *'Bambara leçon'*, *'donsomana'*\n"
                            "- Fula: *'Fulfulde leçon'*, *'Pular conversation'*, *'Fula radio'*\n\n"
                            "Then upload the MP3/WAV file below with its transcription."
                        )
                        yt_lang = gr.Dropdown(
                            choices=["Bambara (bam)", "Fula (ful)"],
                            value="Bambara (bam)",
                            label="Language spoken in the audio",
                        )
                        yt_audio = gr.Audio(
                            sources=["upload"],
                            type="filepath",
                            label="Upload audio file (MP3 or WAV)",
                        )
                        yt_transcript = gr.Textbox(
                            lines=5,
                            placeholder="Type what is said in the audio (as much as you can).\n"
                                        "Example:\nJam waali. No mbadda. Mi woni ɗoo wallude ma.",
                            label="Transcription — what is said in this audio",
                        )
                        yt_source = gr.Textbox(
                            placeholder="e.g. YouTube: Bambara lesson by Moussa Kouyaté",
                            label="Source (optional — for your records)",
                        )
                        yt_btn = gr.Button("💾 Save Audio for Training", variant="secondary")
                        yt_status = gr.Textbox(label="Status", interactive=False, lines=4)

                kb_import_btn.click(
                    fn=_import_phrase_pairs,
                    inputs=[kb_lang, kb_pairs],
                    outputs=[kb_status],
                )
                yt_btn.click(
                    fn=_save_audio_for_training,
                    inputs=[yt_lang, yt_audio, yt_transcript, yt_source],
                    outputs=[yt_status],
                )

            # ── Tab 4: Model Training ─────────────────────────────────────────
            with gr.TabItem("🔧 Model Training"):
                gr.Markdown(
                    "After collecting audio corrections and YouTube samples, "
                    "run the training notebook to fine-tune the speech model."
                )
                adapter_status_md = gr.Markdown(value=_get_adapter_status())
                reload_btn = gr.Button("🔄 Reload Adapters from Hub")
                reload_out = gr.Markdown()

                gr.Markdown("---")
                gr.Markdown(
                    "**Training notebook**: "
                    "`notebooks/train_colab.ipynb` — open in Kaggle or Colab, run all cells.\n\n"
                    "**Feedback dataset**: "
                    f"`{FEEDBACK_REPO_ID}` (auto-updated on each save)\n\n"
                    "**Adapter repo**: "
                    f"`{ADAPTER_REPO_ID}` (updated after training)"
                )

                reload_btn.click(fn=_reload_adapters_from_hub, outputs=[reload_out])
                reload_btn.click(fn=_get_adapter_status, outputs=[adapter_status_md])

    return demo


# ── Entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    from dotenv import load_dotenv
    load_dotenv()

    # Re-read env after dotenv
    HF_TOKEN         = os.environ.get("HF_TOKEN")
    FEEDBACK_REPO_ID = os.environ.get("FEEDBACK_REPO_ID", "ous-sow/sahel-agri-feedback")
    ADAPTER_REPO_ID  = os.environ.get("ADAPTER_REPO_ID",  "ous-sow/sahel-agri-adapters")
    WHISPER_MODEL_ID = os.environ.get("WHISPER_MODEL_ID",  "openai/whisper-small")

    if HF_TOKEN:
        from huggingface_hub import HfApi
        _hf_api = HfApi(token=HF_TOKEN)

    # Load any previously saved phrase additions from HF Hub
    _load_phrase_additions_from_hub()

    # Kick off background model load immediately
    _ensure_whisper_loaded()

    print(f"Whisper model : {WHISPER_MODEL_ID}")
    print(f"Feedback repo : {FEEDBACK_REPO_ID}")
    print(f"Adapter repo  : {ADAPTER_REPO_ID}")
    print(f"HF_TOKEN set  : {'yes' if HF_TOKEN else 'no (local-only mode)'}")
    print()

    demo = build_ui()
    demo.launch(
        server_port=7860,  # HF Spaces standard port
        inbrowser=False,
        share=False,
        show_api=False,
        ssr_mode=False,    # SSR starts a Node.js process that hangs in HF Spaces containers
    )