import os
import subprocess
import zipfile
import tempfile
import shutil
import xml.etree.ElementTree as ET
import magic
import requests
from flask import Flask, request, jsonify, send_file, abort, render_template
from werkzeug.utils import secure_filename

app = Flask(__name__)

# Configuration
UPLOAD_FOLDER = 'uploads'
APKTOOL_VERSION = '2.9.1'  # Latest stable version
APKTOOL_URL = f'https://github.com/iBotPeaches/Apktool/releases/download/v{APKTOOL_VERSION}/apktool_{APKTOOL_VERSION}.jar'
APKTOOL_PATH = 'apktool.jar'
KEYSTORE_PATH = 'debug.keystore'
KEYSTORE_PASS = 'android'
KEY_ALIAS = 'androiddebugkey'

# Ensure directories exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def download_apktool():
    """Download Apktool if not present"""
    if not os.path.exists(APKTOOL_PATH):
        print(f"Downloading Apktool v{APKTOOL_VERSION}...")
        try:
            response = requests.get(APKTOOL_URL, stream=True)
            with open(APKTOOL_PATH, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            print("Apktool downloaded successfully")
        except Exception as e:
            print(f"Failed to download Apktool: {str(e)}")
            raise

def verify_apk(file_path):
    """Verify if file is a valid APK"""
    mime = magic.Magic(mime=True)
    file_type = mime.from_file(file_path)
    return file_type == 'application/vnd.android.package-archive'

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/upload', methods=['POST'])
def upload_apk():
    """Upload an APK file"""
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    
    if file and file.filename.endswith('.apk'):
        filename = secure_filename(file.filename)
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        file.save(filepath)
        
        # Verify it's actually an APK
        if not verify_apk(filepath):
            os.remove(filepath)
            return jsonify({'error': 'Invalid APK file'}), 400
        
        return jsonify({
            'success': True,
            'filename': filename,
            'filepath': filepath
        })
    else:
        return jsonify({'error': 'Invalid file type. Only APK files are allowed.'}), 400

@app.route('/api/decompile', methods=['POST'])
def decompile_apk():
    """Decompile an APK using apktool"""
    data = request.get_json()
    if not data or 'apk_path' not in data:
        return jsonify({'error': 'APK path required'}), 400
    
    apk_path = data['apk_path']
    if not os.path.exists(apk_path):
        return jsonify({'error': 'APK file not found'}), 404
    
    output_dir = apk_path.replace('.apk', '_decompiled')
    
    try:
        # Ensure apktool is available
        download_apktool()
        
        # Run apktool to decompile
        result = subprocess.run(
            ['java', '-jar', APKTOOL_PATH, 'd', apk_path, '-o', output_dir, '-f'],
            capture_output=True,
            text=True
        )
        
        if result.returncode != 0:
            return jsonify({
                'error': 'Decompilation failed',
                'details': result.stderr
            }), 500
            
        return jsonify({
            'success': True,
            'output_dir': output_dir
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/recompile', methods=['POST'])
def recompile_apk():
    """Recompile a decompiled APK"""
    data = request.get_json()
    if not data or 'decompiled_dir' not in data:
        return jsonify({'error': 'Decompiled directory required'}), 400
    
    decompiled_dir = data['decompiled_dir']
    output_apk = decompiled_dir.replace('_decompiled', '_modified.apk')
    
    try:
        result = subprocess.run(
            ['java', '-jar', APKTOOL_PATH, 'b', decompiled_dir, '-o', output_apk],
            capture_output=True,
            text=True
        )
        
        if result.returncode != 0:
            return jsonify({
                'error': 'Recompilation failed',
                'details': result.stderr
            }), 500
            
        return jsonify({
            'success': True,
            'output_apk': output_apk
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/sign', methods=['POST'])
def sign_apk():
    """Sign an APK file"""
    data = request.get_json()
    if not data or 'apk_path' not in data:
        return jsonify({'error': 'APK path required'}), 400
    
    apk_path = data['apk_path']
    if not os.path.exists(apk_path):
        return jsonify({'error': 'APK file not found'}), 404
    
    signed_apk = apk_path.replace('.apk', '_signed.apk')
    
    try:
        # Sign the APK
        result = subprocess.run(
            ['jarsigner', '-verbose', '-sigalg', 'SHA1withRSA', '-digestalg', 'SHA1',
             '-keystore', KEYSTORE_PATH, '-storepass', KEYSTORE_PASS,
             apk_path, KEY_ALIAS],
            capture_output=True,
            text=True
        )
        
        if result.returncode != 0:
            return jsonify({
                'error': 'Signing failed',
                'details': result.stderr
            }), 500
            
        # Zipalign the APK
        result = subprocess.run(
            ['zipalign', '-v', '4', apk_path, signed_apk],
            capture_output=True,
            text=True
        )
        
        if result.returncode != 0:
            return jsonify({
                'error': 'Zipalign failed',
                'details': result.stderr
            }), 500
            
        return jsonify({
            'success': True,
            'signed_apk': signed_apk
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/download', methods=['GET'])
def download_apk():
    """Download a processed APK file"""
    apk_path = request.args.get('apk_path')
    if not apk_path:
        abort(400, description="APK path parameter is required")
    
    if not os.path.exists(apk_path):
        abort(404, description="APK file not found")
    
    if not apk_path.endswith('.apk'):
        abort(400, description="Invalid file type")
    
    return send_file(apk_path, as_attachment=True)

@app.route('/api/extract_icon', methods=['POST'])
def extract_icon():
    """Extract icon from APK"""
    data = request.get_json()
    if not data or 'apk_path' not in data:
        return jsonify({'error': 'APK path required'}), 400
    
    apk_path = data['apk_path']
    if not os.path.exists(apk_path):
        return jsonify({'error': 'APK file not found'}), 404
    
    try:
        with zipfile.ZipFile(apk_path, 'r') as zip_ref:
            icon_paths = [
                'res/mipmap-hdpi-v4/ic_launcher.png',
                'res/mipmap-mdpi-v4/ic_launcher.png',
                'res/mipmap-xhdpi-v4/ic_launcher.png',
                'res/mipmap-xxhdpi-v4/ic_launcher.png',
                'res/mipmap-xxxhdpi-v4/ic_launcher.png',
                'res/drawable-hdpi/ic_launcher.png',
                'res/drawable/ic_launcher.png'
            ]
            
            icon_file = None
            for path in icon_paths:
                if path in zip_ref.namelist():
                    icon_file = path
                    break
            
            if not icon_file:
                return jsonify({'error': 'Icon not found in APK'}), 404
            
            temp_dir = tempfile.mkdtemp()
            extracted_path = os.path.join(temp_dir, os.path.basename(icon_file))
            with zip_ref.open(icon_file) as source, open(extracted_path, 'wb') as target:
                shutil.copyfileobj(source, target)
            
            return send_file(extracted_path, mimetype='image/png')
    except Exception as e:
        return jsonify({'error': str(e)}), 500
    finally:
        if 'temp_dir' in locals():
            shutil.rmtree(temp_dir, ignore_errors=True)

def create_debug_keystore():
    """Create debug keystore if it doesn't exist"""
    if not os.path.exists(KEYSTORE_PATH):
        print("Creating debug keystore...")
        try:
            subprocess.run([
                'keytool', '-genkey', '-v',
                '-keystore', KEYSTORE_PATH,
                '-alias', KEY_ALIAS,
                '-keyalg', 'RSA',
                '-keysize', '2048',
                '-validity', '10000',
                '-storepass', KEYSTORE_PASS,
                '-keypass', KEYSTORE_PASS,
                '-dname', 'CN=Android Debug,O=Android,C=US'
            ], check=True)
            print("Debug keystore created successfully")
        except subprocess.CalledProcessError as e:
            print(f"Failed to create debug keystore: {str(e)}")
            raise

if __name__ == '__main__':
    # Ensure we have all required tools
    try:
        download_apktool()
        create_debug_keystore()
        app.run(host="0.0.0.0", port=7860)
    except Exception as e:
        print(f"Failed to start application: {str(e)}")