import gradio as gr
import librosa
import soundfile as sf
import torch
import json
import csv
import os
import tempfile
import warnings
from datetime import datetime
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F
from docx import Document
from reportlab.platypus import SimpleDocTemplate, Paragraph
from reportlab.lib.styles import getSampleStyleSheet
from pyannote.audio import Pipeline
import whisper

warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)
warnings.filterwarnings("ignore", category=RuntimeWarning)

# =========================
# CONFIG
# =========================
HF_TOKEN = os.getenv("HF_TOKEN")
device = "cuda" if torch.cuda.is_available() else "cpu"

# Whisper for transcription
whisper_model = whisper.load_model("base")

# Pyannote for speaker diarization
diarization_pipeline = Pipeline.from_pretrained(
    "pyannote/speaker-diarization-3.1",
    use_auth_token=HF_TOKEN
)
if device == "cuda":
    diarization_pipeline.to(torch.device("cuda"))

# BERT for sentiment
tokenizer = AutoTokenizer.from_pretrained(
    "nlptown/bert-base-multilingual-uncased-sentiment"
)
sentiment_model = AutoModelForSequenceClassification.from_pretrained(
    "nlptown/bert-base-multilingual-uncased-sentiment"
)
sentiment_model.to(device)
sentiment_model.eval()

# =========================
# HELPERS
# =========================
def format_time(seconds):
    s = int(seconds)
    return f"{s // 60:02d}:{s % 60:02d}"


def analyze_sentiment(text):
    inputs = tokenizer(text[:512], return_tensors="pt", truncation=True).to(device)
    with torch.no_grad():
        logits = sentiment_model(**inputs).logits
    probs = F.softmax(logits, dim=-1)[0]
    return torch.argmax(probs).item() + 1  # 1–5


# =========================
# MAIN PROCESS
# =========================
def process_audio(file, speakers, state):
    if file is None:
        return "❌ No audio provided", "", "", state

    temp_wav = None

    try:
        # Normalise to 16kHz mono WAV (required by Whisper and pyannote)
        audio, sr = librosa.load(file, sr=16000, mono=True)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
            sf.write(tmp.name, audio, 16000)
            temp_wav = tmp.name

        # --- Transcription (Whisper) ---
        result = whisper_model.transcribe(temp_wav)
        transcript_text = result["text"].strip()

        # --- Diarization (pyannote) ---
        num_speakers = int(speakers) if speakers > 0 else None
        diarization = diarization_pipeline(
            temp_wav,
            num_speakers=num_speakers
        )

        # Map raw pyannote speaker IDs → Speaker 1, 2, 3…
        speaker_map = {}
        speaker_counter = 1

        label_map = {
            1: ("🔴", "Very Negative"),
            2: ("🟠", "Negative"),
            3: ("🟡", "Neutral"),
            4: ("🟢", "Positive"),
            5: ("🟢", "Very Positive"),
        }

        segments = []
        conversation = ""

        for i, (turn, _, raw_speaker) in enumerate(
            diarization.itertracks(yield_label=True), start=1
        ):
            if raw_speaker not in speaker_map:
                speaker_map[raw_speaker] = speaker_counter
                speaker_counter += 1

            speaker_id = speaker_map[raw_speaker]
            start = format_time(turn.start)
            end = format_time(turn.end)

            score = analyze_sentiment(transcript_text)
            emoji, label = label_map.get(score, ("⚪", "Unknown"))

            segments.append({
                "speaker": speaker_id,
                "start": start,
                "end": end,
                "text": transcript_text,
                "sentiment": label,
            })

            conversation += (
                f"Speaker {speaker_id} | Utterance {i}\n"
                f"({start} - {end})\n"
                f"{emoji} {label}: {transcript_text}\n\n"
            )

        speaker_count = len(speaker_map)
        new_state = {"segments": segments, "conversation": conversation}

        return (
            "✅ Done",
            conversation,
            f"Speakers: {speaker_count} | Utterances: {len(segments)}",
            new_state,
        )

    except Exception as e:
        return f"❌ Error: {str(e)}", "", "", state

    finally:
        if temp_wav and os.path.exists(temp_wav):
            os.remove(temp_wav)


# =========================
# EXPORT
# =========================
def export_file(format_type, state):
    segments = state.get("segments", [])
    conversation = state.get("conversation", "")

    if not conversation and not segments:
        return None

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    if format_type == "TXT":
        path = f"/tmp/conversation_{timestamp}.txt"
        with open(path, "w", encoding="utf-8") as f:
            f.write(conversation)

    elif format_type == "JSON":
        path = f"/tmp/conversation_{timestamp}.json"
        with open(path, "w", encoding="utf-8") as f:
            json.dump(segments, f, indent=4)

    elif format_type == "CSV":
        path = f"/tmp/conversation_{timestamp}.csv"
        with open(path, "w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(
                f, fieldnames=["speaker", "start", "end", "text", "sentiment"]
            )
            writer.writeheader()
            writer.writerows(segments)

    elif format_type == "WORD":
        path = f"/tmp/conversation_{timestamp}.docx"
        doc = Document()
        doc.add_heading("Conversation Transcript", 0)
        doc.add_paragraph(conversation)
        doc.save(path)

    elif format_type == "PDF":
        path = f"/tmp/conversation_{timestamp}.pdf"
        doc = SimpleDocTemplate(path)
        styles = getSampleStyleSheet()
        content = [
            Paragraph(conversation.replace("\n", "<br/>"), styles["Normal"])
        ]
        doc.build(content)

    else:
        return None

    return path


# =========================
# UI
# =========================
with gr.Blocks(title="AI Conversation Sentiment Analyzer") as app:

    gr.Markdown("# 🎙 AI Conversation Sentiment Analyzer")

    state = gr.State({"segments": [], "conversation": ""})

    with gr.Group():
        gr.Markdown("### 🎙 Input Audio")
        audio = gr.Audio(sources=["upload", "microphone"], type="filepath")

    with gr.Group():
        gr.Markdown("### ⚙ Settings")
        speakers = gr.Number(value=0, label="Number of speakers (0 = auto-detect)")

    analyze_btn = gr.Button("🚀 Analyze", variant="primary")

    with gr.Group():
        gr.Markdown("### 💬 Conversation Output")
        status = gr.Textbox(label="Status")
        conversation_box = gr.Textbox(lines=18, label="Conversation + Sentiment")
        info = gr.Textbox(label="Info")

    with gr.Group():
        gr.Markdown("### 📁 Export")
        with gr.Row():
            export_format = gr.Dropdown(
                ["TXT", "JSON", "CSV", "WORD", "PDF"],
                value="TXT",
                label="Format"
            )
            export_btn = gr.Button("⬇ Export")
        download = gr.File()

    analyze_btn.click(
        process_audio,
        inputs=[audio, speakers, state],
        outputs=[status, conversation_box, info, state],
    )

    export_btn.click(
        export_file,
        inputs=[export_format, state],
        outputs=[download],
    )

if __name__ == "__main__":
    app.launch(theme=gr.themes.Soft())