"""
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

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

# 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 transformers import WhisperForConditionalGeneration, WhisperProcessor
    from src.engine.adapter_manager import AdapterManager

    _model_status = "loading…"
    try:
        _whisper_processor = WhisperProcessor.from_pretrained(
            WHISPER_MODEL_ID, token=HF_TOKEN
        )
        _whisper_model = WhisperForConditionalGeneration.from_pretrained(
            WHISPER_MODEL_ID,
            torch_dtype=torch.float32,
            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:
        if _whisper_model is None and "loading" not in _model_status and "error" not in _model_status:
            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. Intent + sensor data (CPU) ─────────────────────────────────────────
    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 = responder.generate_response(intent, sensor_data)

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

    return transcript, response_text, (sample_rate, wav_np)


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

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

    if not corrected_text.strip():
        return "⚠️ Corrected text is empty."

    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(),
        "response_text": response_text,
        "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)


# ── 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, response_text, audio_out = _run_pipeline(audio_path, language_code)
        return transcript, 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():

            # ── Tab 1: Voice Assistant ────────────────────────────────────────
            with gr.TabItem("🎙️ Voice Assistant"):
                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",
                            lines=3,
                            placeholder="Your words will appear here…",
                            interactive=False,
                        )
                        response_box = gr.Textbox(
                            label="Response / Jaabi",
                            lines=3,
                            placeholder="Agricultural advice will appear here…",
                            interactive=False,
                        )
                        audio_output = gr.Audio(
                            label="Voice response",
                            autoplay=True,
                            interactive=False,
                        )

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

            # ── Tab 2: Feedback & Correction ─────────────────────────────────
            with gr.TabItem("📝 Feedback & Correction"):
                gr.Markdown(
                    "Help improve the model by correcting transcription errors. "
                    "Your audio and corrections are saved to the training dataset."
                )
                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 (re-record or upload)",
                        )
                        fb_transcript = gr.Textbox(
                            label="Whisper output (what it heard)",
                            lines=3,
                            placeholder="Paste or type what Whisper said…",
                        )
                        fb_corrected = gr.Textbox(
                            label="Corrected transcription (what was actually said)",
                            lines=3,
                            placeholder="Type the correct text here…",
                        )

                    with gr.Column():
                        fb_response = gr.Textbox(
                            label="Response text (optional — for rating)",
                            lines=2,
                            placeholder="Copy the response from Tab 1…",
                        )
                        fb_rating = gr.Slider(
                            minimum=1, maximum=5, step=1, value=3,
                            label="Response 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="secondary")
                        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_response, fb_rating, fb_notes, fb_lang,
                    ],
                    outputs=[save_status],
                )

            # ── Tab 3: Training Status ────────────────────────────────────────
            with gr.TabItem("🔧 Training Status"):
                gr.Markdown(
                    "After collecting ≥10 corrections per language, run the training "
                    "notebook on Google Colab (free GPU), then reload adapters here."
                )
                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 Colab, run all cells."
                )
                gr.Markdown(
                    "**Feedback dataset**: "
                    f"`{FEEDBACK_REPO_ID}` (private, auto-updated on each save)"
                )
                gr.Markdown(
                    "**Adapter repo**: "
                    f"`{ADAPTER_REPO_ID}` (private, updated after each training run)"
                )

                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)

    # 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=9001,
        inbrowser=True,
        share=False,
    )