import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
import gtts

# === Load Pipeline ===
sentiment_pipeline = pipeline("sentiment-analysis", verbose = 0, model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")

# === Sentiment Analysis ===
def analyze_sentiment(text):
    result = sentiment_pipeline(text)[0]
    return f"{result['label']} (Score: {result['score']:.2f})"


# Initialize the summarization pipeline, tokenizer, and model
model_name = "sshleifer/distilbart-cnn-12-6"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
summary_pipe = pipeline("summarization", model=model, tokenizer=tokenizer)


# Function to chunk text with an overlap
def chunk_text_with_overlap(tokens, max_length, overlap):
    chunks = []
    for i in range(0, len(tokens), max_length - overlap):
        chunk = tokens[i:i + max_length]
        chunks.append(chunk)
    return chunks


# === Summarization ===
def summarize_text(text):

    # Get the maximum length from the model configuration
    max_length = model.config.max_position_embeddings
    print('max_length:', max_length)

    # Define the overlap
    overlap = 50  # Adjust overlap as needed

    # Tokenize the text
    tokens = tokenizer(text, return_tensors='pt', truncation=False)['input_ids'][0]

    # Chunk the tokens with overlap
    chunks = chunk_text_with_overlap(tokens, max_length, overlap)

    # Summarize each chunk
    summaries = []
    for chunk in chunks:
        input_ids = chunk.unsqueeze(0)  # Add batch dimension
        summary_ids = model.generate(input_ids, max_length=max_length, num_beams=4, length_penalty=2.0, early_stopping=True)
        summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
        summaries.append(summary)

    # Combine the summaries into a final summary
    final_summary = ' '.join(summaries)

    return final_summary


# === Text-to-Speech ===
def text_to_speech(text):
    tts = gtts.gTTS(text)
    tts.save("output.mp3")
    return "output.mp3"


chat_pipeline = pipeline("text-generation", model="yasserrmd/Human-Like-Qwen2.5-1.5B-Instruct")
# === Chatbot ===
def chatbot(message, chat_history):
    # Generate response
    result = chat_pipeline(message, max_new_tokens=10)

    # Extract only the reply
    bot_reply = result[0]["generated_text"]

    return bot_reply

chat_pipeline1 = pipeline("text-generation", model="yasserrmd/Human-Like-Qwen2.5-1.5B-Instruct")

# === Chatbot ===
def chatbot1(message, chat_history):
    # Generate response
    result = chat_pipeline1(message, max_new_tokens=10)

    # Extract only the reply
    bot_reply = result[0]["generated_text"]

    return bot_reply


# === Build Gradio Interface ===
with gr.Blocks() as demo:
    with gr.Tabs():
        with gr.Tab("Sentiment Analysis"):
            gr.Markdown("### 🔍 Sentiment Analysis")
            sentiment_input = gr.Textbox(label="Enter text", lines=3, placeholder="Type a sentence to analyze...")
            sentiment_button = gr.Button("Analyze")
            sentiment_output = gr.Textbox(label="Sentiment")
            sentiment_button.click(analyze_sentiment, inputs=sentiment_input, outputs=sentiment_output)
    
        with gr.Tab("Summarization"):
            gr.Markdown("### ✂️ Summarization")
            summary_input = gr.Textbox(label="Enter text", lines=8, placeholder="Paste long text here...")
            summary_button = gr.Button("Summarize")
            summary_output = gr.Textbox(label="Summary")
            summary_button.click(summarize_text, inputs=summary_input, outputs=summary_output)
    
        with gr.Tab("Text-to-Speech"):
            gr.Markdown("### 🗣️ Text-to-Speech (using Bark)")
            tts_input = gr.Textbox(label="Enter text to speak", lines=3, placeholder="Try something like: 'Hello, how are you?'")
            tts_button = gr.Button("Generate Speech")
            tts_output = gr.Audio(label="Generated Speech", type="numpy")
            tts_button.click(text_to_speech, inputs=tts_input, outputs=tts_output)

    
        with gr.TabItem("Chatbot 1"):
            gr.ChatInterface(
                chatbot,
                type="messages",
                title="Chatbot 1",
                description="Start the conversation in Chatbot 1!")


        with gr.TabItem("Chatbot 2"):
            gr.ChatInterface(
                chatbot1,
                type="messages",
                title="Chatbot 2",
                description="Start the conversation in Chatbot 2!")

    
demo.launch()