from __future__ import annotations

import os
os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")

import spaces
from pathlib import Path
import math

import gradio as gr
import soundfile as sf
import torch
from funasr import AutoModel as FunASRAutoModel
from transformers import AutoModelForAudioFrameClassification, AutoProcessor

from pipeline_utils import (
    SAMPLE_RATE,
    clean_segments,
    export_json,
    format_time,
    prepare_audio,
    prepare_youtube,
    timeline_html,
)

DIAR_MODEL = "nvidia/Nemotron-3-Diarization"
EMOTION_MODEL = "emotion2vec/emotion2vec_plus_large"

EMOTIONS = {
    "开心": "Happy", "happy": "Happy",
    "生气": "Angry", "angry": "Angry",
    "难过": "Sad", "sad": "Sad",
    "中立": "Neutral", "neutral": "Neutral",
    "吃惊": "Surprised", "surprised": "Surprised", "surprise": "Surprised",
    "恐惧": "Fearful", "fearful": "Fearful", "fear": "Fearful",
    "厌恶": "Disgusted", "disgusted": "Disgusted", "disgust": "Disgusted",
    "其他": "Other", "other": "Other",
    "<unk>": "Unknown", "unknown": "Unknown",
}


def english_emotion(label: str) -> str:
    parts = str(label).strip().lower().split("/")
    for part in reversed(parts):
        if part.strip() in EMOTIONS:
            return EMOTIONS[part.strip()]
    return "Unknown"


def parse_emotion(prediction, threshold: float):
    if not isinstance(prediction, list) or not prediction:
        raise ValueError("No emotion prediction returned.")

    record = prediction[0]
    labels = record.get("labels") or []
    values = record.get("scores") or []
    if len(labels) != len(values) or not labels:
        raise ValueError("Invalid emotion prediction.")

    scores = {}
    for raw_label, raw_value in zip(labels, values):
        value = float(raw_value)
        if not math.isfinite(value):
            continue
        label = english_emotion(raw_label)
        scores[label] = max(value, scores.get(label, float("-inf")))

    if not scores:
        raise ValueError("No valid emotion scores returned.")

    scores = dict(sorted(scores.items(), key=lambda pair: pair[1], reverse=True))
    top = next(iter(scores))
    score = scores[top]
    emotion = top if score >= threshold and top not in {"Other", "Unknown"} else "Uncertain"
    return emotion, round(score, 5), {name: round(value, 5) for name, value in scores.items()}


print("Loading diarization model...")
diar_processor = AutoProcessor.from_pretrained(DIAR_MODEL)
diar_model = AutoModelForAudioFrameClassification.from_pretrained(DIAR_MODEL).to("cuda").eval()

print("Loading emotion model...")
emotion_model = FunASRAutoModel(
    model=EMOTION_MODEL,
    hub="hf",
    device="cuda",
    disable_update=True,
    disable_pbar=True,
    ncpu=2,
)


def prepare(upload_path, youtube_url):
    try:
        if youtube_url and youtube_url.strip():
            wav_path, duration, _ = prepare_youtube(youtube_url)
        else:
            wav_path, duration, _ = prepare_audio(upload_path)
        return wav_path, duration
    except Exception as exc:
        raise gr.Error(str(exc)) from exc


