import gradio as gr
from huggingface_hub import InferenceClient

# Initialize the Inference Client
client = InferenceClient(model="HuggingFaceH4/zephyr-7b-beta")

# Function for text summarization
def summarize(text):
    try:
        prompt = f"Summarize the following text:\n\n{text}\n\nSummary:"
        response = client.text_generation(prompt, max_new_tokens=100)
        return response
    except Exception as e:
        return f"Error in summarization: {str(e)}"

# Function for generating flashcards
def generate_flashcards(text):
    try:
        prompt = f"Generate flashcards for the following text:\n\n{text}\n\nFlashcards:"
        response = client.text_generation(prompt, max_new_tokens=200)
        return response
    except Exception as e:
        return f"Error in flashcard generation: {str(e)}"

# Function for question answering
def answer_question(text, question):
    try:
        prompt = f"Answer the following question based on the text:\n\nText: {text}\n\nQuestion: {question}\n\nAnswer:"
        response = client.text_generation(prompt, max_new_tokens=100)
        return response
    except Exception as e:
        return f"Error in question answering: {str(e)}"

# Function to handle task selection
def study_assistant(text, task, question=None):
    if task == "Summarize":
        return summarize(text)
    elif task == "Generate Flashcards":
        return generate_flashcards(text)
    elif task == "Answer Question":
        if not question:
            return "Please enter a question."
        return answer_question(text, question)
    else:
        return "Invalid task selected."

# Gradio Blocks for advanced UI
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="teal", secondary_hue="pink"),  # Use a vibrant theme
    css=".gradio-container {background: linear-gradient(135deg, #f5f7fa, #c3cfe2);} "
        ".output-text {font-family: 'Arial', sans-serif; font-size: 16px; color: #333;} "
        ".input-text {font-family: 'Arial', sans-serif; font-size: 16px; color: #555;} "
        "button {background: linear-gradient(135deg, #6a11cb, #2575fc); color: white; border: none; padding: 10px 20px; border-radius: 5px;} "
        "button:hover {background: linear-gradient(135deg, #2575fc, #6a11cb);} "
) as demo:
    # Title and description
    gr.Markdown(
        """
        # 🚀 **AI-Powered Study Assistant**
        **Summarize text, generate flashcards, or answer questions using AI!**
        """
    )

    # Inputs
    with gr.Row():
        text_input = gr.Textbox(
            lines=10,
            label="📝 **Input Text**",
            placeholder="Paste your text here...",
            elem_classes="input-text"
        )
        question_input = gr.Textbox(
            lines=2,
            label="❓ **Question (for Answer Question task)**",
            placeholder="Enter your question here...",
            elem_classes="input-text"
        )

    # Task selection
    task_radio = gr.Radio(
        choices=["Summarize", "Generate Flashcards", "Answer Question"],
        label="🎯 **Task**",
        value="Summarize"
    )

    # Output
    output_text = gr.Textbox(
        label="📄 **Output**",
        lines=10,
        elem_classes="output-text"
    )

    # Submit button
    submit_button = gr.Button("✨ **Submit**")

    # Examples
    gr.Examples(
        examples=[
            ["The French Revolution was a period of radical social and political upheaval in France that lasted from 1789 to 1799. It led to the rise of Napoleon Bonaparte and the eventual decline of the French monarchy.", "Summarize"],
            ["Photosynthesis is the process by which green plants use sunlight to synthesize foods with the help of chlorophyll. It converts carbon dioxide and water into glucose and oxygen.", "Generate Flashcards"],
            ["The Industrial Revolution began in the 18th century and marked a major turning point in history. Almost every aspect of daily life was influenced in some way.", "Answer Question", "When did the Industrial Revolution begin?"],
        ],
        inputs=[text_input, task_radio, question_input],
        outputs=output_text,
        fn=study_assistant,
        label="📚 **Examples**"
    )

    # Link button to function
    submit_button.click(
        study_assistant,
        inputs=[text_input, task_radio, question_input],
        outputs=output_text
    )

# Launch the app
demo.launch()