import gradio as gr
import torch
import librosa
import soundfile as sf
import numpy as np

from transformers import Wav2Vec2Processor, AutoModelForCTC

import zipfile
import os
import firebase_admin
from firebase_admin import credentials, firestore, storage, get_app
from datetime import datetime
import json
import tempfile
import uuid
from typing import List, Dict, Any, Optional
import pandas as pd
import re
import pympi
import spaces

# =========================
# Firebase initialization
# =========================

def init_firebase():
    # If already initialized (e.g., Gradio reload), skip.
    try:
        _ = get_app()
        return
    except ValueError:
        pass

    # Prefer deployed: env var 'firebase_creds' contains the *full* JSON, not a path
    svc_json = os.getenv("firebase_creds")
    if svc_json:
        firebase_config = json.loads(svc_json)
        cred = credentials.Certificate(firebase_config)
        # IMPORTANT: Admin SDK expects the **GCS bucket name** (appspot.com), not the web domain.
        bucket_name = os.getenv(
            "FIREBASE_STORAGE_BUCKET",
            f"{firebase_config.get('project_id')}.appspot.com"
        )
    else:
        # Local fallback (file on disk)
        cred = credentials.Certificate("serviceAccountKey.json")
        bucket_name = os.getenv("FIREBASE_STORAGE_BUCKET", "<your-project-id>.appspot.com")

    firebase_admin.initialize_app(cred, {"storageBucket": bucket_name})


# Call once
init_firebase()
db = firestore.client()
bucket = storage.bucket()

# =========================
# Models
# =========================
MODEL_NAME = "FormosanBank/formosan-asr-paiwan"
lang = "pwn"

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

processor = Wav2Vec2Processor.from_pretrained(MODEL_NAME)
model = AutoModelForCTC.from_pretrained(MODEL_NAME)
model.to(device)
model.eval()

# =========================
# Utility helpers
# =========================
def fmt_time(s: float) -> str:
    ms = int(round(s * 1000))
    hh = ms // 3_600_000
    mm = (ms % 3_600_000) // 60_000
    ss = (ms % 60_000) / 1000.0
    if hh:
        return f"{hh:02d}:{mm:02d}:{ss:06.3f}"
    return f"{mm:02d}:{ss:06.3f}"

def parse_ts(ts_str: str) -> float:
    """
    Accepts mm:ss.mmm or hh:mm:ss.mmm or seconds as string, returns seconds (float).
    """
    ts_str = str(ts_str).strip()
    if ts_str == "":
        return 0.0
    parts = ts_str.split(":")
    try:
        if len(parts) == 1:
            return float(parts[0])
        elif len(parts) == 2:
            m, s = parts
            return float(m) * 60 + float(s)
        elif len(parts) == 3:
            h, m, s = parts
            return float(h) * 3600 + float(m) * 60 + float(s)
    except ValueError:
        pass
    try:
        return float(ts_str)
    except ValueError:
        return 0.0

def slugify_filename(path: str) -> str:
    """
    Turn an arbitrary filename into a safe, lowercase slug.
    e.g. '/tmp/My File 01.wav' -> 'my_file_01'
    """
    if not path:
        return "audio"
    base = os.path.basename(path)
    stem, _ = os.path.splitext(base)
    stem = stem.strip().lower()
    stem = re.sub(r"[^a-z0-9]+", "_", stem)
    slug = stem.strip("_")
    return slug or "audio"

def make_base_name(audio_file: Optional[str], lang_code: str) -> str:
    """
    Build a descriptive, unique base name shared by .wav / .eaf / .txt.
    Example: 'pwn_story_1_20251202_191430_ab12cd34'
    """
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    stem = slugify_filename(audio_file) if audio_file else lang_code
    short_id = uuid.uuid4().hex[:8]
    return f"{lang_code}_{stem}_{ts}_{short_id}"

