import gradio as gr
from transformers import pipeline
import torch
import spaces
import os
import shutil
from datetime import datetime

# Import dataset manager
from dataset_manager import upload_audio, update_feedback

# ---------------------------------------------------------------------------
# Setup
# ---------------------------------------------------------------------------
SAVE_DIR = "saved_audio"
os.makedirs(SAVE_DIR, exist_ok=True)

MODELS = {
    "Tonga": "buumba641/tonga-whisper-medium-shona-proxy",
    "Bemba": "buumba641/bemba-whisper-medium-siwahili-final",
    "Nyanja": "buumba641/Nyanja-whisper-medium-shona-proxy",
}

_pipeline_cache = {}

def get_pipeline(language: str):
    if language not in _pipeline_cache:
        repo_id = MODELS[language]
        device = 0 if torch.cuda.is_available() else -1
        _pipeline_cache[language] = pipeline(
            "automatic-speech-recognition",
            model=repo_id,
            device=device,
        )
    return _pipeline_cache[language]

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
STAR_FULL = "★"
STAR_EMPTY = "☆"
NUM_STARS = 5

def star_updates(n: int):
    """Return a gr.update() for each of the 5 star buttons given a rating n."""
    return [gr.update(value=(STAR_FULL if i <= n else STAR_EMPTY)) for i in range(1, NUM_STARS + 1)]

def set_star_rating(n: int):
    """Update stars and reveal the correction box if rating < 5."""
    updates = star_updates(n)
    try:
        r = int(n)
    except Exception:
        r = 0
        
    correction_update = gr.update(visible=True, value="") if 1 <= r < NUM_STARS else gr.update(visible=False, value="")
    return (*updates, n, correction_update)

def submit_feedback_to_dataset(record_id, correction, rating):
    """Pushes the updated feedback to the Hugging Face dataset."""
    if not record_id:
        return gr.update(value="⚠️ No record ID found. Did the upload fail?", visible=True)
    if not rating or rating == 0:
        return gr.update(value="⚠️ Please select a star rating.", visible=True)
        
    try:
        update_feedback(record_id, correction, rating)
        return gr.update(value="✅ Feedback submitted! Thank you.", visible=True)
    except Exception as e:
        return gr.update(value=f"⚠️ Error saving feedback: {e}", visible=True)

# ---------------------------------------------------------------------------
# Core Transcription Logic
# ---------------------------------------------------------------------------
@spaces.GPU
def transcribe(audio_path: str, language: str, progress=gr.Progress()):
    def reset_state(msg):
        return (
            msg, gr.update(visible=True), gr.update(visible=True), *star_updates(0), 0,
            gr.update(visible=False, value=""), gr.update(value=""), gr.update(interactive=True, value="✨ Transcribe"),
            "", gr.update(value="", visible=False)
        )

    if not audio_path:
        return reset_state("Please record or upload audio first.")
    if not language:
        return reset_state("Please select a language.")

    progress(0.1, desc="Preparing audio...")
    
    # Preserve original file extension (.wav, .mp3, .webm, etc.)
    ext = os.path.splitext(audio_path)[1] or ".wav"
    save_path = ""
    record_id = ""

    try:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        save_path = os.path.join(SAVE_DIR, f"{timestamp}_{language}{ext}")
        shutil.copy(audio_path, save_path)
    except Exception as e:
        print(f"Save warning: {e}")

    try:
        progress(0.35, desc=f"Loading {language} model...")
        asr = get_pipeline(language)

        progress(0.60, desc="Transcribing...")
        result = asr(audio_path, batch_size=4, return_timestamps=False)
        text = result.get("text", "").strip() or "No transcription produced."

        progress(0.85, desc="Uploading to dataset...")
        try:
            record_id = upload_audio(save_path, language, text)
            print(f"Successfully committed record: {record_id}")
        except Exception as dataset_err:
            print(f"Dataset upload failed (ensure HF_TOKEN is set): {dataset_err}")
            record_id = ""

        progress(1.0, desc="Done.")
        return (
            text,
            gr.update(visible=True),   # transcript_section
            gr.update(visible=True),   # feedback_section
            *star_updates(0),          # reset 5 stars
            0,                         # rating_value
            gr.update(visible=False, value=""), # correction_box
            gr.update(value=""),       # copy_status
            gr.update(interactive=True, value="✨ Transcribe"), # submit_btn
            record_id,                 # HIDDEN record_id
            gr.update(value="", visible=False) # feedback_status
        )

    except Exception as e:
        return reset_state(f"Error during transcription: {e}")

def start_transcribe_ui():
    return gr.update(interactive=False, value="⏳ Transcribing...")

def copy_js_status(msg):
    return gr.update(value=msg, visible=True)

# ---------------------------------------------------------------------------
# UI Definition
# ---------------------------------------------------------------------------
custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="blue",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
    radius_size="md",
)

