import gradio as gr
import assemblyai as aai
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

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

# =========================
# CONFIG
# =========================
aai.settings.api_key = os.getenv("ASSEMBLYAI_API_KEY")
device = "cuda" if torch.cuda.is_available() else "cpu"

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(ms):
    s = ms / 1000
    return f"{int(s // 60):02d}:{int(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


def build_segments(transcript):
    speaker_map = {}
    counter = 1
    segments = []
    for u in transcript.utterances:
        raw = str(u.speaker)
        if raw not in speaker_map:
            speaker_map[raw] = counter
            counter += 1
        segments.append({
            "speaker": speaker_map[raw],
            "start": format_time(u.start or 0),
            "end": format_time(u.end or 0),
            "text": u.text,
        })
    return segments


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

    temp_wav = None
    try:
        audio, sr = librosa.load(file, sr=None, mono=True)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
            sf.write(tmp.name, audio, sr)
            temp_wav = tmp.name

        config = aai.TranscriptionConfig(
            speaker_labels=True,
            speakers_expected=int(speakers) if speakers > 0 else None,
            language_code=None if language == "auto" else language,
        )
        transcript = aai.Transcriber().transcribe(temp_wav, config)

        if transcript.error:
            return f"❌ {transcript.error}", "", "", state

        segments = build_segments(transcript)
        speaker_count = len(set(s["speaker"] for s in segments))

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

        conversation = ""
        for i, seg in enumerate(segments, start=1):
            score = analyze_sentiment(seg["text"])
            emoji, label = label_map.get(score, ("⚪", "Unknown"))
            seg["sentiment"] = label
            conversation += (
                f"Speaker {seg['speaker']} | Utterance {i}\n"
                f"({seg['start']} - {seg['end']})\n"
                f"{emoji} {label}: {seg['text']}\n\n"
            )

        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", theme=gr.themes.Soft()) 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")
        with gr.Row():
            speakers = gr.Number(value=0, label="Speakers (0 = auto-detect)")
            language = gr.Dropdown(
                ["auto", "en", "fr", "es", "de"], value="auto", label="Language"
            )

    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, language, state],
        outputs=[status, conversation_box, info, state],
    )
    export_btn.click(
        export_file,
        inputs=[export_format, state],
        outputs=[download],
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860)