def make_segment_audio_file(audio_file: str, start: float, end: float) -> Optional[str]:
    """
    Create a temporary WAV file containing [start, end) from the original audio.
    Returns the filepath or None if creation fails.
    """
    if not audio_file or not os.path.exists(audio_file):
        return None
    if end is None or end <= start:
        return None
    duration = max(0.0, end - start)
    if duration <= 0:
        return None

    # Preserve original sample rate for playback
    y, sr = librosa.load(
        audio_file,
        sr=None,
        mono=True,
        offset=max(0.0, start),
        duration=duration
    )
    if y.size == 0:
        return None

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
    tmp.close()
    sf.write(tmp.name, y, sr)
    return tmp.name

@spaces.GPU(duration=180)
def asr_transcribe_array(audio_1d_16k: np.ndarray) -> str:
    if audio_1d_16k.size == 0:
        return ""
    with torch.inference_mode():
        inputs = processor(audio_1d_16k, sampling_rate=asr_sr, return_tensors="pt")
        input_values = inputs.input_values.to(device)
        logits = model(input_values).logits
        pred_ids = torch.argmax(logits, dim=-1)
        text = processor.batch_decode(pred_ids)[0]
    return text.replace("[UNK]", "").strip()

# Convert a DF (rows) or list-of-dicts into normalized segments
def ensure_segments(segments_like: Any) -> List[Dict]:
    """
    Accepts either:
      - a pandas.DataFrame from Gradio Dataframe
      - a list of row lists/tuples: [Start, End, Speaker, Text]
      - a list of dicts: {"start","end","speaker","text"}
    Returns a normalized list[dict].
    """
    out: List[Dict] = []

    if segments_like is None:
        return out

    if pd is not None and isinstance(segments_like, pd.DataFrame):
        if segments_like.empty:
            return out
        rows = segments_like.values.tolist()
        for row in rows:
            if len(row) < 4:
                row = list(row) + [""] * (4 - len(row))
            start_s = parse_ts(row[0])
            end_s = parse_ts(row[1])
            spk = "UNK" if row[2] is None else str(row[2])
            txt = "" if row[3] is None else str(row[3])
            out.append({"start": float(start_s), "end": float(end_s), "speaker": spk, "text": txt})
        return out

    if isinstance(segments_like, list):
        if len(segments_like) == 0:
            return out

        if isinstance(segments_like[0], (list, tuple)):
            for row in segments_like:
                if len(row) < 4:
                    row = list(row) + [""] * (4 - len(row))
                start_s = parse_ts(row[0])
                end_s = parse_ts(row[1])
                spk = "UNK" if row[2] is None else str(row[2])
                txt = "" if row[3] is None else str(row[3])
                out.append({"start": float(start_s), "end": float(end_s), "speaker": spk, "text": txt})
            return out

        if isinstance(segments_like[0], dict):
            for seg in segments_like:
                start_s = parse_ts(seg.get("start", 0))
                end_s = parse_ts(seg.get("end", 0))
                spk = str(seg.get("speaker", "UNK"))
                txt = str(seg.get("text", ""))
                out.append({"start": float(start_s), "end": float(end_s), "speaker": spk, "text": txt})
            return out

    return out

# =========================
# Fast streaming: VAD -> chunk -> transcribe
# =========================

def _resample_if_needed(audio: np.ndarray, sr_in: int, sr_out: int) -> np.ndarray:
    if sr_in == sr_out:
        return audio.astype(np.float32, copy=False)
    # librosa.resample expects float
    audio_f = audio.astype(np.float32, copy=False)
    return librosa.resample(audio_f, orig_sr=sr_in, target_sr=sr_out).astype(np.float32, copy=False)

