import gradio as gr
import os
from huggingface_hub import InferenceClient

# Get the token from the "HF_TOKEN" environment variable
token = os.getenv("HF_TOKEN")

# Create a client for the Salesforce/codet5-large model using the token
client = InferenceClient("Salesforce/codet5-large", token=token)


def generate_code(
        task_description,
        max_tokens,
        temperature,
        top_p,
):
    # 2. Create a prompt using task description
    prompt = task_description

    # 3. Generate code based on the description
    response = client.text_generation(
        prompt,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )

    # Since the response is already a string, just return it
    generated_code = response
    return generated_code


# 4. Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 CodeT5 Code Generator")

    with gr.Row():
        task_input = gr.Textbox(
            lines=3, placeholder="Describe the task in natural language...", label="Task Description"
        )

    with gr.Row():
        max_tokens = gr.Slider(1, 2048, value=100, step=1, label="Max Tokens")
        temperature = gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")

    with gr.Row():
        submit_button = gr.Button("Generate Code 🚀")

    output = gr.Textbox(lines=10, label="Generated Code")

    # 5. Button click triggers code generation
    submit_button.click(
        generate_code,
        inputs=[task_input, max_tokens, temperature, top_p],
        outputs=output,
    )

if __name__ == "__main__":
    demo.launch()