import gradio as gr
import whisper
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import uuid
import csv
import os

# Load models
whisper_model = whisper.load_model("base")
tokenizer = AutoTokenizer.from_pretrained("yiyanghkust/finbert-tone")
model = AutoModelForSequenceClassification.from_pretrained("yiyanghkust/finbert-tone")
labels = ["Positive", "Negative", "Neutral"]

# Analysis Function
def analyze_sentiment_auto(audio_path, client_type):
    transcription_result = whisper_model.transcribe(audio_path, task="translate")
    transcript = transcription_result["text"].strip()

    inputs = tokenizer(transcript, return_tensors="pt", truncation=True)
    with torch.no_grad():
        logits = model(**inputs).logits
        prediction = torch.argmax(logits, dim=1).item()
    sentiment = labels[prediction]

    client_id = str(uuid.uuid4())[:8]
    result = {
        "client_id": client_id,
        "client_type": client_type,
        "sentiment": sentiment,
        "transcript": transcript
    }

    # Save to CSV
    file_exists = os.path.isfile("sentiment_results.csv")
    with open("sentiment_results.csv", mode='a', newline='', encoding='utf-8') as csv_file:
        fieldnames = ["client_id", "client_type", "sentiment", "transcript"]
        writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
        if not file_exists:
            writer.writeheader()
        writer.writerow(result)

    return (
        transcript,
        f"🆔 {client_id}",
        f"📋 {client_type}",
        f"{sentiment}"
    )

# Color map for sentiment
def sentiment_color(sentiment):
    return {
        "Positive": "#22c55e",  # Green
        "Negative": "#ef4444",  # Red
        "Neutral": "#facc15"    # Yellow
    }.get(sentiment, "#e5e7eb")

# Gradio UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    <h1 style='text-align: center; color: #1d4ed8;'>🎙️ Multilingual Voice Sentiment Analyzer</h1>
    <p style='text-align: center; font-size: 16px;'>
        Upload a voice note in any language. It will be auto-translated to English and analyzed using FinBERT.
    </p>
    """)

    with gr.Row():
        audio_input = gr.Audio(sources=["upload"], type="filepath", label="🔊 Upload Voice File")
        client_type = gr.Dropdown(choices=["New", "Renewal", "Support"], value="New", label="👤 Client Type")

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

    with gr.Row():
        transcript_output = gr.Textbox(label="📝 English Transcript", lines=4, interactive=False)
        client_id_output = gr.Textbox(label="🆔 Client ID", interactive=False)
        client_type_output = gr.Textbox(label="👤 Client Type", interactive=False)
        sentiment_output = gr.Textbox(label="📊 Sentiment", interactive=False)

    def process(audio, client_type):
        transcript, client_id, ctype, sentiment = analyze_sentiment_auto(audio, client_type)
        return (
            gr.update(value=transcript),
            gr.update(value=client_id),
            gr.update(value=ctype),
            gr.update(value=sentiment, label=f"📊 Sentiment: {sentiment}", elem_id="sentiment-output")
        )

    submit_btn.click(
        fn=process,
        inputs=[audio_input, client_type],
        outputs=[transcript_output, client_id_output, client_type_output, sentiment_output]
    )

    gr.Markdown("""
    <style>
    #sentiment-output textarea {
        font-weight: bold;
        text-align: center;
        font-size: 1.2em;
        color: white;
        background-color: #1f2937;
    }
    </style>
    """)

demo.launch()