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")
if token:
    login(token)
    print("Login is successful")
else:
    print("Token not found. Please set your token in the environment variables.")

# Model and tokenizer setup
MODEL_NAME = "google/flan-t5-base"
tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, use_auth_token=token)
config = PeftConfig.from_pretrained("Komal-patra/results")
base_model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
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
background_image_path = '/mnt/data/image.png'
custom_css = f"""
.message.pending {{
    background: rgba(168, 196, 214, 0.5);
}}
/* Response message */
.message.bot.svelte-1s78gfg.message-bubble-border {{
    border-color: rgba(38, 107, 153, 0.5);
    background: rgba(255, 255, 255, 0.1);
}}
/* User message */
.message.user.svelte-1s78gfg.message-bubble-border {{
    background: rgba(157, 221, 249, 0.5);
    border-color: rgba(157, 221, 249, 0.5);
}}   
/* For both user and response message as per the document */
span.md.svelte-8tpqd2.chatbot.prose p {{
    color: #266B99;
}}
/* Chatbot container */
.gradio-container {{
    color: white; /* Light text color */
    background-image: url('{background_image_path}'); /* Add background image */
    background-size: cover; /* Cover the entire container */
    background-position: center; /* Center the image */
    background-repeat: no-repeat; /* Do not repeat the image */
    background-color: rgba(0, 0, 0, 0.5); /* Semi-transparent dark background */
}}
/* 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;
}}
/* Add icons to messages */
.message.user.svelte-1s78gfg {{
    display: flex;
    align-items: center;
}}
.message.user.svelte-1s78gfg:before {{
    content: url('file=Komal-patra/EU_AI_ACT/user icon.jpeg');
    margin-right: 8px;
}}
.message.bot.svelte-1s78gfg {{
    display: flex;
    align-items: center;
}}
.message.bot.svelte-1s78gfg:before {{
    content: url('file=Komal-patra/EU_AI_ACT/orcawise image.png');
    margin-right: 8px;
}}
/* Enable scrolling for the chatbot messages */
.chatbot .messages {{
    max-height: 500px;  /* Adjust as needed */
    overflow-y: auto;
}}
"""

# Gradio interface setup
with gr.Blocks(css=custom_css) as demo:
    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")

    # Function to handle user input
    def user(user_message, history):
        return "", history + [[user_message, None]]

    # Function to handle bot response
    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()