@spaces.GPU(duration=120)
def analyze(wav_path, duration, min_seconds, min_score):
    if not wav_path or not Path(wav_path).is_file():
        raise gr.Error("Upload a file first.")

    waveform, sr = sf.read(wav_path, dtype="float32")
    if waveform.ndim != 1 or sr != SAMPLE_RATE:
        raise gr.Error("Could not process the audio.")

    model_input = diar_processor(waveform, sampling_rate=sr).to(
        diar_model.device, dtype=diar_model.dtype
    )
    with torch.inference_mode():
        logits = diar_model(**model_input).logits
    raw_turns = diar_processor.extract_speaker_dict(logits, model_input.attention_mask)[0]
    del logits, model_input

    turns, skipped = clean_segments(raw_turns, duration, float(min_seconds))
    run_dir = Path(wav_path).parent
    rows, selections, segments = [], [], []

    for turn in turns:
        left = max(0, round(turn["start"] * sr))
        right = min(len(waveform), round(turn["end"] * sr))
        if right <= left:
            continue

        index = len(segments) + 1
        speaker = f"Speaker {turn['speaker'] + 1}"
        clip_path = run_dir / f"clip_{index:03d}.wav"
        sf.write(clip_path, waveform[left:right], sr, subtype="PCM_16")

        emotion, score, scores = "Overlap", None, {}
        if not turn["too_much_overlap"]:
            with torch.inference_mode():
                prediction = emotion_model.generate(
                    input=str(clip_path),
                    granularity="utterance",
                    extract_embedding=False,
                )
            emotion, score, scores = parse_emotion(prediction, float(min_score))

        turn["emotion"] = emotion
        start, end = format_time(turn["start"]), format_time(turn["end"])
        rows.append([speaker, f"{start} – {end}", emotion, score])

        segment = {
            "speaker": speaker,
            "start": round(turn["start"], 3),
            "end": round(turn["end"], 3),
            "emotion": emotion,
            "score": score,
            "class_scores": scores,
            "overlap_ratio": round(turn["overlap_ratio"], 4),
        }
        segments.append(segment)
        selections.append({"path": str(clip_path), "title": f"{speaker} · {emotion} · {start}–{end}"})

    speakers = len({turn["speaker"] for turn in turns})
    summary = f"**{speakers} speakers** · **{len(segments)} segments** · {skipped} skipped"
    if not segments:
        summary = "No segments found. Try a shorter minimum duration."

    export_path = export_json(run_dir / "results.json", {
        "models": {"diarization": DIAR_MODEL, "emotion": EMOTION_MODEL},
        "settings": {"minimum_seconds": min_seconds, "minimum_score": min_score},
        "segments": segments,
    })

    first = selections[0] if selections else None
    return (
        summary,
        timeline_html(turns, max(duration, 0.01)),
        rows,
        first["path"] if first else None,
        first["title"] if first else "Select a segment",
        export_path,
        selections,
        gr.update(visible=True),
    )


def select_segment(selections, evt: gr.SelectData):
    index = evt.index[0] if isinstance(evt.index, (tuple, list)) else evt.index
    if not selections or not isinstance(index, int) or not 0 <= index < len(selections):
        return gr.skip(), gr.skip()
    item = selections[index]
    return item["path"], item["title"]


CSS = """
.gradio-container {max-width: 940px !important; margin: 0 auto;}
#page-title h1 {font-size: 1.65rem; letter-spacing: -.025em; margin-bottom: 0;}
.timeline-wrapper {padding: 12px 8px;}
.speaker-timeline-row {display: flex; align-items: center; gap: 12px; margin: 9px 0;}
.speaker-timeline-row > span {flex: 0 0 76px; font-size: .8rem;}
.timeline-track {flex: 1; height: 20px; background: var(--background-fill-secondary); border-radius: 4px; position: relative; overflow: hidden;}
.timeline-block {position: absolute; top: 0; height: 100%; border-radius: 3px; opacity: .8; filter: saturate(.4);}
.timeline-axis {display: flex; justify-content: space-between; margin: 6px 0 0 88px; font-size: .7rem; opacity: .65;}
.empty-timeline {font-size: .85rem; opacity: .6;}
"""

with gr.Blocks(
    title="Voice Emotions",
    delete_cache=(3600, 7200),
) as demo:
    gr.Markdown("# Voice Emotions", elem_id="page-title")
    with gr.Row():
        upload = gr.File(label="Audio / Video", type="filepath")
        youtube_url = gr.Textbox(label="YouTube", placeholder="Video URL", lines=1)

    with gr.Row():
        min_seconds = gr.Slider(0.5, 10, value=2, step=0.5, label="Min. segment (s)")
        min_score = gr.Slider(0, 1, value=0.4, step=0.05, label="Min. score")

    button = gr.Button("Analyze", variant="primary")
    wav_state = gr.State()
    duration_state = gr.State()
    segments_state = gr.State([])

    with gr.Column(visible=False) as results:
        summary = gr.Markdown()
        timeline = gr.HTML()
        with gr.Row():
            with gr.Column(scale=3):
                table = gr.Dataframe(
                    headers=["Speaker", "Time", "Emotion", "Score"],
                    datatype=["str", "str", "str", "number"],
                    type="array",
                    interactive=False,
                    label="Segments",
                )
            with gr.Column(scale=2):
                player = gr.Audio(label="Clip", type="filepath", interactive=False)
                clip_title = gr.Markdown()
                download = gr.File(label="JSON", interactive=False)

    button.click(
        prepare,
        inputs=[upload, youtube_url],
        outputs=[wav_state, duration_state],
        concurrency_limit=2,
    ).success(
        analyze,
        inputs=[wav_state, duration_state, min_seconds, min_score],
        outputs=[summary, timeline, table, player, clip_title, download, segments_state, results],
        concurrency_limit=1,
    )
    table.select(select_segment, inputs=segments_state, outputs=[player, clip_title])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        show_error=True, theme=gr.themes.Base(primary_hue="slate"), css=CSS
    )