import gradio as gr
from transformers import pipeline
import torch
from huggingface_hub import InferenceClient


# Initialize pipeline
sentiment_analyzer = pipeline("text-classification", model="tabularisai/multilingual-sentiment-analysis")
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
chatbot = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
speech_model = pipeline("automatic-speech-recognition", model="openai/whisper-small")


# Sentiment analysis functoin that takes text and returns sentiment and confidence
def analyze_sentiment(text):
    result = sentiment_analyzer(text)[0]
    
    return f"Sentiment: {result['label']}", f"Confidence: {result['score']:.2f}"

# Function where users can engage in a conversational interface.
def chat_response(
    message,
    max_tokens = 500,
    temperature = 0.8,
    top_p = 0.8,
):
    messages = [{"role": "system", "content": 'You are an assistant for a company called TrailTrek Gears Co specialized in Hiking'}]

    
    messages.append({"role": "user", "content": message})

    response = ""

    for message in chatbot.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response



# A summarization function
def summarize_text(text):
    # Ensure text is within model's max length
    max_length = 1024
    if len(text.split()) > max_length:
        text = " ".join(text.split()[:max_length])
    
    summary = summarizer(text, max_length=130, min_length=30, do_sample=False)
    return summary[0]['summary_text']

# A function that takes text and outputs a speech of the text
def speech_to_text(audio):
    if audio is None:
        return "Please upload an audio file or record audio to transcribe."
    
    # Transcribe the audio
    result = speech_model(audio)
    return result["text"]



with gr.Blocks(title="TrailTrek Gears AI Suite") as demo:
    gr.Markdown("# TrailTrek Gears AI Suite")
    
    with gr.Tab("Sentiment Analysis"):
        gr.Markdown("### Analyze the sentiment of your text")
        with gr.Row():
            sentiment_input = gr.Textbox(label="Enter text to analyze", lines=4)
            with gr.Column():
                sentiment_label = gr.Textbox(label="Sentiment")
                confidence_score = gr.Textbox(label="Confidence Score")
        sentiment_button = gr.Button("Analyze Sentiment")
        sentiment_button.click(
            analyze_sentiment,
            inputs=sentiment_input,
            outputs=[sentiment_label, confidence_score]
        )
    
    with gr.Tab("Chatbot"):
        gr.Markdown("# Stateless Chatbot using Zephyr-7b-beta")
    
        with gr.Row():
            chat_input = gr.Textbox(
                placeholder="Type your message here...", 
                label="Your Message"
            )
            chat_output = gr.Textbox(
                label="Assistant Response", 
                interactive=False
            )
        chat_input.submit(chat_response, inputs=chat_input, outputs=chat_output)


    
    with gr.Tab("Summarization"):
        gr.Markdown("### Get a concise summary of your text")
        with gr.Row():
            summary_input = gr.Textbox(label="Enter text to summarize", lines=8)
            summary_output = gr.Textbox(label="Summary", lines=4)
        summary_button = gr.Button("Generate Summary")
        summary_button.click(
            summarize_text,
            inputs=summary_input,
            outputs=summary_output
        )
    
    with gr.Tab("Speech-to-Text"):
        gr.Markdown("### Convert speech to text using Whisper")
        with gr.Row():
            with gr.Column():
                audio = gr.Audio(
                    "microphone",
                    type="filepath",
                    label="Record or upload audio"
                )
                transcribe_button = gr.Button("Transcribe Audio")
            text_output = gr.Textbox(
                label="Transcribed Text",
                lines=4,
                placeholder="Transcribed text will appear here..."
            )
        
        transcribe_button.click(
            speech_to_text,
            inputs=audio,
            outputs=text_output
        )

# Launch the app
if __name__ == "__main__":
    demo.launch()