import spaces
import gradio as gr
import torch
from transformers import pipeline

model_id = "meta-llama/Llama-3.2-3B-Instruct"
pipe = pipeline(
    "text-generation",
    model=model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)


messages = [
    {"role": "system", "content": """
        You are an insightful and empathetic journal assistant. Your goal is to help the user reflect on their daily experiences by offering thoughtful responses.
        Encourage the user to think deeply about their day, acknowledge their feelings, suggest possible improvements, and reinforce positive achievements.
        Your responses should feel warm, supportive, and introspective, promoting a balanced perspective and personal growth.
        Use a conversational, reflective tone, and offer subtle suggestions or questions for further thought when appropriate.
    """},
    {"role": "user", "content": "Today, I learned a valuable lesson about time management and how it can improve productivity."},  # User's journal entry
]

@spaces.GPU
def generate_text(prompt):
    response = pipe(prompt, max_new_tokens=1024, num_return_sequences=1)
    return response[0]['generated_text']

# Custom CSS for the app
css = """
body {background-color: #f0f8ff; font-family: 'Arial', sans-serif;}
.gradio-container {max-width: 900px; margin: 0 auto; border-radius: 15px; padding: 30px; background-color: white; box-shadow: 0 4px 8px rgba(0,0,0,0.1);}
textarea {border-radius: 10px; padding: 15px; font-size: 16px; border: 2px solid #ddd; width: 100%;}
button {background-color: #4caf50; color: #000000; border-radius: 10px; padding: 10px 15px; font-size: 16px; border: none; cursor: pointer; margin-top: 10px; margin-right: 5px;}
button:hover {background-color: #45a049;}
h1 {color: #333; font-size: 36px; text-align: center; margin-bottom: 20px;}
h2 {color: #555; font-size: 24px; margin-bottom: 10px;}
footer {text-align: center; padding: 20px; font-size: 14px;}
"""

# Define example entries
example_entries = [
    "Today, I finished a major project that I’ve been working on for weeks. It felt great to finally get it done, and I’m proud of the effort I put in.",
    "I faced a tough challenge today when I had to give a presentation with very little preparation. Despite feeling nervous, I managed to pull through by staying calm and focused.",
    "Tomorrow, my goal is to focus on improving my communication skills during team meetings. I want to make sure I’m clear and concise in expressing my ideas.",
    "I’m grateful for the support I received from my coworkers today. Their encouragement really helped me push through a stressful moment.",
    "I learned a new approach to managing my time more effectively today. I’m going to apply this to my daily routine to see if it helps me stay more productive."
]

# Function to return the selected prompt
def fill_prompt(selected_prompt):
    return selected_prompt

# Create the Gradio interface
with gr.Blocks(css=css) as journal_app:
    
    # Header and introduction
    gr.Markdown("# 📝 Productivity Journal")
    gr.Markdown("Welcome to your personalized productivity journal. Click an example below or write your thoughts!")

    # Example journal entries section
    gr.Markdown("### **Select an example to start writing:**")

    # Input text box for writing thoughts
    prompt_textbox = gr.Textbox(label="Write your thoughts here:", placeholder="Start writing or select an example...", elem_id="prompt_box", lines=5)
    
    # Align buttons using rows and columns for better layout
    with gr.Row():
        for i, entry in enumerate(example_entries):
            gr.Button(entry).click(fn=lambda e=entry: fill_prompt(e), outputs=prompt_textbox)

    with gr.Row():
        gr.Markdown("Your AI Journal Entry ✨")
        output = gr.Markdown()
    with gr.Row():
        # Generate button
        generate_button = gr.Button("Generate Entry ✍️")
        # Define the click event for text generation
        generate_button.click(fn=generate_text, inputs=prompt_textbox, outputs=output)

    # Footer
    gr.Markdown("🌞 Keep your spirits high and stay productive! 😊")
    gr.Markdown("Made with ❤️ using Gradio and Llama-3.2")

# Launch the app
journal_app.launch()