import gradio as gr
import zipfile
import os

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

# Carregando o modelo StarCoder (especialista em programação)
tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoderbase")
model = AutoModelForCausalLM.from_pretrained("bigcode/starcoder", trust_remote_code=True)
gen = pipeline("text-generation", model=model, tokenizer=tokenizer)

def process_input(user_input, files):
    all_text = ""

    if files:
        for file in files:
            if file.name.endswith(".zip"):
                with zipfile.ZipFile(file.name, 'r') as zip_ref:
                    zip_ref.extractall("temp_zip")
                for root, dirs, filenames in os.walk("temp_zip"):
                    for fname in filenames:
                        if fname.endswith((".py", ".txt")):
                            with open(os.path.join(root, fname), 'r', encoding='utf-8', errors='ignore') as f:
                                all_text += f"\n\n# Arquivo: {fname}\n" + f.read()
            else:
                with open(file.name, 'r', encoding='utf-8', errors='ignore') as f:
                    all_text += f"\n\n# Arquivo: {file.name}\n" + f.read()

    prompt = f"""
Você é um assistente especialista em Python. Analise a seguinte mensagem do usuário e os arquivos enviados (se houver). Corrija, melhore ou explique o código conforme necessário.

Mensagem do usuário:
{user_input}

Arquivos enviados:
{all_text}

Responda com sugestões claras e organizadas.
"""

    output = gen(prompt, max_new_tokens=512, temperature=0.2)[0]["generated_text"]
    return output[len(prompt):]  # Remove o prompt da resposta final

iface = gr.Interface(
    fn=process_input,
    inputs=[
        gr.Textbox(label="Sua pergunta ou instrução", placeholder="Ex: Corrija esse código..."),
        gr.File(label="Arquivos (.py, .txt ou .zip)", file_types=[".py", ".txt", ".zip"], file_count="multiple"),
    ],
    outputs=gr.Textbox(label="Resposta da IA"),
    title="Bot IA Python",
    description="Envie sua mensagem e arquivos para corrigir, melhorar ou entender código Python."
)

iface.launch()