from transformers import pipeline

# Load once
reviewer = pipeline("text2text-generation", model="google/flan-t5-base")

def review_code(code_text):
    # Ask model for both feedback + improved code
    prompt = f"""
    You are a code review assistant. Analyze the following code diff or snippet:
    {code_text}

    1. Give clear review feedback with improvements.
    2. Then provide a refactored/improved version of the code.
    Format the response as:
    FEEDBACK:
    <your feedback here>

    CODE:
    <improved code here>
    """

    output = reviewer(prompt, max_length=800, do_sample=False)
    response = output[0]['generated_text']

    # Split into feedback & code
    if "CODE:" in response:
        feedback, code = response.split("CODE:", 1)
        feedback = feedback.replace("FEEDBACK:", "").strip()
        return feedback, code.strip()
    else:
        return response, ""