#!/usr/bin/env python3

import subprocess
import sys
from dataclasses import dataclass
from pprint import pprint
from functools import lru_cache
from tempfile import TemporaryDirectory
from zipfile import ZipFile
import os

@dataclass
class AndroidIdent:
    ndk: str
    ver: int

@lru_cache
def get_objdump():
    # Stub in case I need to do something more clever to find
    # objdump at some point in the future
    return "objdump"

def get_android_ident_for_so_file(filename: str):
    process_result = subprocess.run(args=[
        get_objdump(), "--full-contents", "--section", ".note.android.ident", filename
    ], capture_output=True)

    if process_result.returncode != 0:
        print("Couldn't get Android NDK version info from file", filename, file=sys.stderr)
        return AndroidIdent("Unknown", -1)

    # Get all lines of the output from objdump as strings, discarding empty lines
    all_ident_lines = [line for line in process_result.stdout.decode("utf-8").splitlines() if line != ""]

    # Strip off the header
    ident_lines = all_ident_lines[3:]

    # Each line looks like:
    #  0200 08000000 84000000 01000000 416e6472  ............Andr

    # Split on whitespace to get components
    # ['0200', '08000000', '84000000', '01000000', '416e6472', '............Andr']
    # [0] is file offset
    # [1:-1] are hex data
    # [-1] is decoded data
    # Offsets from the end because last line has less data in it
    split_lines = [line.split() for line in ident_lines]
    
    # Join decoded together into a single line, which looks like
    # ............Android........r23b........23948123
    decoded = "".join((line[-1] for line in split_lines))

    # Pull out ["Android", "r23b", 23948123]
    decoded_split = [x for x in decoded.split(".") if x != ""]

    if len(decoded_split) < 3:
        print("No Android NDK version info in file:", filename, decoded, decoded_split, process_result.stdout.decode("utf-8"), file=sys.stderr)
        return AndroidIdent("Unknown", -1)
    
    return AndroidIdent(decoded_split[1], int(decoded_split[2]))

def get_needed_shared_libs_from_so(filename: str):
    process_result = subprocess.run(args=[
        get_objdump(), "-p", filename
    ], capture_output=True)

    lines = [line.split() for line in process_result.stdout.decode("utf-8").splitlines()]

    return [line[1] for line in lines if len(line) >=2 and line[0] == "NEEDED"]


def get_from_aar(filename: str):
    files = {}

    with ZipFile(filename) as aarfile:
        with TemporaryDirectory() as temp_dir:
            for f in aarfile.infolist(): 
                if f.filename.endswith(".so"):
                    # Extract to temporary
                    temp_file = aarfile.extract(f, temp_dir)

                    idents = get_android_ident_for_so_file(temp_file)
                    needed_libs = get_needed_shared_libs_from_so(temp_file)

                    synthetic_filename = os.path.join(filename, f.filename)

                    files[synthetic_filename] = {
                        "idents": idents,
                        "needed_libs": needed_libs
                    }

    return files
    

def get_from_so(filename: str):
    return {
        filename: {
            "idents": get_android_ident_for_so_file(filename),
            "needed_libs": get_needed_shared_libs_from_so(filename)
        }
    }

def get_from_path(path: str):
    if not os.path.exists(path):
        print("Path doesn't exist:", path, file=sys.stderr)
        return {}
        
    if os.path.isfile(path):
        if path.endswith(".aar"):
            return get_from_aar(path)
        if path.endswith(".so"):
            return get_from_so(path)
    elif os.path.isdir(path):
        files = {}
        for dirpath, _dirnames, dirfilenames in os.walk(path):
            for filename in dirfilenames:
                files |= get_from_path(os.path.join(dirpath, filename))
        return files
        
    return {}

def main():
    files = {}

    # Run commands on each path
    for path in sys.argv[1:]:
        files |= get_from_path(path)

    idents = {f: files[f]["idents"] for f in files}
    needed_libs = {f: files[f]["needed_libs"] for f in files}

    # Analyse results
    invert_dict = {}
    
    for f in idents:
        ndk = idents[f].ndk

        # Filter to libc++shared.so itself and any libs which need it
        if("libc++_shared.so" in needed_libs[f] or f.endswith("libc++_shared.so")):
            if ndk not in invert_dict:
                invert_dict[ndk] = []
            invert_dict[ndk].append(f)

    pprint(invert_dict)

if __name__ == "__main__":
    exit(main())