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

# Load pre-trained Hugging Face model for recommendation tasks
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # Specialized model for general tasks
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Increase the tokenizer's max length to handle larger inputs
tokenizer.model_max_length = 2048  # Increase this to a higher number if needed
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

def laptop_recommendation(user_input, task):
    """
    Handles laptop recommendation tasks based on user preferences.

    Parameters:
    - user_input: str, the input text from the user with laptop preferences.
    - task: str, the type of task (e.g., "Recommendation", "Compare", "Budget Recommendation").

    Returns:
    - str: The generated response.
    """
    if task == "Recommendation":
        prompt = f"Recommend a laptop based on the following preferences:\n{user_input}\nRecommended Laptop:"
    elif task == "Compare":
        prompt = f"Compare two laptops based on the following specifications:\n{user_input}\nComparison:"
    elif task == "Budget Recommendation":
        prompt = f"Recommend the best laptop for the following budget:\n{user_input}\nRecommended Laptop for Budget:"
    else:
        return "Invalid task selected."

    response = text_generator(
        prompt,
        max_length=97,
        num_return_sequences=1,
        pad_token_id=tokenizer.eos_token_id,
        temperature=0.7,
        top_p=0.9
    )[0]["generated_text"]

    return response[len(prompt):].strip()

def gradio_interface(user_input, task):
    if not user_input.strip():
        return "Please enter some input."
    return laptop_recommendation(user_input, task)

with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="green", neutral_hue="slate")) as laptop_recommendation_ui:
    gr.Markdown(
        """
        <h1 style="text-align:center; color: #3b82f6;">💻 Laptop Recommendation Chatbot</h1>
        <p style="text-align:center; color: #10b981; font-size: 18px;">
        Get personalized laptop recommendations, detailed comparisons, and budget-friendly options.
        </p>
        """,
        elem_id="header"
    )

    with gr.Row():
        with gr.Column(scale=1):
            user_input = gr.Textbox(
                lines=5,
                placeholder="Enter your laptop preferences here...",
                label="Your Input",
                elem_id="user-input",
            )
            task = gr.Radio(
                ["Recommendation", "Compare", "Budget Recommendation"],
                label="Select Task",
                elem_id="task-selector"
            )
        with gr.Column(scale=1):
            output = gr.Textbox(
                lines=10,
                label="Chatbot Response",
                elem_id="response",
            )

    with gr.Row():
        submit_button = gr.Button("Submit", elem_id="submit-btn")
        clear_button = gr.Button("Clear", elem_id="clear-btn")

    submit_button.click(gradio_interface, inputs=[user_input, task], outputs=output)
    clear_button.click(lambda: ("", ""), None, [user_input, output])

    gr.Markdown(
        """
        <style>
            #header h1 {
                font-family: 'Arial', sans-serif;
                font-weight: bold;
            }
            #user-input, #response {
                font-family: 'Courier New', monospace;
                border: 2px solid #3b82f6;
                border-radius: 10px;
            }
            #task-selector label {
                font-family: 'Verdana', sans-serif;
                color: #3b82f6;
            }
            #submit-btn {
                background-color: #3b82f6;
                color: white;
                border-radius: 10px;
            }
            #clear-btn {
                background-color: #10b981;
                color: white;
                border-radius: 10px;
            }
        </style>
        """
    )

laptop_recommendation_ui.launch()