import gradio as gr
import subprocess
import os
import javalang
from transformers import pipeline
import shutil

# Função para decompilar o APK
def decompile_apk(apk_path):
    subprocess.run(["apktool", "d", apk_path, "-o", "decompiled_apk"])

# Função para analisar o código
def analyze_code():
    java_files = []
    smali_files = []
    
    for root, dirs, files in os.walk("decompiled_apk"):
        for file in files:
            if file.endswith(".java"):
                with open(os.path.join(root, file), "r") as f:
                    java_files.append(f.read())
            elif file.endswith(".smali"):
                with open(os.path.join(root, file), "r") as f:
                    smali_files.append(f.read())
    
    return java_files, smali_files

# Função para modificar o código
def modify_code(smali_code, initial_score):
    modified_smali_code = []
    
    for code in smali_code:
        modified_code = code.replace("initial_score = 0", f"initial_score = {initial_score}")
        modified_smali_code.append(modified_code)
    
    return modified_smali_code

# Função para compilar o APK
def compile_apk(modified_smali_code):
    for i, code in enumerate(modified_smali_code):
        with open(f"decompiled_apk/smali/ModifiedClass{i}.smali", "w") as f:
            f.write(code)
    
    subprocess.run(["apktool", "b", "decompiled_apk", "-o", "modified_apk.apk"])
    
    return "modified_apk.apk"

# Função para o chat com a IA
def chat_with_ai(user_message):
    chatbot = pipeline("conversational")
    response = chatbot(user_message)
    return response[0]["generated_text"]

# Função principal para processar o APK
def process_apk(apk_file, initial_score):
    apk_path = apk_file.name
    decompile_apk(apk_path)
    java_code, smali_code = analyze_code()
    modified_smali_code = modify_code(smali_code, initial_score)
    modified_apk_path = compile_apk(modified_smali_code)
    return modified_apk_path

# Interface Gradio
def main():
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column():
                apk_input = gr.File(label="Upload APK")
                initial_score = gr.Textbox(label="Initial Score")
                process_button = gr.Button("Process APK")
                process_button.click(fn=process_apk, inputs=[apk_input, initial_score], outputs="file")
            
            with gr.Column():
                chat_input = gr.Textbox(label="Chat with AI")
                chat_output = gr.Textbox(label="AI Response")
                chat_button = gr.Button("Send")
                chat_button.click(fn=chat_with_ai, inputs=chat_input, outputs=chat_output)
    
    demo.launch()

if __name__ == "__main__":
    main()