#!/usr/bin/env python

import argparse
import json
from os import path
import os
import shutil
import subprocess

from build_java_shared_aar import cleanup, fill_template, get_compiled_aar_path, get_opencv_version


ANDROID_PROJECT_TEMPLATE_DIR = path.join(path.dirname(__file__), "aar-template")
TEMP_DIR = "build_static"
ANDROID_PROJECT_DIR = path.join(TEMP_DIR, "AndroidProject")
COMPILED_AAR_PATH_1 = path.join(ANDROID_PROJECT_DIR, "OpenCV/build/outputs/aar/OpenCV-release.aar") # original package name
COMPILED_AAR_PATH_2 = path.join(ANDROID_PROJECT_DIR, "OpenCV/build/outputs/aar/opencv-release.aar") # lower case package name
AAR_UNZIPPED_DIR = path.join(TEMP_DIR, "aar_unzipped")
FINAL_AAR_PATH_TEMPLATE = "outputs/opencv_static_<OPENCV_VERSION>.aar"
FINAL_REPO_PATH = "outputs/maven_repo"
MAVEN_PACKAGE_NAME = "opencv-static"


def get_list_of_opencv_libs(sdk_dir):
    files = os.listdir(path.join(sdk_dir, "sdk/native/staticlibs/arm64-v8a"))
    libs = [f[3:-2] for f in files if f[:3] == "lib" and f[-2:] == ".a"]
    return libs

def get_list_of_3rdparty_libs(sdk_dir, abis):
    libs = []
    for abi in abis:
        files = os.listdir(path.join(sdk_dir, "sdk/native/3rdparty/libs/" + abi))
        cur_libs = [f[3:-2] for f in files if f[:3] == "lib" and f[-2:] == ".a"]
        for lib in cur_libs:
            if lib not in libs:
                libs.append(lib)
    return libs

def add_printing_linked_libs(sdk_dir, opencv_libs):
    """
    Modifies CMakeLists.txt file in Android project, so it prints linked libraries for each OpenCV library"
    """
    sdk_jni_dir = sdk_dir + "/sdk/native/jni"
    with open(path.join(ANDROID_PROJECT_DIR, "OpenCV/src/main/cpp/CMakeLists.txt"), "a") as f:
        f.write('\nset(OpenCV_DIR "' + sdk_jni_dir + '")\n')
        f.write('find_package(OpenCV REQUIRED)\n')
        for lib_name in opencv_libs:
            output_filename_prefix = "linkedlibs." + lib_name + "."
            f.write('get_target_property(OUT "' + lib_name + '" INTERFACE_LINK_LIBRARIES)\n')
            f.write('file(WRITE "' + output_filename_prefix + '${ANDROID_ABI}.txt" "${OUT}")\n')

def read_linked_libs(lib_name, abis):
    """
    Reads linked libs for each OpenCV library from files, that was generated by gradle. See add_printing_linked_libs()
    """
    deps_lists = []
    for abi in abis:
         with open(path.join(ANDROID_PROJECT_DIR, "OpenCV/src/main/cpp", f"linkedlibs.{lib_name}.{abi}.txt")) as f:
            text = f.read()
            linked_libs = text.split(";")
            linked_libs = [x.replace("$<LINK_ONLY:", "").replace(">", "") for x in linked_libs]
            deps_lists.append(linked_libs)

    return merge_dependencies_lists(deps_lists)

def merge_dependencies_lists(deps_lists):
    """
    One library may have different dependencies for different ABIS.
    We need to merge them into one list with all the dependencies preserving the order.
    """
    result = []
    for d_list in deps_lists:
        for i in range(len(d_list)):
            if d_list[i] not in result:
                if i == 0:
                    result.append(d_list[i])
                else:
                    index = result.index(d_list[i-1])
                    result = result[:index + 1] + [d_list[i]] + result[index + 1:]

    return result

def convert_deps_list_to_prefab(linked_libs, opencv_libs, external_libs):
    """
    Converting list of dependencies into prefab format.
    """
    prefab_linked_libs = []
    for lib in linked_libs:
        if (lib in opencv_libs) or (lib in external_libs):
            prefab_linked_libs.append(":" + lib)
        elif (lib[:3] == "lib" and lib[3:] in external_libs):
            prefab_linked_libs.append(":" + lib[3:])
        elif lib == "ocv.3rdparty.android_mediandk":
            prefab_linked_libs += ["-landroid", "-llog", "-lmediandk"]
            print("Warning: manualy handled ocv.3rdparty.android_mediandk dependency")
        elif lib == "ocv.3rdparty.flatbuffers":
            print("Warning: manualy handled ocv.3rdparty.flatbuffers dependency")
        elif lib.startswith("ocv.3rdparty"):
            raise Exception("Unknown lib " + lib)
        else:
            prefab_linked_libs.append("-l" + lib)
    return prefab_linked_libs

