import sys
import shutil
import os
from pathlib import Path
import subprocess

class BuildError(Exception):
    def __init__(self):
        super.__init__()

def getDexName(smaliName):
    if smaliName == 'smali':
        return 'classes.dex'
    elif smaliName.startswith('smali_classes'):
        return 'classes' + smaliName[13:] + '.dex'
    print("Not a smali name")
    raise BuildError()

class Task:
    def clean():
        shutil.rmtree('build')
    
    def extractApk():
        if Path('build/extracted').exists():
            return
        Path('build').mkdir(parents=True, exist_ok=True)

        apk = list(filter(lambda f: f.endswith('.apk'), os.listdir('source-apk')))
        if len(apk) == 0:
            print('Please place the APK inside the source-apk directory.')
            raise BuildError()
        elif len(apk) > 1:
            print('More than one apk found in source-apk. Please remove all but one.')
            raise BuildError()

        # Extract the apk
        print("Extracting APK...")
        subprocess.run(['java', '-jar', './tools/apktool_2.9.3.jar', 'd', 'source-apk/' + apk[0], '-o', 'build/extracted'])

    def copySmali():
        Task.extractApk()

        Path('build/smali').mkdir(parents=True, exist_ok=True)
        # Copy smali classes that haven't been yet
        for file in os.listdir('build/extracted'):
            if file != 'smali' and not file.startswith('smali_classes'):
                continue
            
            # Skip if already exists
            if Path(f'build/smali/{file}').exists():
                continue

            print(f'Copying {file}...')
            shutil.copytree(f'build/extracted/{file}', f'build/smali/{file}')

        # Copy from src
        print('Copying ykit smali...')
        shutil.copytree('src/smali', 'build/smali', dirs_exist_ok=True)

    def compileSmali():
        Task.copySmali()

        Path('build/dex').mkdir(parents=True, exist_ok=True)

        def shouldCompileSmali(file, dex):
            src = f'build/smali/{file}'
            ykitsrc = f'src/smali/{file}'
            dst = f'build/dex/{dex}'

            # If dest doesn't exist, always build
            if not Path(dst).exists():
                return True

            # Check if ykit src dir exists for it, if not then don't recompile
            if not Path(ykitsrc).exists():
                return False

            # Get latest file timestamp in ykitsrc
            files = list(Path(ykitsrc).rglob('*.smali'))
            files = [os.path.getmtime(x) for x in files]
            latest = max(files)

            dsttime = os.path.getmtime(dst)

            # If it is newer than the dest, then do recompile
            if latest > dsttime:
                return True

            # Otherwise, don't
            return False

        for file in os.listdir('build/smali'):
            dex = getDexName(file)

            if not shouldCompileSmali(file, dex):
                print(f"Skipping {file}...")
                continue

            print(f"Compiling {file}...")
            subprocess.run(['java', '-jar', 'tools/smali-2.5.2.jar', 'a', f'build/smali/{file}', '-o', f'build/dex/{dex}'])

def main():

    match sys.argv[1]:
        case 'clean':
            Task.clean()
        case 'extract':
            Task.extractApk()
        case 'build':
            Task.compileSmali()

if __name__ == '__main__':
    main()