import gradio as gr
import os
import tempfile
import whisper
from pyannote.audio import Pipeline
from transformers import pipeline as hf_pipeline

# =========================
# LOAD MODELS
# =========================

HF_TOKEN = os.getenv("HF_TOKEN")

diarization_pipeline = Pipeline.from_pretrained(
    "pyannote/speaker-diarization-3.1", # Highly recommended to use the latest version
    token=HF_TOKEN
)


whisper_model = whisper.load_model("base")

sentiment_pipeline = hf_pipeline(
    "sentiment-analysis",
    model="nlptown/bert-base-multilingual-uncased-sentiment"
)

# =========================
# MAIN FUNCTION
# =========================

def analyze_audio(audio_file):

    if audio_file is None:
        return "❌ No audio uploaded"

    # Save temp file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
        temp_path = tmp.name

    os.system(f"ffmpeg -i \"{audio_file}\" -ar 16000 -ac 1 \"{temp_path}\" -y")

    # =========================
    # TRANSCRIPTION
    # =========================
    result = whisper_model.transcribe(temp_path)
    transcript_text = result["text"]

    # =========================
    # DIARIZATION
    # =========================
    diarization = diarization_pipeline(temp_path)

    # 🔥 MAP RAW SPEAKERS → Speaker 1, 2, 3...
    speaker_map = {}
    speaker_counter = 1

    output = "🎙 TRANSCRIPT + SPEAKERS + SENTIMENT\n\n"

    for turn, _, speaker in diarization.itertracks(yield_label=True):

        # Assign clean speaker labels
        if speaker not in speaker_map:
            speaker_map[speaker] = f"Speaker {speaker_counter}"
            speaker_counter += 1

        clean_speaker = speaker_map[speaker]

        # Simple text (you can upgrade alignment later)
        segment_text = transcript_text

        sentiment = sentiment_pipeline(segment_text[:512])[0]

        output += (
            f"{clean_speaker} ({turn.start:.2f}s - {turn.end:.2f}s)\n"
            f"Sentiment: {sentiment['label']} ({round(sentiment['score'],2)})\n"
            f"Text: {segment_text}\n\n"
        )

    return output


# =========================
# UI
# =========================

app = gr.Interface(
    fn=analyze_audio,
    inputs=gr.Audio(type="filepath", label="Upload Audio"),
    outputs=gr.Textbox(lines=25, label="Results"),
    title="🎙 AI Conversation Analyzer",
    description="Speaker Diarization + Sentiment Analysis"
)

app.launch()