def vad_chunk_stream(
    audio_path: str,
    chunk_len_sec: float = 10.0,
    overlap_sec: float = 0.2,
    block_sec: float = 0.5,
    min_speech_sec: float = 0.6,
    hangover_sec: float = 0.3,
) -> List[Dict]:
    """
    Incrementally reads audio and yields speech chunks with approximate VAD.
    Returns list of dict chunks: {start, end, speaker, audio_16k}
    """
    if not audio_path or not os.path.exists(audio_path):
        return []

    speaker = "UNK"

    with sf.SoundFile(audio_path, "r") as f:
        sr_in = int(f.samplerate)
        n_channels = int(f.channels)

        block_size = max(1, int(round(block_sec * sr_in)))
        chunk_size_16k = max(1, int(round(chunk_len_sec * asr_sr)))
        overlap_size_16k = int(round(overlap_sec * asr_sr))
        min_speech_16k = int(round(min_speech_sec * asr_sr))
        hangover_blocks = max(0, int(round(hangover_sec / block_sec)))

        noise_floor = 0.003  # adaptive baseline
        speaking = False
        hang = 0

        # buffers in input sr, then we resample when we finalize a chunk
        buf_in = np.zeros((0,), dtype=np.float32)
        buf_start_sample_in: Optional[int] = None

        # we track position in input samples
        pos_in = 0

        emitted = []

        while True:
            block = f.read(block_size, dtype="float32", always_2d=True)
            if block.size == 0:
                break

            if n_channels > 1:
                block_mono = block.mean(axis=1)
            else:
                block_mono = block[:, 0]

            # Simple energy VAD on this block
            rms = float(np.sqrt(np.mean(block_mono ** 2) + 1e-12))
            thr = max(0.01, noise_floor * 3.0)

            is_speech = rms > thr

            if is_speech:
                speaking = True
                hang = hangover_blocks
                if buf_start_sample_in is None:
                    buf_start_sample_in = pos_in
                buf_in = np.concatenate([buf_in, block_mono], axis=0)
            else:
                # update noise floor slowly when not speech
                noise_floor = 0.98 * noise_floor + 0.02 * rms

                if speaking:
                    if hang > 0:
                        hang -= 1
                        # still keep a bit of trailing context for natural cuts
                        buf_in = np.concatenate([buf_in, block_mono], axis=0)
                    else:
                        # speech ended -> finalize buffered speech into chunks
                        if buf_start_sample_in is not None:
                            speech_audio_16k = _resample_if_needed(buf_in, sr_in, asr_sr)
                            # chop into chunk_len with overlap
                            i = 0
                            # Convert start in seconds based on input sr, but adjust as we chunk in 16k domain
                            speech_start_sec = buf_start_sample_in / sr_in
                            while i + min_speech_16k <= len(speech_audio_16k):
                                j = min(len(speech_audio_16k), i + chunk_size_16k)
                                chunk = speech_audio_16k[i:j]
                                if len(chunk) >= min_speech_16k:
                                    start_sec = speech_start_sec + (i / asr_sr)
                                    end_sec = speech_start_sec + (j / asr_sr)
                                    emitted.append({
                                        "start": float(start_sec),
                                        "end": float(end_sec),
                                        "speaker": speaker,
                                        "audio_16k": chunk
                                    })
                                if j >= len(speech_audio_16k):
                                    break
                                i = max(0, j - overlap_size_16k)

                        # reset
                        speaking = False
                        buf_in = np.zeros((0,), dtype=np.float32)
                        buf_start_sample_in = None

            pos_in += len(block_mono)

        # EOF flush
        if buf_start_sample_in is not None and buf_in.size > 0:
            speech_audio_16k = _resample_if_needed(buf_in, sr_in, asr_sr)
            min_speech_16k = int(round(min_speech_sec * asr_sr))
            chunk_size_16k = int(round(chunk_len_sec * asr_sr))
            overlap_size_16k = int(round(overlap_sec * asr_sr))

            i = 0
            speech_start_sec = buf_start_sample_in / sr_in
            while i + min_speech_16k <= len(speech_audio_16k):
                j = min(len(speech_audio_16k), i + chunk_size_16k)
                chunk = speech_audio_16k[i:j]
                if len(chunk) >= min_speech_16k:
                    start_sec = speech_start_sec + (i / asr_sr)
                    end_sec = speech_start_sec + (j / asr_sr)
                    emitted.append({
                        "start": float(start_sec),
                        "end": float(end_sec),
                        "speaker": speaker,
                        "audio_16k": chunk
                    })
                if j >= len(speech_audio_16k):
                    break
                i = max(0, j - overlap_size_16k)

        return emitted

