#!/usr/bin/env python3
# Copyright (c) 2018, the R8 project authors. Please see the AUTHORS file
# for details. All rights reserved. Use of this source code is governed by a
# BSD-style license that can be found in the LICENSE file.

# Script for checking impact of a change by comparing the sizes of generated
# classes in an apk.

import glob
import optparse
import os
import shutil
import sys
import threading
import time
import zipfile

import toolhelper
import utils

USAGE = """%prog [options] app1 app2
  NOTE: This only makes sense if minification is disabled"""

MAX_THREADS = 40


def parse_options():
    result = optparse.OptionParser(usage=USAGE)
    result.add_option('--no-build',
                      help='Run without building first',
                      default=False,
                      action='store_true')
    result.add_option('--temp',
                      help='Temporary directory to store extracted classes in')
    result.add_option(
        '--use_code_size',
        help=
        'Use the size of code segments instead of the full size of the dex.',
        default=False,
        action='store_true')
    result.add_option(
        '--report', help='Print comparison to this location instead of stdout')
    return result.parse_args()


def extract_apk(apk, output):
    if os.path.exists(output):
        shutil.rmtree(output)
    zipfile.ZipFile(apk).extractall(output)
    with utils.ChangedWorkingDirectory(output):
        dex = glob.glob('*.dex')
        return [os.path.join(output, dexfile) for dexfile in dex]


def ensure_exists(files):
    for f in files:
        if not os.path.exists(f):
            raise Exception('%s does not exist' % f)


def extract_classes(input, output, options):
    if os.path.exists(output):
        shutil.rmtree(output)
    os.makedirs(output)
    args = ['--file-per-class', '--output', output]
    if options.no_build:
        args.extend(['--no-build'])
    args.extend(input)
    if toolhelper.run('d8', args) is not 0:
        raise Exception('Failed running d8')


def get_code_size(path):
    segments = toolhelper.run('dexsegments', [path],
                              build=False,
                              return_stdout=True)
    for line in segments.splitlines():
        if 'Code' in line:
            # The code size line looks like:
            #  - Code: 264 / 4
            splits = line.split(' ')
            return int(splits[3])
    # Some classes has no code.
    return 0


class FileInfo:

    def __init__(self, path, root):
        self.path = path
        self.full_path = os.path.join(root, path)

    def __eq__(self, other):
        return self.full_path == other.full_path

    def set_size(self, use_code_size):
        if use_code_size:
            self.size = get_code_size(self.full_path)
        else:
            self.size = os.path.getsize(self.full_path)


def generate_file_info(path, options):
    file_info_map = {}
    with utils.ChangedWorkingDirectory(path):
        for root, dirs, files in os.walk('.'):
            for f in files:
                assert f.endswith('dex')
                file_path = os.path.join(root, f)
                entry = FileInfo(file_path, path)
                if not options.use_code_size:
                    entry.set_size(False)
                file_info_map[file_path] = entry
    threads = []
    file_infos = file_info_map.values() if options.use_code_size else []
    while len(file_infos) > 0 or len(threads) > 0:
        for t in threads:
            if not t.is_alive():
                threads.remove(t)
        # sleep
        if len(threads) == MAX_THREADS or len(file_infos) == 0:
            time.sleep(0.5)
        while len(threads) < MAX_THREADS and len(file_infos) > 0:
            info = file_infos.pop()
            print('Added %s for size calculation' % info.full_path)
            t = threading.Thread(target=info.set_size,
                                 args=(options.use_code_size,))
            threads.append(t)
            t.start()
        print('Missing %s files, threads=%s ' % (len(file_infos), len(threads)))

    return file_info_map


def print_info(app, app_files, only_in_app, bigger_in_app, output):
    output.write('Only in %s\n' % app)
    only_app_sorted = sorted(only_in_app,
                             key=lambda a: app_files[a].size,
                             reverse=True)
    output.write('\n'.join(
        ['  %s %s bytes' % (x, app_files[x].size) for x in only_app_sorted]))
    output.write('\n\n')
    output.write('Bigger in %s\n' % app)
    # Sort by the percentage diff compared to size
    percent = lambda a: (0.0 + bigger_in_app.get(a)) / app_files.get(a
                                                                    ).size * 100
    for bigger in sorted(bigger_in_app, key=percent, reverse=True):
        output.write('  {0:.3f}% {1} bytes {2}\n'.format(
            percent(bigger), bigger_in_app[bigger], bigger))
    output.write('\n\n')


def compare(app1_classes_dir, app2_classes_dir, app1, app2, options):
    app1_files = generate_file_info(app1_classes_dir, options)
    app2_files = generate_file_info(app2_classes_dir, options)
    only_in_app1 = [k for k in app1_files if k not in app2_files]
    only_in_app2 = [k for k in app2_files if k not in app1_files]
    in_both = [k for k in app2_files if k in app1_files]
    assert len(app1_files) == len(only_in_app1) + len(in_both)
    assert len(app2_files) == len(only_in_app2) + len(in_both)
    bigger_in_app1 = {}
    bigger_in_app2 = {}
    same_size = []
    for f in in_both:
        app1_entry = app1_files[f]
        app2_entry = app2_files[f]
        if app1_entry.size > app2_entry.size:
            bigger_in_app1[f] = app1_entry.size - app2_entry.size
        elif app2_entry.size > app1_entry.size:
            bigger_in_app2[f] = app2_entry.size - app1_entry.size
        else:
            same_size.append(f)
    output = open(options.report, 'w') if options.report else sys.stdout
    print_info(app1, app1_files, only_in_app1, bigger_in_app1, output)
    print_info(app2, app2_files, only_in_app2, bigger_in_app2, output)
    output.write('Same size\n')
    output.write('\n'.join(['  %s' % x for x in same_size]))
    if options.report:
        output.close()


def Main():
    (options, args) = parse_options()
    if len(args) is not 2:
        print(args)
        print('Takes exactly two arguments, the two apps to compare')
        return 1
    app1 = args[0]
    app2 = args[1]
    ensure_exists([app1, app2])
    with utils.TempDir() as temporary:
        # If a temp dir is passed in, use that instead of the generated temporary
        output = options.temp if options.temp else temporary
        ensure_exists([output])
        app1_input = [app1]
        app2_input = [app2]
        if app1.endswith('apk'):
            app1_input = extract_apk(app1, os.path.join(output, 'app1'))
        if app2.endswith('apk'):
            app2_input = extract_apk(app2, os.path.join(output, 'app2'))
        app1_classes_dir = os.path.join(output, 'app1_classes')
        app2_classes_dir = os.path.join(output, 'app2_classes')

        extract_classes(app1_input, app1_classes_dir, options)
        extract_classes(app2_input, app2_classes_dir, options)
        compare(app1_classes_dir, app2_classes_dir, app1, app2, options)


if __name__ == '__main__':
    sys.exit(Main())