import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

def load_model():
    model_name = "Salesforce/codet5-small"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
    return pipeline("summarization", model=model, tokenizer=tokenizer)

def explain_code(files):
    explainer = load_model()
    explanations = []
    for file in files:
        filename = file.name if hasattr(file, 'name') else 'Uploaded File'
        if isinstance(file, bytes):
            code = file.decode('utf-8')
        elif hasattr(file, 'read'):
            code = file.read().decode('utf-8')
        else:
            code = str(file)
        lines = code.splitlines()
        line_explanations = []
        for idx, line in enumerate(lines):
            if not line.strip():
                continue
            input_length = len(line.split())
            max_length = max(1, input_length // 2)
            result = explainer(line, max_length=max_length, do_sample=False)[0].get('summary_text', '')
            line_explanations.append(f"Line {idx+1}: {line}\nExplanation: {result.strip()}\n")
        explanations.append(f"## {filename}\n" + "\n".join(line_explanations))
    return "\n\n".join(explanations)

def main():
    with gr.Blocks() as demo:
        gr.Markdown("# CodeExplainerBot (Local CodeT5-small)")
        code_input = gr.File(label="Upload code files", file_count="multiple", type="binary")
        output = gr.Textbox(label="Line-by-line Explanation", lines=20)
        explain_btn = gr.Button("Explain Code")
        explain_btn.click(fn=explain_code, inputs=[code_input], outputs=output)
    demo.launch()

if __name__ == "__main__":
    main()