# =========================
# Transcription (streaming)
# =========================
def transcribe_streaming(audio_file, chunk_len=10.0, overlap=0.2, max_chunks=None):
    """
    Fast streaming: VAD -> chunk -> transcribe.
    Speaker labels are always UNK.
    """
    if audio_file is None:
        yield "", [], []
        return

    chunks = vad_chunk_stream(
        audio_path=audio_file,
        chunk_len_sec=float(chunk_len),
        overlap_sec=float(overlap),
    )
    if max_chunks is not None:
        chunks = chunks[:max_chunks]

    rows, running_text, realized = [], "", []
    for ch in chunks:
        y16k = ch.get("audio_16k", None)
        if y16k is None or len(y16k) == 0:
            continue

        text = asr_transcribe_array(y16k).strip()
        ch_out = {"start": ch["start"], "end": ch["end"], "speaker": "UNK", "text": text}
        realized.append(ch_out)

        running_text += f"[{fmt_time(ch_out['start'])}-{fmt_time(ch_out['end'])}] {ch_out['speaker']}: {ch_out['text']}\n"
        rows.append([fmt_time(ch_out["start"]), fmt_time(ch_out["end"]), ch_out["speaker"], ch_out["text"]])

        yield running_text, rows, realized

def transcribe_both(audio_file):
    transcription = ""
    if audio_file:
        y, _ = librosa.load(audio_file, sr=asr_sr, mono=True)
        transcription = asr_transcribe_array(y)
    return transcription, transcription

# =========================
# ELAN (.eaf) writer per ELAN best practices
# =========================
def write_eaf(
    segments: List[Dict],
    path: str,
    audio_path: Optional[str] = None,
    lang_code: str = "pwn",
    trim_overlaps: bool = True
):
    """
    Write an ELAN .eaf from the edited segments only.
    - One time-alignable tier per speaker with the *corrected* text.
    - Optionally link the primary media.
    - No dependent/REF tiers, no full-span fallbacks.
    """
    eaf = pympi.Elan.Eaf()

    def ms(x):
        try:
            return int(round(float(x) * 1000))
        except Exception:
            return 0

    if audio_path:
        try:
            if os.path.exists(audio_path):
                eaf.add_linked_file(audio_path)
            else:
                eaf.add_linked_file(audio_path, mimetype="audio/wav")
        except Exception:
            pass

    if "ASR" not in eaf.linguistic_types:
        eaf.add_linguistic_type("ASR", timealignable=True)

    by_spk: Dict[str, List[Dict]] = {}
    for seg in segments or []:
        by_spk.setdefault(seg.get("speaker", "UNK"), []).append(seg)

    if trim_overlaps:
        for spk, lst in by_spk.items():
            lst.sort(key=lambda d: (d["start"], d["end"]))
            last_end = -float("inf")
            for s in lst:
                if s["start"] < last_end:
                    s["start"] = last_end
                if s["end"] < s["start"]:
                    s["end"] = s["start"]
                last_end = s["end"]

    speakers = sorted(by_spk.keys())
    for spk in speakers:
        if spk not in eaf.tiers:
            eaf.add_tier(spk, ling="ASR", part=spk, language=lang_code)
        for seg in by_spk[spk]:
            eaf.add_annotation(spk, ms(seg["start"]), ms(seg["end"]), seg.get("text", ""))

    eaf.to_file(path)

# =========================
# Download helpers
# =========================
def segments_to_txt(segments: List[Dict]) -> str:
    lines = []
    for seg in sorted(segments, key=lambda d: (d["start"], d["end"])):
        lines.append(f"[{fmt_time(seg['start'])}-{fmt_time(seg['end'])}] {seg['speaker']}: {seg.get('text','')}")
    return "\n".join(lines) + ("\n" if lines else "")

def prepare_download_from_df(audio_file, df_rows, eaf_trim, eaf_link_audio):
    segs = ensure_segments(df_rows)
    return prepare_download_core(audio_file, segs, eaf_trim, eaf_link_audio)

