from flask import jsonify, send_file, request
from main import *

def generate_code(prompt, output_path="output_code.py"):
    if codegen_model is None or codegen_tokenizer is None:
        return "Code generation model or tokenizer not initialized."
    input_ids = codegen_tokenizer(prompt, return_tensors='pt').to(device)
    output = codegen_model.generate(input_ids, max_length=2048, temperature=0.7, top_p=0.9)
    code = codegen_tokenizer.decode(output[0], skip_special_tokens=True)
    with open(output_path, "w") as file:
        file.write(code)
    return output_path

def codegen_api():
    data = request.get_json()
    prompt = data.get('prompt')
    if not prompt:
        return jsonify({"error": "Prompt is required"}), 400
    output_file = generate_code(prompt)
    if output_file == "Code generation model or tokenizer not initialized.":
        return jsonify({"error": "Code generation failed"}), 500
    return send_file(output_file, mimetype="text/x-python", as_attachment=True, download_name="output.py")