import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import spaces
import io


# Initialize the Qwen model and tokenizer
model_name = "Qwen/Qwen2.5-Coder-7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Function to generate README and documentation
@spaces.GPU
def generate_documentation(code_input):
    prompt = f"Generate README for the following code:\n\n{code_input}"
    
    messages = [
        {"role": "system", "content": "You are CodeDocify, a highly efficient and intelligent assistant designed to analyze code and generate comprehensive, clear, and concise documentation. Your purpose is to help developers by producing well-structured README files and detailed explanations of their code. You aim to simplify complex code into easily understandable documentation, ensuring that your responses are accurate, professional, and easy to follow."},
        {"role": "user", "content": prompt}
    ]
    
    # Prepare inputs for the model
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    # Generate the documentation
    generated_ids = model.generate(**model_inputs, max_new_tokens=4000)
    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
    documentation = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(documentation)
    return documentation


# Gradio interface
def process_code(code_input):
    documentation = generate_documentation(code_input)
    readme_file = io.BytesIO()
    readme_file.write(documentation.encode('utf-8'))
    readme_file.seek(0)
    
    return documentation, readme_file

# Set up the Gradio app with Bootstrap, icons, and smiley
with gr.Blocks(css=".container { font-family: 'Roboto', sans-serif; } .btn-primary { background-color: #007bff; } .icon { margin-right: 10px; }") as app:
    gr.Markdown("""
    #Code Documentation Generator
    
    Paste your code below, and the app will generate the Documentation for you.
    """)
    
    with gr.Row():
        code_input = gr.Textbox(lines=10, label="Paste your code here", placeholder="Enter your code...", show_label=False, elem_classes="form-control")
    
    with gr.Row():
        generate_button = gr.Button("Generate Documentation", elem_classes="btn btn-primary")
    
    with gr.Row():
        output_text = gr.Markdown(label="Generated Documentation")
        download_button = gr.File(label="Download README.md")

    # Bind function to button click
    generate_button.click(process_code, inputs=code_input, outputs=[output_text, download_button])

# Launch the Gradio app
app.launch()