import gradio as gr
import os
import subprocess
import tempfile
import shutil
import logging
import re
from typing import Dict, List, Optional
from dataclasses import dataclass
from transformers import pipeline

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler()],
)


class APKError(Exception):
    """Erro personalizado para problemas com APK."""


@dataclass
class PointsSystem:
    currency_name: str
    modification_points: List[str]


class APKAnalyzer:
    def __init__(self):
        self.smali_files = {}
        self.java_files = {}

    def analyze_apk(self, apk_file: bytes) -> Dict:
        """Analisa o APK usando apktool e jadx."""
        temp_dir = tempfile.mkdtemp()
        apk_path = os.path.join(temp_dir, "uploaded.apk")
        smali_dir = os.path.join(temp_dir, "smali")
        java_dir = os.path.join(temp_dir, "java")

        with open(apk_path, "wb") as f:
            f.write(apk_file)

        try:
            # Decompilar Smali com apktool
            logging.info("Descompilando APK para Smali...")
            subprocess.run(["apktool", "d", "-o", smali_dir, apk_path], check=True)
            self.smali_files = self._load_files(smali_dir, ".smali")

            # Decompilar Java com jadx
            logging.info("Descompilando APK para Java...")
            subprocess.run(["jadx", "-d", java_dir, apk_path], check=True)
            self.java_files = self._load_files(java_dir, ".java")

            # Detectar pontos no código
            points_system = self.detect_points_system()

            return {
                "smali_count": len(self.smali_files),
                "java_count": len(self.java_files),
                "currency_name": points_system.currency_name,
                "modification_points": points_system.modification_points,
            }
        except Exception as e:
            logging.error(f"Erro ao analisar o APK: {e}")
            raise APKError("Erro durante a análise do APK.")
        finally:
            shutil.rmtree(temp_dir, ignore_errors=True)

    def _load_files(self, base_dir: str, extension: str) -> Dict[str, str]:
        files = {}
        for root, _, filenames in os.walk(base_dir):
            for filename in filenames:
                if filename.endswith(extension):
                    with open(os.path.join(root, filename), "r", encoding="utf-8", errors="ignore") as f:
                        files[filename] = f.read()
        return files

    def detect_points_system(self) -> PointsSystem:
        """Detecta moedas ou pontos no APK."""
        currency_name = None
        modification_points = []

        patterns = [
            r"(coins?|gems?|points?|credits?|gold|money|cash)",
            r"(score|balance|currency|wealth)",
        ]

        for file_name, content in self.smali_files.items():
            for pattern in patterns:
                match = re.search(pattern, content, re.IGNORECASE)
                if match:
                    currency_name = match.group(0)
                    modification_points.append(file_name)

        if not currency_name:
            raise APKError("Nenhum sistema de pontuação detectado.")

        return PointsSystem(currency_name, list(set(modification_points)))


class APKModifier:
    def modify_apk(self, smali_files: Dict[str, str], points: int, multiplier: Optional[float]) -> bytes:
        """Modifica arquivos Smali e recompila o APK."""
        temp_dir = tempfile.mkdtemp()
        smali_dir = os.path.join(temp_dir, "smali")
        modified_apk_path = os.path.join(temp_dir, "modified.apk")

        os.makedirs(smali_dir, exist_ok=True)

        try:
            for file_name, content in smali_files.items():
                # Modificar conteúdo
                content = content.replace("INITIAL_POINTS", str(points))
                if multiplier:
                    content = re.sub(
                        r"ADD_POINTS (\d+)",
                        lambda m: f"ADD_POINTS {int(m.group(1)) * multiplier}",
                        content,
                    )
                with open(os.path.join(smali_dir, file_name), "w", encoding="utf-8") as f:
                    f.write(content)

            # Reempacotar com apktool
            logging.info("Recompilando APK...")
            subprocess.run(["apktool", "b", smali_dir, "-o", modified_apk_path], check=True)

            with open(modified_apk_path, "rb") as f:
                return f.read()
        finally:
            shutil.rmtree(temp_dir, ignore_errors=True)


def create_interface():
    analyzer = APKAnalyzer()
    modifier = APKModifier()
    chat_pipeline = pipeline("text-generation", model="bigscience/bloom-560m")

    with gr.Blocks(title="APK Analysis & Modification Tool") as interface:
        gr.Markdown("## APK Analysis & Modification Tool 🚀")

        with gr.Tab("Analisar APK"):
            apk_file = gr.File(label="Envie o APK", file_types=[".apk"], type="binary")
            analyze_btn = gr.Button("Analisar APK")
            analysis_output = gr.JSON(label="Resultado da Análise")

        with gr.Tab("Modificar APK"):
            points_input = gr.Slider(0, 10000, label="Pontos Iniciais")
            multiplier_input = gr.Slider(1.0, 10.0, label="Multiplicador de Pontos")
            modify_btn = gr.Button("Modificar APK")
            modified_apk_output = gr.File(label="Baixar APK Modificado")

        analyze_btn.click(
            fn=analyzer.analyze_apk, inputs=apk_file, outputs=analysis_output
        )

        modify_btn.click(
            fn=modifier.modify_apk,
            inputs=[gr.State(analyzer.smali_files), points_input, multiplier_input],
            outputs=modified_apk_output,
        )

    return interface


if __name__ == "__main__":
    interface = create_interface()
    interface.launch()