import os
import subprocess
import tempfile
from androguard.misc import AnalyzeAPK
from transformers import pipeline
from sentence_transformers import SentenceTransformer, util
import gradio as gr

# Contexto global para armazenar dados do APK
apk_context = {"smali": {}, "java": {}, "info": ""}

def check_java():
    """Função para verificar se o Java está instalado e acessível."""
    try:
        result = subprocess.run(["java", "-version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        print("Java está disponível:", result.stderr.decode())  # Exibe informações sobre a versão do Java
    except FileNotFoundError:
        raise EnvironmentError("Java não está instalado ou não foi encontrado no PATH.")
    except Exception as e:
        raise EnvironmentError(f"Erro inesperado ao verificar a instalação do Java: {str(e)}")

def install_tools():
    """Instalar ferramentas como Baksmali e JADX."""
    baksmali_path = "/usr/local/bin/baksmali.jar"
    jadx_path = "/usr/local/bin/jadx/bin/jadx"

    # Verificar se o Java está disponível
    check_java()

    # Instalar o Baksmali caso não esteja presente
    if not os.path.exists(baksmali_path):
        print("Instalando o Baksmali...")
        subprocess.run(
            [
                "curl",
                "-L",
                "-o",
                baksmali_path,
                "https://bitbucket.org/JesusFreke/smali/downloads/baksmali-2.5.2.jar",
            ],
            check=True,
        )

    # Instalar o JADX caso não esteja presente
    jadx_zip_path = "/usr/local/bin/jadx.zip"
    if not os.path.exists(jadx_path):
        print("Instalando o JADX...")
        subprocess.run(
            [
                "curl",
                "-L",
                "-o",
                jadx_zip_path,
                "https://github.com/skylot/jadx/releases/download/v1.4.7/jadx-1.4.7.zip",
            ],
            check=True,
        )
        subprocess.run(["unzip", "-o", jadx_zip_path, "-d", "/usr/local/bin/jadx"], check=True)
        if os.path.exists(jadx_path):
            subprocess.run(["chmod", "+x", jadx_path], check=True)
        else:
            raise FileNotFoundError("Executável do JADX não encontrado no caminho esperado.")

install_tools()

def decompile_apk(apk_file):
    if apk_file is None:
        return "Nenhum arquivo enviado. Por favor, envie um arquivo APK."

    temp_apk_path = apk_file.name  # Usar o caminho do arquivo diretamente
    output_dir = tempfile.mkdtemp()
    try:
        # Decompilar usando o Baksmali
        smali_output = os.path.join(output_dir, "smali")
        subprocess.run(
            ["java", "-jar", "/usr/local/bin/baksmali.jar", "d", temp_apk_path, "-o", smali_output], 
            check=True
        )

        # Decompilar usando o JADX
        java_output = os.path.join(output_dir, "java")
        subprocess.run(
            ["/usr/local/bin/jadx/bin/jadx", "-d", java_output, temp_apk_path], 
            check=True
        )

        # Coletar arquivos Smali decompilados
        smali_files = {}
        for root, _, files in os.walk(smali_output):
            for file in files:
                if file.endswith(".smali"):
                    with open(os.path.join(root, file), "r") as f:
                        smali_files[file] = f.read()

        # Coletar arquivos Java decompilados
        java_files = {}
        for root, _, files in os.walk(java_output):
            for file in files:
                if file.endswith(".java"):
                    with open(os.path.join(root, file), "r") as f:
                        java_files[file] = f.read()

        # Armazenar resultados no contexto global
        apk_context["smali"] = smali_files
        apk_context["java"] = java_files

        return f"Decompilação bem-sucedida. Extraídos {len(smali_files)} arquivos Smali e {len(java_files)} arquivos Java."

    except subprocess.CalledProcessError as e:
        return f"Erro durante a decompilação: {e.stderr.decode('utf-8') if e.stderr else str(e)}"
    except Exception as e:
        return f"Erro durante a decompilação: {str(e)}"

def build_search_index():
    smali_texts = [f"{k}\n{v}" for k, v in apk_context["smali"].items()]
    java_texts = [f"{k}\n{v}" for k, v in apk_context["java"].items()]
    model = SentenceTransformer("all-MiniLM-L6-v2")
    smali_embeddings = model.encode(smali_texts, convert_to_tensor=True)
    java_embeddings = model.encode(java_texts, convert_to_tensor=True)
    return model, smali_texts, smali_embeddings, java_texts, java_embeddings

def query_apk_chat(user_message):
    if not apk_context["smali"] and not apk_context["java"]:
        return "Nenhum APK decompilado disponível. Por favor, envie e decompile um APK primeiro."

    try:
        model, smali_texts, smali_embeddings, java_texts, java_embeddings = build_search_index()
        query_embedding = model.encode(user_message, convert_to_tensor=True)
        smali_scores = util.pytorch_cos_sim(query_embedding, smali_embeddings).squeeze(0)
        java_scores = util.pytorch_cos_sim(query_embedding, java_embeddings).squeeze(0)
        smali_result = smali_texts[smali_scores.argmax().item()]
        java_result = java_texts[java_scores.argmax().item()]
        response = f"**Código Smali relevante:**\n\n{smali_result[:1000]}\n\n"
        response += f"**Código Java relevante:**\n\n{java_result[:1000]}"
        return response

    except Exception as e:
        return f"Erro durante a busca: {str(e)}"

apk_upload_interface = gr.Interface(
    fn=decompile_apk,
    inputs=gr.File(label="Enviar arquivo APK", file_types=[".apk"]),
    outputs="text",
    title="Analisador de APK",
    description="Envie um arquivo APK para decompilá-lo em código Smali e Java.",
)

chat_interface = gr.Interface(
    fn=query_apk_chat,
    inputs=gr.Textbox(lines=3, placeholder="Faça uma pergunta sobre o código do APK..."),
    outputs=gr.Textbox(lines=10, label="Resposta do AI"),
    title="Chat com APK",
    description="Faça perguntas sobre o código-fonte do APK em Smali ou Java.",
)

iface = gr.TabbedInterface([apk_upload_interface, chat_interface], ["Enviar & Analisar", "Conversar com AI"])
iface.launch()