import os
import gradio as gr
import torch
import numpy as np

# ── Constants ──────────────────────────────────────────────────────────────────
BASE_MODEL_ID = "openai/whisper-tiny"
LORA_MODEL_ID = "dungca/whisper-tiny-ja-lora"
SAMPLING_RATE = 16000
HF_TOKEN = os.environ.get("HF_TOKEN", None)

# ── Lazy model loading ─────────────────────────────────────────────────────────
_processor = None
_model = None

def get_model():
    global _processor, _model
    if _model is None:
        from transformers import AutoProcessor, WhisperForConditionalGeneration
        from peft import PeftModel

        print("Loading processor...")
        _processor = AutoProcessor.from_pretrained(BASE_MODEL_ID, token=HF_TOKEN)

        print("Loading base model...")
        base_model = WhisperForConditionalGeneration.from_pretrained(
            BASE_MODEL_ID, token=HF_TOKEN
        )

        print("Loading LoRA adapter...")
        _model = PeftModel.from_pretrained(base_model, LORA_MODEL_ID, token=HF_TOKEN)
        _model.eval()

        if torch.cuda.is_available():
            _model = _model.cuda()
            print("Using GPU ✓")
        else:
            print("Using CPU")

    return _processor, _model


# ── Inference ──────────────────────────────────────────────────────────────────
def transcribe(audio):
    if audio is None:
        return "⚠️ Vui lòng cung cấp audio (upload file hoặc record từ mic)."

    try:
        processor, model = get_model()
    except Exception as e:
        return f"❌ Lỗi load model: {str(e)}"

    sr, audio_array = audio

    audio_array = audio_array.astype(np.float32)
    if audio_array.ndim > 1:
        audio_array = audio_array.mean(axis=1)

    max_val = np.abs(audio_array).max()
    if max_val > 0:
        audio_array = audio_array / max_val

    if sr != SAMPLING_RATE:
        import librosa
        audio_array = librosa.resample(audio_array, orig_sr=sr, target_sr=SAMPLING_RATE)

    inputs = processor(audio_array, sampling_rate=SAMPLING_RATE, return_tensors="pt")

    if torch.cuda.is_available():
        inputs = {k: v.cuda() for k, v in inputs.items()}

    with torch.no_grad():
        predicted_ids = model.generate(
            input_features=inputs["input_features"],
            language="japanese",
            task="transcribe",
            max_new_tokens=256,
        )

    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
    return transcription.strip() if transcription.strip() else "（音声が認識できませんでした）"


# ── UI — Gradio 6.x compatible ────────────────────────────────────────────────
CSS = """
.gradio-container { max-width: 800px !important; margin: auto !important; }
footer { display: none !important; }
"""

with gr.Blocks(title="Whisper Japanese ASR Demo") as demo:

    gr.HTML("""
        <div style="text-align:center; margin-bottom:1.5rem">
            <h1 style="font-size:2rem; font-weight:700; margin-bottom:0.25rem">
                🎙️ Whisper Japanese ASR
            </h1>
            <p style="color:#666">
                LoRA fine-tuned on <b>ReazonSpeech</b> ·
                Base: <code>openai/whisper-tiny</code> ·
                Adapter: <a href="https://huggingface.co/dungca/whisper-tiny-ja-lora" target="_blank">
                dungca/whisper-tiny-ja-lora</a>
            </p>
        </div>
    """)

    with gr.Tab("🎤 Record từ Mic"):
        mic_input = gr.Audio(
            sources=["microphone"],
            type="numpy",
            label="Nói tiếng Nhật vào mic...",
        )
        mic_btn = gr.Button("📝 Transcribe", variant="primary")
        mic_output = gr.Textbox(
            label="Kết quả phiên âm (日本語)",
            placeholder="Kết quả sẽ hiện ở đây... (lần đầu load model ~30s)",
            lines=3,
        )
        mic_btn.click(fn=transcribe, inputs=mic_input, outputs=mic_output)

    with gr.Tab("📁 Upload File"):
        file_input = gr.Audio(
            sources=["upload"],
            type="numpy",
            label="Upload file audio (wav, mp3, m4a...)",
        )
        file_btn = gr.Button("📝 Transcribe", variant="primary")
        file_output = gr.Textbox(
            label="Kết quả phiên âm (日本語)",
            placeholder="Kết quả sẽ hiện ở đây... (lần đầu load model ~30s)",
            lines=3,
        )
        file_btn.click(fn=transcribe, inputs=file_input, outputs=file_output)

    gr.HTML("""
        <div style="margin-top:1.5rem; padding:1rem; background:#f8f9fa;
                    border-radius:8px; font-size:0.9rem; color:#555">
            <b>📊 Model Info</b><br>
            • Base: openai/whisper-tiny (39M params)<br>
            • Fine-tuning: LoRA (r=16, α=32) trên ReazonSpeech small<br>
            • CER: 0.525 · Eval loss: 1.177 · Trained on Kaggle P100<br>
            • <i>whisper-small version đang được train để cải thiện độ chính xác</i>
        </div>
    """)

if __name__ == "__main__":
    demo.launch(css=CSS, server_name="0.0.0.0", server_port=7860, show_error=True)