from transformers import pipeline
import streamlit as st

# Load once, cache in Streamlit
@st.cache_resource
def load_model():
    return pipeline("text2text-generation", model="google/flan-t5-base")

reviewer = load_model()

def review_code(code_text):
    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, ""