def prepare_download_core(audio_file, segments, eaf_trim, eaf_link_audio):
    segs = ensure_segments(segments)
    if audio_file is None and not segs:
        return None

    base_name = make_base_name(audio_file, lang)

    tmp_eaf = tempfile.NamedTemporaryFile(delete=False, suffix=".eaf")
    tmp_eaf.close()

    media_href = f"{base_name}.wav" if (eaf_link_audio and audio_file) else None

    write_eaf(
        segments=segs,
        path=tmp_eaf.name,
        audio_path=media_href,
        lang_code=lang,
        trim_overlaps=bool(eaf_trim),
    )

    tmp_txt = tempfile.NamedTemporaryFile(delete=False, suffix=".txt", mode="w", encoding="utf-8")
    txt_path = tmp_txt.name
    tmp_txt.write(segments_to_txt(segs))
    tmp_txt.close()

    tmp_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
    tmp_zip.close()
    with zipfile.ZipFile(tmp_zip.name, "w") as zf:
        if audio_file and os.path.exists(audio_file):
            zf.write(audio_file, arcname=f"{base_name}.wav")
        zf.write(tmp_eaf.name, arcname=f"{base_name}.eaf")
        zf.write(txt_path, arcname=f"{base_name}.txt")

    try:
        os.unlink(tmp_eaf.name)
    except Exception:
        pass
    try:
        os.unlink(txt_path)
    except Exception:
        pass

    return tmp_zip.name

# =========================
# Save correction (segments from DF at click-time)
# =========================
def store_correction_from_df(use_zh, original_transcription, audio_file, age, native_speaker, df_rows):
    return store_correction_core(bool(use_zh), original_transcription, audio_file, age, native_speaker, ensure_segments(df_rows))

def store_correction_core(use_zh: bool, original_transcription, audio_file, age, native_speaker, segments):
    try:
        unique_id = str(uuid.uuid4())

        base_name = make_base_name(audio_file, lang)
        audio_ref = None

        audio_metadata, audio_url = {}, None
        if audio_file and os.path.exists(audio_file):
            audio, sr = librosa.load(audio_file, sr=44100, mono=True)
            audio_metadata = {
                "duration": float(librosa.get_duration(y=audio, sr=sr)),
                "file_size": os.path.getsize(audio_file),
            }
            audio_ref = f"audio/{lang}/{base_name}.wav"
            blob_audio = bucket.blob(audio_ref)
            blob_audio.upload_from_filename(audio_file)
            audio_url = blob_audio.public_url

        tmp_eaf = tempfile.NamedTemporaryFile(delete=False, suffix=".eaf")
        tmp_eaf.close()

        media_href = f"{base_name}.wav" if audio_file else None

        write_eaf(
            segments=segments or [],
            path=tmp_eaf.name,
            audio_path=media_href,
            lang_code=lang,
            trim_overlaps=True,
        )

        eaf_ref = f"elan/{lang}/{base_name}.eaf"
        blob_eaf = bucket.blob(eaf_ref)
        blob_eaf.upload_from_filename(tmp_eaf.name)
        eaf_url = blob_eaf.public_url
        try:
            os.unlink(tmp_eaf.name)
        except Exception:
            pass

        db.collection("paiwan_transcriptions").document(unique_id).set({
            "id": unique_id,
            "language": lang,
            "files": {
                "base_name": base_name,
                "audio_path": audio_ref,
                "audio_url": audio_url,
                "eaf_path": eaf_ref,
                "eaf_url": eaf_url,
            },
            "audio_metadata": audio_metadata,
            "user_info": {
                "native_paiwan_speaker": native_speaker,
                "age": age,
            },
            "timestamp": datetime.now().isoformat(),
            "model_name": MODEL_NAME,
            "schema": "eaf+audio_only_v3_vad_segments_source",
        })

        if use_zh:
            return "已儲存：EAF 與音訊已上傳至 Firebase。"
        else:
            return "Saved: uploaded EAF + audio to Firebase."
    except Exception as e:
        if use_zh:
            return f"儲存失敗：{e}"
        else:
            return f"Save failed: {e}"

# =========================
# Segment playback helper
# =========================
def play_selected_segment(audio_file, df_rows, selected_option):
    if audio_file is None or df_rows is None or not selected_option:
        return None

    segs = ensure_segments(df_rows)
    if not segs:
        return None

    try:
        idx_str = str(selected_option).split(":", 1)[0]
        idx = int(idx_str)
    except Exception:
        return None

    if idx < 0 or idx >= len(segs):
        return None

    seg = segs[idx]
    return make_segment_audio_file(audio_file, seg["start"], seg["end"])