def main(args):
    opencv_version = get_opencv_version(args.opencv_sdk_path)
    abis = os.listdir(path.join(args.opencv_sdk_path, "sdk/native/libs"))
    final_aar_path = FINAL_AAR_PATH_TEMPLATE.replace("<OPENCV_VERSION>", opencv_version)
    sdk_dir = args.opencv_sdk_path

    print("Removing data from previous runs...")
    cleanup([TEMP_DIR, final_aar_path, path.join(FINAL_REPO_PATH, "org/opencv", MAVEN_PACKAGE_NAME)])

    print("Preparing Android project...")
    # ANDROID_PROJECT_TEMPLATE_DIR contains an Android project template that creates AAR
    shutil.copytree(ANDROID_PROJECT_TEMPLATE_DIR, ANDROID_PROJECT_DIR)

    # Configuring the Android project to static C++ libs version
    fill_template(path.join(ANDROID_PROJECT_DIR, "OpenCV/build.gradle.template"),
                  path.join(ANDROID_PROJECT_DIR, "OpenCV/build.gradle"),
                  {"LIB_NAME": "templib",
                   "LIB_TYPE": "c++_static",
                   "PACKAGE_NAME": MAVEN_PACKAGE_NAME,
                   "OPENCV_VERSION": opencv_version,
                   "COMPILE_SDK": args.android_compile_sdk,
                   "MIN_SDK": args.android_min_sdk,
                   "TARGET_SDK": args.android_target_sdk,
                   "ABI_FILTERS": ", ".join(['"' + x + '"' for x in abis]),
                   "JAVA_VERSION": args.java_version,
                   })
    fill_template(path.join(ANDROID_PROJECT_DIR, "OpenCV/src/main/cpp/CMakeLists.txt.template"),
                  path.join(ANDROID_PROJECT_DIR, "OpenCV/src/main/cpp/CMakeLists.txt"),
                  {"LIB_NAME": "templib", "LIB_TYPE": "STATIC"})

    local_props = ""
    if args.ndk_location:
        local_props += "ndk.dir=" + args.ndk_location + "\n"
    if args.cmake_location:
        local_props += "cmake.dir=" + args.cmake_location + "\n"

    if local_props:
        with open(path.join(ANDROID_PROJECT_DIR, "local.properties"), "wt") as f:
            f.write(local_props)

    opencv_libs = get_list_of_opencv_libs(sdk_dir)
    external_libs = get_list_of_3rdparty_libs(sdk_dir, abis)

    add_printing_linked_libs(sdk_dir, opencv_libs)

    print("Running gradle assembleRelease...")
    cmd = ["./gradlew", "assembleRelease"]
    if args.offline:
        cmd = cmd + ["--offline"]
    # Running gradle to build the Android project
    subprocess.run(cmd, shell=False, cwd=ANDROID_PROJECT_DIR, check=True)

    # The created AAR package contains only one empty libtemplib.a library.
    # We need to add OpenCV libraries manually.
    # AAR package is just a zip archive
    complied_aar_path = get_compiled_aar_path(COMPILED_AAR_PATH_1, COMPILED_AAR_PATH_2) # two possible paths
    shutil.unpack_archive(complied_aar_path, AAR_UNZIPPED_DIR, "zip")

    print("Adding libs to AAR...")

    # Copying 3rdparty libs from SDK into the AAR
    for lib in external_libs:
        for abi in abis:
            os.makedirs(path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/libs/android." + abi))
            if path.exists(path.join(sdk_dir, "sdk/native/3rdparty/libs/" + abi, "lib" + lib + ".a")):
                shutil.copy(path.join(sdk_dir, "sdk/native/3rdparty/libs/" + abi, "lib" + lib + ".a"),
                            path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/libs/android." + abi, "lib" + lib + ".a"))
            else:
                # One OpenCV library may have different dependency lists for different ABIs, but we can write only one
                # full dependency list for all ABIs. So we just add empty .a library if this ABI doesn't have this dependency.
                shutil.copy(path.join(AAR_UNZIPPED_DIR, "prefab/modules/templib/libs/android." + abi, "libtemplib.a"),
                            path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/libs/android." + abi, "lib" + lib + ".a"))
            shutil.copy(path.join(AAR_UNZIPPED_DIR, "prefab/modules/templib/libs/android." + abi + "/abi.json"),
                        path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/libs/android." + abi + "/abi.json"))
        shutil.copy(path.join(AAR_UNZIPPED_DIR, "prefab/modules/templib/module.json"),
                    path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/module.json"))

    # Copying OpenV libs from SDK into the AAR
    for lib in opencv_libs:
        for abi in abis:
            os.makedirs(path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/libs/android." + abi))
            shutil.copy(path.join(sdk_dir, "sdk/native/staticlibs/" + abi, "lib" + lib + ".a"),
                        path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/libs/android." + abi, "lib" + lib + ".a"))
            shutil.copy(path.join(AAR_UNZIPPED_DIR, "prefab/modules/templib/libs/android." + abi + "/abi.json"),
                        path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/libs/android." + abi + "/abi.json"))
        os.makedirs(path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/include/opencv2"))
        shutil.copy(path.join(sdk_dir, "sdk/native/jni/include/opencv2/" + lib.replace("opencv_", "") + ".hpp"),
                    path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/include/opencv2/" + lib.replace("opencv_", "") + ".hpp"))
        module_include_folder = path.join(sdk_dir, "sdk/native/jni/include/opencv2/" + lib.replace("opencv_", ""))
        if os.path.exists(module_include_folder):
            shutil.copytree(module_include_folder,
                            path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/include/opencv2/" + lib.replace("opencv_", "")))

        # Adding dependencies list
        module_json_text = {
            "export_libraries": convert_deps_list_to_prefab(read_linked_libs(lib, abis), opencv_libs, external_libs),
            "android": {},
        }
        with open(path.join(AAR_UNZIPPED_DIR, "prefab/modules/" + lib + "/module.json"), "w") as f:
            json.dump(module_json_text, f)

    for h_file in ("cvconfig.h", "opencv.hpp", "opencv_modules.hpp"):
        shutil.copy(path.join(sdk_dir, "sdk/native/jni/include/opencv2/" + h_file),
                    path.join(AAR_UNZIPPED_DIR, "prefab/modules/opencv_core/include/opencv2/" + h_file))


    shutil.rmtree(path.join(AAR_UNZIPPED_DIR, "prefab/modules/templib"))

    # Creating final AAR zip archive
    os.makedirs("outputs", exist_ok=True)
    shutil.make_archive(final_aar_path, "zip", AAR_UNZIPPED_DIR, ".")
    os.rename(final_aar_path + ".zip", final_aar_path)

    print("Creating local maven repo...")

    shutil.copy(final_aar_path, path.join(ANDROID_PROJECT_DIR, "OpenCV/opencv-release.aar"))

    print("Creating a maven repo from project sources (with sources jar and javadoc jar)...")
    cmd = ["./gradlew", "publishReleasePublicationToMyrepoRepository"]
    if args.offline:
        cmd = cmd + ["--offline"]
    subprocess.run(cmd, shell=False, cwd=ANDROID_PROJECT_DIR, check=True)

    os.makedirs(path.join(FINAL_REPO_PATH, "org/opencv"), exist_ok=True)
    shutil.move(path.join(ANDROID_PROJECT_DIR, "OpenCV/build/repo/org/opencv", MAVEN_PACKAGE_NAME),
                path.join(FINAL_REPO_PATH, "org/opencv", MAVEN_PACKAGE_NAME))

    print("Creating a maven repo from modified AAR (with cpp libraries)...")
    cmd = ["./gradlew", "publishModifiedPublicationToMyrepoRepository"]
    if args.offline:
        cmd = cmd + ["--offline"]
    subprocess.run(cmd, shell=False, cwd=ANDROID_PROJECT_DIR, check=True)

    # Replacing AAR from the first maven repo with modified AAR from the second maven repo
    shutil.copytree(path.join(ANDROID_PROJECT_DIR, "OpenCV/build/repo/org/opencv", MAVEN_PACKAGE_NAME),
                    path.join(FINAL_REPO_PATH, "org/opencv", MAVEN_PACKAGE_NAME),
                    dirs_exist_ok=True)

    print("Done")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Builds AAR with static C++ libs from OpenCV SDK")
    parser.add_argument('opencv_sdk_path')
    parser.add_argument('--android_compile_sdk', default="31")
    parser.add_argument('--android_min_sdk', default="21")
    parser.add_argument('--android_target_sdk', default="31")
    parser.add_argument('--java_version', default="1_8")
    parser.add_argument('--ndk_location', default="")
    parser.add_argument('--cmake_location', default="")
    parser.add_argument('--offline', action="store_true", help="Force Gradle use offline mode")
    args = parser.parse_args()

    main(args)