import gradio as gr
from transformers import pipeline
import random

# Load sentiment analysis model
sentiment_pipe = pipeline("sentiment-analysis")

# Load summarization model
summarizer = pipeline("summarization")

# Load text-to-speech model
tts_pipe = pipeline("text-to-speech", model="suno/bark-small")

## real work now

# Sentiment Analysis Function
def get_sentiment(input_text):
    analysis = sentiment_pipe(input_text)[0]
    return analysis['label'], str(round(analysis['score'], 4))

# Summarization Function
def summarize_text(input_text):
    summary = summarizer(input_text, max_length=150, min_length=30, do_sample=False)[0]
    return summary['summary_text']

# Text-to-Speech Function
def text_to_speech(input_text):
    speech = tts_pipe(input_text)
    return speech["path"]

# Chatbot Logic
def chat(message, history):
    history = history or []
    if message.startswith("How many"):
        response = str(random.randint(1, 10))
    elif message.startswith("How"):
        response = random.choice(["Great", "Good", "Okay", "Bad"])
    elif message.startswith("Where"):
        response = random.choice(["Here", "There", "Somewhere"])
    else:
        response = "I don't know"
    history.append((message, response))
    return history, history

# Create Gradio Interface
with gr.Blocks(title="TrailTrek AI Suite") as demo:
    gr.Markdown("# TrailTrek Gears Co. AI Prototype")

    with gr.Tabs():
        # Sentiment Analysis Tab
        with gr.Tab("Sentiment Analysis"):
            gr.Markdown("## Analyze Text Sentiment")
            with gr.Row():
                text_input = gr.Textbox(label="Input Text")
                with gr.Column():
                    sentiment_label = gr.Textbox(label="Sentiment")
                    score_output = gr.Textbox(label="Confidence Score")
            analyze_btn = gr.Button("Analyze")
            analyze_btn.click(
                fn=get_sentiment,
                inputs=text_input,
                outputs=[sentiment_label, score_output]
            )

        # Chatbot Tab
        with gr.Tab("Chatbot"):
            gr.Markdown("## Interactive Chat")
            chatbot = gr.Chatbot()
            msg = gr.Textbox(label="Your Message")
            clear = gr.Button("Clear")
            msg.submit(chat, [msg, chatbot], [chatbot, msg])
            clear.click(lambda: None, None, chatbot, queue=False)

        # Summarization Tab
        with gr.Tab("Summarization"):
            gr.Markdown("## Text Summarization")
            with gr.Row():
                long_text = gr.Textbox(label="Input Text", lines=5)
                summary = gr.Textbox(label="Summary", lines=5)
            summarize_btn = gr.Button("Summarize")
            summarize_btn.click(
                fn=summarize_text,
                inputs=long_text,
                outputs=summary
            )

        # Text-to-Speech Tab
        with gr.Tab("Text-to-Speech"):
            gr.Markdown("## Web Accessibility Prototype")
            with gr.Row():
                tts_input = gr.Textbox(label="Enter Text")
                tts_output = gr.Audio(label="Generated Speech")
            tts_btn = gr.Button("Convert to Speech")
            tts_btn.click(
                fn=text_to_speech,
                inputs=tts_input,
                outputs=tts_output
            )


# Launch the Gradio app
demo.launch()