import gradio as gr
from transformers import pipeline
try:
    from transformers import Conversation
except ImportError:                              # older Transformers fallback
    from transformers.pipelines.conversational import Conversation

from speechbrain.pretrained import Tacotron2, HIFIGAN
import torch, numpy as np

# ───────────────────────────  Model loading  ────────────────────────────
tacotron2 = Tacotron2.from_hparams(source="speechbrain/tts-tacotron2-ljspeech")
hifigan   = HIFIGAN.from_hparams(source="speechbrain/tts-hifigan-ljspeech")

sentiment_pipe = pipeline(
    "sentiment-analysis",
    model="distilbert-base-uncased-finetuned-sst-2-english"
)
chat_pipe = pipeline(
    "conversational",
    model="microsoft/DialoGPT-medium",
    device_map="auto"
)
summ_pipe = pipeline(
    "summarization",
    model="sshleifer/distilbart-cnn-12-6",
    device_map="auto"
)

print("✅  All models are loaded!")

# ─────────────────────  Sentiment-analysis tab  ─────────────────────────
def analyze(text):
    res = sentiment_pipe(text)[0]
    return res["label"], round(res["score"], 4)

with gr.Blocks() as sentiment_tab:
    inp   = gr.Textbox(label="Enter text")
    btn   = gr.Button("Analyze")
    label = gr.Textbox(label="Sentiment")
    conf  = gr.Number(label="Confidence")
    btn.click(analyze, inp, [label, conf])

# ─────────────────────  Summarization tab  ──────────────────────────────
def summarize(text):
    return summ_pipe(text, max_length=130, min_length=30, do_sample=False)[0]["summary_text"]

with gr.Blocks() as summarization_tab:
    long  = gr.Textbox(lines=12, label="Paste long text")
    btn   = gr.Button("Summarize")
    short = gr.Textbox(lines=8, label="Summary")
    btn.click(summarize, long, short)

# ─────────────────────  Text-to-speech tab  ─────────────────────────────
def speak(text):
    text = text.strip()
    if not text:
        return None
    mel, _, _ = tacotron2.encode_text(text)
    wav = hifigan.decode_batch(mel)[0].cpu().numpy().astype(np.float32).squeeze()
    return 22050, wav                   # (sample-rate, numpy array)

with gr.Blocks() as tts_tab:
    txt = gr.Textbox(label="Text to speak")
    btn = gr.Button("Generate voice")
    aud = gr.Audio(label="Speech", type="numpy")
    btn.click(speak, txt, aud)

# ─────────────────────  Chatbot tab (live echo)  ────────────────────────
with gr.Blocks() as chatbot_tab:
    state   = gr.State([])           # list[tuple[str, str]]
    chatbox = gr.Chatbot()
    msg     = gr.Textbox(placeholder="Type here…")
    send    = gr.Button("Send")

    def add_user(user_text, history):
        user_text = user_text.strip()
        if not user_text:
            return history, history, ""
        history = history + [(user_text, "")]
        return history, history, ""          # clear textbox

    def add_bot(history):
        past_users = [u for u, _ in history[:-1]]
        past_bots  = [b for _, b in history[:-1]]
        last_user  = history[-1][0]

        conv = Conversation(
            text=last_user,
            past_user_inputs=past_users,
            generated_responses=past_bots
        )
        chat_pipe(conv)
        bot_reply = conv.generated_responses[-1]
        history[-1] = (last_user, bot_reply)
        return history, history

    send.click(
        add_user,
        [msg, state],
        [chatbox, state, msg],
        queue=False
    ).then(
        add_bot,
        state,
        [chatbox, state],
        queue=True
    )

    msg.submit(
        add_user,
        [msg, state],
        [chatbox, state, msg],
        queue=False
    ).then(
        add_bot,
        state,
        [chatbox, state],
        queue=True
    )

# ────────────────────  Assemble & launch app  ──────────────────────────
with gr.Blocks(title="TrailTrek Gears – AI Demo") as demo:
    with gr.TabItem("Sentiment Analysis"): sentiment_tab.render()
    with gr.TabItem("Summarization"):      summarization_tab.render()
    with gr.TabItem("Text-to-Speech"):     tts_tab.render()
    with gr.TabItem("Chatbot"):            chatbot_tab.render()

if __name__ == "__main__":
    demo.launch()