#!/usr/bin/env python3
"""PolyWhisper — Gradio demo for multilingual Indic ASR.

Run locally:
    python app.py                    # local only
    python app.py --share            # public URL via gradio.live
    python app.py --share --port 7860
"""

import os
os.environ.setdefault("TRANSFORMERS_NO_ADVISORY_WARNINGS", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import gradio as gr
import numpy as np
import time
import json
from pathlib import Path

# Global model cache
_model = None
_processor = None

LANGUAGES = {
    "Hindi (हिन्दी)": "hi",
    "Tamil (தமிழ்)": "ta",
    "Telugu (తెలుగు)": "te",
    "Bengali (বাংলা)": "bn",
    "Marathi (मराठी)": "mr",
}

SAMPLE_AUDIO_DIR = Path(__file__).parent / "samples"


def load_model():
    global _model, _processor
    if _model is None:
        from polywhisper.model import PolyWhisper
        _model = PolyWhisper(backbone="small", device="cpu")
        # Pre-load all adapters
        for lang in ["hi", "ta", "te", "bn", "mr"]:
            try:
                _model.auto_load_adapter(lang, variant="prod")
            except FileNotFoundError:
                pass
        _processor = _model.processor
    return _model, _processor


def transcribe_audio(audio, language, max_tokens, beam_width, output_format):
    """Transcribe uploaded or recorded audio."""
    if audio is None:
        return "Please upload or record audio first.", "", None

    lang_code = LANGUAGES.get(language, "hi")
    model, processor = load_model()

    # Load audio
    import soundfile as sf
    if isinstance(audio, tuple):
        # Gradio returns (sample_rate, numpy_array) for recorded audio
        sr, audio_array = audio
        audio_array = audio_array.astype(np.float32)
        if audio_array.ndim > 1:
            audio_array = audio_array.mean(axis=1)
        if sr != 16000:
            import resampy
            audio_array = resampy.resample(audio_array, sr, 16000)
    else:
        # File path
        audio_array, sr = sf.read(str(audio), dtype="float32")
        if audio_array.ndim > 1:
            audio_array = audio_array.mean(axis=1)
        if sr != 16000:
            import resampy
            audio_array = resampy.resample(audio_array, sr, 16000)

    # Transcribe
    t0 = time.time()
    result = _transcribe(model, processor, audio_array, lang_code, max_tokens, int(beam_width))
    dt = time.time() - t0

    # Format output
    text = result["text"]
    duration = result["duration"]

    # Build output based on format
    if output_format == "Plain Text":
        output_text = text
        srt_output = None
    elif output_format == "JSON":
        output_text = json.dumps(result, ensure_ascii=False, indent=2)
        srt_output = None
    elif output_format == "SRT Subtitles":
        output_text = text
        srt_lines = []
        for i, seg in enumerate(result.get("segments", []), 1):
            start = _fmt_srt(seg["start"])
            end = _fmt_srt(seg["end"])
            srt_lines.append(f"{i}\n{start} --> {end}\n{seg['text']}\n")
        srt_output = "\n".join(srt_lines)
    else:
        output_text = text
        srt_output = None

    # Stats
    stats = f"Language: {language} | Duration: {duration:.1f}s | Processed in: {dt:.1f}s | Speed: {duration/dt:.1f}× realtime"

    return output_text, stats, srt_output


def _transcribe(model, processor, audio, lang, max_new_tokens=256, num_beams=1):
    """Core transcription function."""
    import torch
    model.set_language(lang)

    feats = processor.feature_extractor(
        [audio], sampling_rate=16000, return_tensors="pt", padding=True
    )["input_features"]
    if feats.shape[-1] < 3000:
        feats = torch.cat(
            [feats, torch.zeros(1, 80, 3000 - feats.shape[-1], dtype=feats.dtype)], -1
        )

    out = model.generate(
        feats, lang=lang, max_new_tokens=max_new_tokens,
        num_beams=num_beams, use_cache=True, task="transcribe",
    )
    text = processor.decode(out[0], skip_special_tokens=True).strip()

    return {
        "text": text,
        "lang": lang,
        "duration": len(audio) / 16000,
        "segments": [{"text": text, "start": 0.0, "end": len(audio) / 16000}],
    }


def _fmt_srt(sec):
    h = int(sec // 3600)
    m = int((sec % 3600) // 60)
    s = int(sec % 60)
    ms = int((sec - int(sec)) * 1000)
    return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"


# ===== Custom CSS =====
CUSTOM_CSS = """
.main-title {
    text-align: center;
    margin-bottom: 0.5em;
}
.main-title h1 {
    font-size: 2.2em;
    font-weight: 700;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 0.2em;
}
.main-title p {
    font-size: 1.1em;
    color: #666;
    margin-top: 0;
}
.lang-badge {
    display: inline-block;
    padding: 2px 8px;
    border-radius: 12px;
    font-size: 0.85em;
    font-weight: 500;
    margin: 2px;
}
.stats-box {
    background: #f0f7ff;
    border: 1px solid #b3d4fc;
    border-radius: 8px;
    padding: 8px 12px;
    font-family: monospace;
    font-size: 0.9em;
}
footer { display: none !important; }
"""


def build_app():
    """Build the Gradio interface."""
    with gr.Blocks(
        title="PolyWhisper — Indic ASR Demo",
    ) as app:
        # Header
        gr.HTML("""
        <div class="main-title">
            <h1>PolyWhisper</h1>
            <p>Efficient Multilingual Indic Speech Recognition</p>
            <p style="font-size:0.9em; color:#888;">
                Hindi · Tamil · Telugu · Bengali · Marathi
                &nbsp;|&nbsp; Powered by Whisper Small + LoRA
            </p>
        </div>
        """)

        with gr.Row():
            # Left column: Input
            with gr.Column(scale=1):
                audio_input = gr.Audio(
                    label="Upload or Record Audio",
                    sources=["upload", "microphone"],
                    type="numpy",
                    waveform_options={"waveform_color": "#667eea"},
                )

                with gr.Row():
                    language = gr.Dropdown(
                        choices=list(LANGUAGES.keys()),
                        value="Hindi (हिन्दी)",
                        label="Language",
                        scale=2,
                    )
                    output_format = gr.Dropdown(
                        choices=["Plain Text", "JSON", "SRT Subtitles"],
                        value="Plain Text",
                        label="Output Format",
                        scale=1,
                    )

                with gr.Row():
                    max_tokens = gr.Slider(
                        minimum=64, maximum=512, value=256, step=64,
                        label="Max Tokens",
                    )
                    beam_width = gr.Slider(
                        minimum=1, maximum=5, value=1, step=1,
                        label="Beam Width",
                    )

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

            # Right column: Output
            with gr.Column(scale=1):
                output_text = gr.Textbox(
                    label="Transcription",
                    lines=8,
                    elem_id="output-text",
                )
                stats = gr.Textbox(
                    label="Stats",
                    interactive=False,
                    elem_classes=["stats-box"],
                )
                srt_output = gr.Textbox(
                    label="SRT Subtitles",
                    lines=8,
                    visible=False,
                )

        # Wire up output format visibility
        output_format.change(
            fn=lambda fmt: gr.update(visible=fmt == "SRT Subtitles"),
            inputs=[output_format],
            outputs=[srt_output],
        )

        # Wire up transcribe button
        transcribe_btn.click(
            fn=transcribe_audio,
            inputs=[audio_input, language, max_tokens, beam_width, output_format],
            outputs=[output_text, stats, srt_output],
        )

        # Examples
        gr.Examples(
            examples=[
                ["polywhisper_output/data/audio_fleurs_hi_in/000000.wav", "Hindi (हिन्दी)"],
                ["polywhisper_output/data/audio_fleurs_ta_in/000000.wav", "Tamil (தமிழ்)"],
            ],
            inputs=[audio_input, language],
            label="Try these examples",
        )

        # Footer
        gr.HTML("""
        <div style="text-align:center; margin-top:2em; padding:1em; color:#888; font-size:0.85em;">
            <p><b>PolyWhisper</b> — Frozen Whisper Small + Per-Language LoRA Adapters</p>
            <p>
                <a href="https://huggingface.co/eulogik/polywhisper" target="_blank">🤗 Model</a> ·
                <a href="https://github.com/eulogikdeveloper/PolyWhisper" target="_blank"> GitHub</a> ·
                <a href="https://huggingface.co/papers" target="_blank">📄 Paper</a>
            </p>
            <p>Hindi WER 37.2 | Tamil WER 60.4 on FLEURS benchmark</p>
        </div>
        """)

    return app


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--share", action="store_true", help="Create public URL via gradio.live")
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--host", default="0.0.0.0")
    args = parser.parse_args()

    app = build_app()
    app.launch(
        server_name=args.host,
        server_port=args.port,
        share=args.share,
        show_error=True,
        css=CUSTOM_CSS,
        theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="purple"),
    )