import gradio as gr
import shutil
import os
from zipfile import ZipFile
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Diretório temporário para os uploads
UPLOAD_DIR = "uploaded_files"

# Variável global para armazenar o código dos arquivos .java
java_files_content = {}

# Carregar o modelo e o tokenizador do CodeBERT
tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
model = AutoModelForCausalLM.from_pretrained("microsoft/codebert-base")

# Função para processar a pasta enviada
def process_uploaded_folder(zip_path):
    # Descompactar o arquivo ZIP
    folder_path = os.path.join(UPLOAD_DIR, "source_code")
    with ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(folder_path)
    
    # Ler os arquivos Java e armazenar seu conteúdo
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.endswith(".java"):
                file_path = os.path.join(root, file)
                with open(file_path, 'r', encoding='utf-8') as f:
                    java_files_content[file] = f.read()  # Armazenar conteúdo dos arquivos .java
    
    shutil.rmtree(folder_path)  # Limpar arquivos temporários

# Função para responder ao chat
def chat_with_codebert(chat_history, user_input):
    try:
        # Se o usuário pedir para buscar algo no código
        if "mostrar" in user_input or "variável" in user_input or "função" in user_input:
            # Pesquisar no código armazenado
            search_term = user_input.split("mostrar")[-1].strip()  # Exemplo de como pegar o termo a ser buscado
            response = search_in_code(search_term)
        else:
            # Caso contrário, apenas gerar uma resposta genérica com o histórico de chat
            chat_prompt = f"Histórico do chat:\n{chat_history}\nUsuário: {user_input}\nIA:"
            response = generate_codebert_response(chat_prompt)
        
        return response
    except Exception as e:
        return f"Erro ao processar a resposta da IA: {e}"

# Função para procurar por termos no código Java
def search_in_code(search_term):
    results = []
    for filename, content in java_files_content.items():
        if search_term.lower() in content.lower():
            # Adiciona o nome do arquivo e um trecho do código onde foi encontrado
            results.append(f"Arquivo: {filename}\n{content[:300]}...")  # Exibe os primeiros 300 caracteres do código
    
    if not results:
        return "Não encontrei nada relacionado ao seu pedido no código."
    
    return "\n\n".join(results)

# Função para gerar respostas do modelo CodeBERT
def generate_codebert_response(prompt):
    # Tokenize the input
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_length=512, num_return_sequences=1)

    # Decodificar a resposta gerada
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Interface de upload
def upload_and_analyze(zip_file_path):
    # Salvar o arquivo ZIP enviado
    if not os.path.exists(UPLOAD_DIR):
        os.makedirs(UPLOAD_DIR)
    
    # Processar o ZIP para extrair os arquivos Java
    try:
        process_uploaded_folder(zip_file_path)
    finally:
        if os.path.exists(zip_file_path):
            os.remove(zip_file_path)  # Remover o arquivo ZIP após processamento
    return "Código Java carregado com sucesso. Pergunte-me sobre o código!"

# Interface gráfica com Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Analisador de Código Java com CodeBERT 🚀")
    
    with gr.Tab("Upload e Análise"):
        upload_input = gr.File(label="Envie um arquivo ZIP com o código Java", type="filepath")
        analysis_output = gr.Textbox(label="Resultado da Análise", lines=5)
        analyze_button = gr.Button("Analisar")
        analyze_button.click(upload_and_analyze, inputs=upload_input, outputs=analysis_output)
    
    with gr.Tab("Chat com a IA"):
        chat_history = gr.State(value="")
        chat_input = gr.Textbox(label="Digite sua mensagem")
        chat_output = gr.Textbox(label="Resposta da IA")
        chat_button = gr.Button("Enviar")
        chat_button.click(chat_with_codebert, inputs=[chat_history, chat_input], outputs=chat_output)

# Executar a aplicação
if __name__ == "__main__":
    demo.launch()