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 mergeResources():
        Task.compileSmali()

        # Copy original resources
        # on first run only
        if not Path('build/merged').exists():
            print('Copying extracted resources...')
            shutil.copytree('build/extracted', 'build/merged', dirs_exist_ok=True, ignore=shutil.ignore_patterns('smali', 'smali_classes*'))

        print("Copying ykit resources...")
        shutil.copytree('src/resources', 'build/merged', dirs_exist_ok=True)

        print("Copying compiled dex files...")
        shutil.copytree('build/dex', 'build/merged', dirs_exist_ok=True)

    def buildApk():
        Task.mergeResources()

        debugFlag = ('-d',) if ('d' in ''.join(sys.argv[2:])) else ()

        print("Building apk...")
        subprocess.run(['java', '-jar', './tools/apktool_2.9.3.jar', 'b', *debugFlag, 'build/merged', '-o', 'build/apk/tumblr-ykit-unsigned.apk'])

    def alignApk():
        Task.buildApk()

        print('Aligning apk...')
        subprocess.run(['tools/zipalign', '-p', '-f', '-v', '4', 'build/apk/tumblr-ykit-unsigned.apk', 'build/apk/tumblr-ykit.apk'])

    def debugKey():
        if Path('keystores/debug.keystore').exists():
            return

        print("Debug keystore does not exist, generating...")
        subprocess.run(['keytool', '-genkey', '-v', '-keystore', 'keystores/debug.keystore', '-alias', 'alias_name', '-keyalg', 'RSA', '-keysize', '2048', '-validity', '10000', '-storepass', '123456', '-keypass', '123456', '-dname', 'CN=, OU=, O=, L=, ST=, C='])

    def signApk():
        Task.alignApk()
        Task.debugKey()

        debugFlag = 'd' in ''.join(sys.argv[2:])
        releaseFlag = 'r' in ''.join(sys.argv[2:])

        if not (debugFlag or releaseFlag):
            print("Neither debug nor release flag was specified so the apk will not be signed.")
            return

        if releaseFlag:
            if not Path('keystores/release.keystore').exists():
                print("release.keystore is missing from keystores directory")
                return
            
            keystoreOpts = ('--ks', 'keystores/release.keystore',)
        else:
            Task.debugKey()
            keystoreOpts = ('--ks', 'keystores/debug.keystore', '--ks-pass', 'pass:123456',)

        print('Signing apk...')
        subprocess.run(['java', '-jar', './tools/apksigner.jar', 'sign', *keystoreOpts, 'build/apk/tumblr-ykit.apk'])
        print("Successfully signed apk in 'build/apk/tumblr-ykit.apk' using " + ('release' if releaseFlag else 'debug') + ' keystore')

    def deployToAndroidStudio():
        projDir = Path.home() / 'ApkProjects/tumblr-ykit'
        if not projDir.exists():
            print("No project named 'tumblr-ykit' in ~/ApkProjects/")

        shutil.copy('build/apk/tumblr-ykit.apk', projDir / 'tumblr-ykit.apk')
        print("Deployed to Android Studio")

def main():

    match sys.argv[1]:
        case 'clean':
            Task.clean()
        case 'extract':
            Task.extractApk()
        case 'assemble':
            Task.buildApk()
        case 'build':
            Task.signApk()
            
            if ('a' in ''.join(sys.argv[2:])):
                Task.deployToAndroidStudio()

if __name__ == '__main__':
    main()