# =========================
# Gradio UI
# =========================
with gr.Blocks() as demo:
    EN = {
        "toggle_off": "切換到繁體中文 (Switch to Traditional Chinese)",
        "toggle_on": "切換到英文 (Switch to English)",
        "title": "Paiwan ASR Transcription & Correction System",
        "step1": "Step 1: Audio Upload & Transcription",
        "step2": "Step 2: Review & Edit Transcription",
        "step3": "Step 3: User Information",
        "step4": "Step 4: Save & Download",
        "audio_label": "Audio Input",
        "stream_btn": "Stream (VAD → Chunk → Transcribe) — Speaker=UNK",
        "chunk_len": "Chunk length (sec)",
        "overlap": "Chunk overlap (sec)",
        "orig_text": "Original Transcription (streaming)",
        "segments": "Segments",
        "headers": ["Start", "End", "Speaker", "Text"],
        "age": "Age",
        "native": "Native Paiwan Speaker?",
        "eaf_trim": "Trim overlaps per speaker for EAF",
        "eaf_link": "Embed media link in EAF",
        "save_btn": "Save Correction",
        "save_status": "Save Status",
        "dl_btn": "Download ZIP File (audio + .txt + .eaf)",
        "dl_label": "Download File",
        "segment_player": "Segment Player",
        "play_segment_btn": "Play selected segment",
        "segment_selector": "Select segment to play",
    }
    ZH = {
        "toggle_off": "切換到繁體中文 (Switch to Traditional Chinese)",
        "toggle_on": "切換到英文 (Switch to English)",
        "title": "排灣語轉錄與修正系統",
        "step1": "步驟一：上傳音訊並進行轉錄",
        "step2": "步驟二：審閱與編輯段落轉錄",
        "step3": "步驟三：使用者資訊",
        "step4": "步驟四：儲存與下載",
        "audio_label": "音訊輸入",
        "stream_btn": "開始串流（VAD→分段→轉錄）— 說話者=UNK",
        "chunk_len": "分段長度（秒）",
        "overlap": "分段重疊（秒）",
        "orig_text": "原始轉錄（串流中）",
        "segments": "分段列表",
        "headers": ["開始", "結束", "說話者", "內容"],
        "age": "年齡",
        "native": "是否為排灣語母語者？",
        "eaf_trim": "為每位說話者裁切重疊區間（輸出 EAF）",
        "eaf_link": "於 EAF 中連結音訊媒體",
        "save_btn": "儲存修正",
        "save_status": "儲存狀態",
        "dl_btn": "下載 ZIP（音訊 + .txt + .eaf）",
        "dl_label": "下載檔案",
        "segment_player": "片段播放器",
        "play_segment_btn": "播放選取的片段",
        "segment_selector": "選擇要播放的片段",
    }

    use_zh_state = gr.State(False)
    lang_switch = gr.Checkbox(label=EN["toggle_off"], value=False)

    title = gr.Markdown(EN["title"])
    step1 = gr.Markdown(EN["step1"])
    with gr.Row():
        audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath", label=EN["audio_label"])

    step2 = gr.Markdown(EN["step2"])
    with gr.Row():
        transcribe_button = gr.Button(EN["stream_btn"])
        chunk_len = gr.Slider(4, 20, value=10, step=1, label=EN["chunk_len"])
        overlap = gr.Slider(0.0, 1.0, value=0.2, step=0.1, label=EN["overlap"])

    original_text = gr.Textbox(label=EN["orig_text"], interactive=False, lines=10)

    segments_table = gr.Dataframe(
        headers=EN["headers"],
        datatype=["str", "str", "str", "str"],
        label=EN["segments"],
        wrap=True,
        interactive=True,
    )
    segments_state = gr.State([])

    with gr.Row():
        segment_selector = gr.Dropdown(
            label=EN["segment_selector"],
            choices=[],
            interactive=True
        )
        play_segment_button = gr.Button(EN["play_segment_btn"])

    segment_player = gr.Audio(label=EN["segment_player"], type="filepath", interactive=False)

    def df_to_segments_and_dropdown(df_rows):
        segs = ensure_segments(df_rows)
        choices = []
        for i, seg in enumerate(segs):
            label = f"{i}: [{fmt_time(seg['start'])}-{fmt_time(seg['end'])}] {seg['speaker']}"
            txt = seg.get("text", "")
            if txt:
                snippet = txt[:40].replace("\n", " ")
                if len(txt) > 40:
                    snippet += "..."
                label += f" – {snippet}"
            choices.append(label)
        value = choices[0] if choices else None
        return segs, gr.update(choices=choices, value=value)

    segments_table.change(
        df_to_segments_and_dropdown,
        inputs=segments_table,
        outputs=[segments_state, segment_selector]
    )

    step3 = gr.Markdown(EN["step3"])
    with gr.Row():
        age_input = gr.Slider(minimum=0, maximum=100, step=1, label=EN["age"], value=25)
        native_speaker_input = gr.Checkbox(label=EN["native"], value=True)

    step4 = gr.Markdown(EN["step4"])
    with gr.Row():
        eaf_trim = gr.Checkbox(value=True, label=EN["eaf_trim"])
        eaf_link_audio = gr.Checkbox(value=True, label=EN["eaf_link"])

    with gr.Row():
        save_button = gr.Button(EN["save_btn"])
        save_status = gr.Textbox(label=EN["save_status"], interactive=False)

    with gr.Row():
        download_button = gr.Button(EN["dl_btn"])
        download_output = gr.File(label=EN["dl_label"])

    def toggle_language(switch):
        use_zh = bool(switch)
        D = ZH if use_zh else EN

        return (
            gr.update(value=D["title"]),
            gr.update(value=D["step1"]),
            gr.update(label=D["audio_label"]),
            gr.update(value=D["step2"]),
            gr.update(value=D["stream_btn"]),
            gr.update(label=D["chunk_len"]),
            gr.update(label=D["overlap"]),
            gr.update(label=D["orig_text"]),
            gr.update(label=D["segments"], headers=D["headers"]),
            gr.update(label=D["segment_selector"]),
            gr.update(value=D["play_segment_btn"]),
            gr.update(label=D["segment_player"]),
            gr.update(value=D["step3"]),
            gr.update(label=D["age"]),
            gr.update(label=D["native"]),
            gr.update(value=D["step4"]),
            gr.update(label=D["eaf_trim"]),
            gr.update(label=D["eaf_link"]),
            gr.update(value=D["save_btn"]),
            gr.update(label=D["save_status"]),
            gr.update(value=D["dl_btn"]),
            gr.update(label=D["dl_label"]),
            gr.update(label=(ZH["toggle_on"] if use_zh else EN["toggle_off"])),
            use_zh
        )

    lang_switch.change(
        toggle_language,
        inputs=lang_switch,
        outputs=[
            title, step1, audio_input,
            step2, transcribe_button, chunk_len, overlap,
            original_text, segments_table,
            segment_selector, play_segment_button, segment_player,
            step3, age_input, native_speaker_input,
            step4, eaf_trim, eaf_link_audio,
            save_button, save_status, download_button, download_output,
            lang_switch, use_zh_state
        ]
    )

    transcribe_button.click(
        fn=transcribe_streaming,
        inputs=[audio_input, chunk_len, overlap],
        outputs=[original_text, segments_table, segments_state],
        queue=True
    )

    play_segment_button.click(
        play_selected_segment,
        inputs=[audio_input, segments_table, segment_selector],
        outputs=segment_player
    )

    save_button.click(
        store_correction_from_df,
        inputs=[use_zh_state, original_text, audio_input, age_input, native_speaker_input, segments_table],
        outputs=save_status
    )

    download_button.click(
        prepare_download_from_df,
        inputs=[audio_input, segments_table, eaf_trim, eaf_link_audio],
        outputs=download_output
    )

if __name__ == "__main__":
    demo.launch()