custom_css = """
html { scroll-behavior: smooth; }
body { background: #ffffff; }

.gradio-container{ max-width: 1240px !important; width: 100% !important; margin: 0 auto !important; padding: 28px 32px 48px !important; }
.hero{ padding: 4px 0 24px; border-bottom: 1px solid #e2e8f0; margin-bottom: 24px; }
.hero h1{ margin: 0; font-size: 2rem; font-weight: 700; color: #0f172a; letter-spacing: -0.02em; }
.hero p{ margin-top: 6px; color: #64748b; font-size: 1rem; }
.card{ background: transparent; border: none; box-shadow: none; padding: 0; margin-bottom: 28px; }
button.primary{ height: 52px; font-weight: 600; border-radius: 10px; font-size: 1rem; background: #2563eb !important; border: none !important; }
button.primary:hover{ background: #1d4ed8 !important; }
#transcript_box textarea{ font-size: 15px !important; line-height: 1.6 !important; }
.feedback-title{ font-weight: 600; margin-bottom: 8px; color: #0f172a; }
.star-row{ display: flex; justify-content: flex-start; gap: 6px; margin: 4px 0 14px; }
.star-btn{ background: transparent !important; border: none !important; box-shadow: none !important; font-size: 34px !important; line-height: 1 !important; padding: 2px 4px !important; min-width: 0 !important; color: #f5b301 !important; transition: transform 0.1s ease; }
.star-btn:hover{ transform: scale(1.15); }

#audio_input { border-radius: 12px !important; }
#audio_input button svg, #audio_input .icon svg, #audio_input [class*="icon"] svg { width: 2.25rem !important; height: 2.25rem !important; }
#audio_input button { min-height: 56px !important; }
#audio_input .wrap { font-size: 1.05rem !important; }
#upload_audio .empty, #mic_audio .empty { min-height: 64px !important; padding: 8px !important; }
#upload_audio button, #mic_audio button { min-height: 44px !important; font-size: 0.95rem !important; }

@media (max-width: 768px){
    .gradio-container{ padding: 20px 16px 32px !important; }
    .hero h1 { font-size: 1.6rem; }
}
"""

# FIXED: Removed theme & css from gr.Blocks() constructor for Gradio 6
with gr.Blocks(title="ZamVoice — Zambian Language ASR Demo") as demo:
    gr.HTML("""
    <div class="hero">
        <h1>🎙️ ZamVoice</h1>
        <p>Record or upload audio in Tonga, Bemba, or Nyanja.</p>
    </div>
    """)

    with gr.Column(elem_classes="card"):
        language = gr.Dropdown(choices=list(MODELS.keys()), value="Bemba", label="🌍 Language")
        
        # FIXED: Removed waveform_options invalid kwarg
        audio = gr.Audio(
            sources=["upload", "microphone"],
            type="filepath",
            format="wav",
            editable=False,
            label="🎤 Upload or Record Audio",
            elem_id="audio_input"
        )
        submit_btn = gr.Button("✨ Transcribe", variant="primary")

    record_id_box = gr.Textbox(visible=False)

    with gr.Column(elem_classes="card", visible=False) as transcript_section:
        gr.HTML('<div id="transcript-anchor"></div>')
        with gr.Row():
            gr.Markdown("### 📝 Transcript")
            copy_btn = gr.Button("📋 Copy", variant="secondary")

        transcript_box = gr.Textbox(label="", lines=8, interactive=False, elem_id="transcript_box", visible=True)
        copy_status = gr.Markdown(visible=False)
        js_status_sink = gr.Textbox(visible=False)

    with gr.Column(elem_classes="card", visible=False) as feedback_section:
        gr.HTML('<div class="feedback-title">Was the transcript accurate? Rate it:</div>')
        rating_value = gr.Number(value=0, visible=False)

        with gr.Row(elem_classes="star-row"):
            star_buttons = [gr.Button(STAR_EMPTY, elem_classes="star-btn") for _ in range(NUM_STARS)]

        correction_box = gr.Textbox(
            label="Please write the correct transcript here to make the model better.",
            lines=3,
            placeholder="Type corrected transcript...",
            visible=False
        )
        
        feedback_submit_btn = gr.Button("📤 Submit Feedback", variant="secondary")
        feedback_status = gr.Markdown(visible=False)

    # ---------------------------------------------------------------------------
    # Event Routing
    # ---------------------------------------------------------------------------
    for idx, btn in enumerate(star_buttons, start=1):
        btn.click(
            fn=lambda n=idx: set_star_rating(n),
            inputs=None,
            outputs=[*star_buttons, rating_value, correction_box],
            queue=False
        )

    feedback_submit_btn.click(
        fn=submit_feedback_to_dataset,
        inputs=[record_id_box, correction_box, rating_value],
        outputs=[feedback_status],
        queue=True
    )

    submit_btn.click(
        fn=start_transcribe_ui,
        inputs=None,
        outputs=[submit_btn],
        queue=False
    ).then(
        fn=transcribe,
        inputs=[audio, language],
        outputs=[
            transcript_box, transcript_section, feedback_section,
            *star_buttons, rating_value, correction_box,
            copy_status, submit_btn, record_id_box, feedback_status
        ],
        queue=True
    ).then(
        fn=None,
        inputs=None,
        outputs=[],
        js="""
        () => {
          setTimeout(() => {
            const el = document.getElementById("transcript-anchor");
            if (el) el.scrollIntoView({ behavior: "smooth", block: "start" });
            const box = document.querySelector("#transcript_box textarea");
            if (box) box.focus({ preventScroll: true });
          }, 150);
          return [];
        }
        """,
        queue=False
    )

    copy_btn.click(
        fn=None,
        inputs=[transcript_box],
        outputs=[js_status_sink],
        js="""
        async (text) => {
          try {
            await navigator.clipboard.writeText(text || "");
            return "✅ Copied to clipboard.";
          } catch (e) {
            return "⚠️ Clipboard blocked by browser. Please copy manually.";
          }
        }
        """,
        queue=False
    ).then(
        fn=copy_js_status,
        inputs=[js_status_sink],
        outputs=[copy_status],
        queue=False
    )

if __name__ == "__main__":
    # FIXED: Passed theme and css to launch() for Gradio 6
    demo.queue(default_concurrency_limit=1, max_size=20).launch(
        theme=custom_theme,
        css=custom_css
    )