import os
import gradio as gr
from huggingface_hub import login
from transformers import AutoModelForSeq2SeqLM, T5Tokenizer
from peft import PeftModel, PeftConfig

# Hugging Face login
token = os.environ.get("token")
login(token)
print("Login is successful")

# Model and tokenizer setup
MODEL_NAME = "google/flan-t5-base"
tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, token=token)
config = PeftConfig.from_pretrained("Komal-patra/results")
base_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
model = PeftModel.from_pretrained(base_model, "Komal-patra/results")

# Text generation function
def generate_text(prompt, max_length=150):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        input_ids=inputs["input_ids"],
        max_length=max_length,
        num_beams=1,
        repetition_penalty=2.2
    )
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# Custom CSS for the UI
custom_css = """
.message.pending {
    background: #A8C4D6;
}
/* Response message */
.message.bot.svelte-1s78gfg.message-bubble-border {
    border-color: #266B99;
}
/* User message */
.message.user.svelte-1s78gfg.message-bubble-border {
    background: #9DDDF9;
    border-color: #9DDDF9;
}   
/* For both user and response message as per the document */
span.md.svelte-8tpqd2.chatbot.prose p {
    color: #266B99;
}
/* Chatbot container */
.gradio-container {
    background: #1c1c1c; /* Dark background */
    color: white; /* Light text color */
}
/* RED (Hex: #DB1616) for action buttons and links only */
.clear-btn {
    background: #DB1616;
    color: white;
}
/* Primary colors are set to be used for all sorts */
.submit-btn {
    background: #266B99;
    color: white;
}
"""

# Gradio interface setup
with gr.Blocks(css=custom_css) as demo:
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("<h2>My chats</h2>")
            chat_topics = gr.Markdown("<!-- Dynamic content -->")

        with gr.Column(scale=3):
            gr.Markdown("<h1>Ask a question about the EU AI Act</h1>")
            chatbot = gr.Chatbot()
            msg = gr.Textbox(placeholder="Ask your question...", show_label=False)  # Add placeholder text
            submit_button = gr.Button("Submit", elem_classes="submit-btn")
            clear = gr.Button("Clear", elem_classes="clear-btn")

            def user(user_message, history):
                return "", history + [[user_message, None]]

            def bot(history):
                if len(history) == 1:  # Check if it's the first interaction
                    bot_message = "Hi there! How can I help you today?"
                    history[-1][1] = bot_message  # Add welcome message to history
                else:
                    history[-1][1] = ""  # Clear the last bot message
                    previous_message = history[-1][0]  # Access the previous user message
                    bot_message = generate_text(previous_message)  # Generate response based on previous message
                    history[-1][1] = bot_message  # Update the last bot message
                return history

            submit_button.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
            msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
            clear.click(lambda: None, None, chatbot, queue=False)

demo.launch()