#!/usr/bin/env python3
# encoding=utf-8
# ============================================================================
# @brief    indie upgrade build file
# Copyright HiSilicon (Shanghai) Technologies Co., Ltd. 2022-2022. All rights reserved.
# ============================================================================

import os
import sys
import shutil
import hashlib
import json
import subprocess

root_dir = os.path.split(os.path.realpath(__file__))[0]
sys.path.append(os.path.join(root_dir, 'build', 'config'))
sys.path.append(os.path.join(root_dir, 'build', 'script'))

from cmake_builder import CMakeBuilder
from utils.build_utils import root_path, output_root

sys.path.append(os.path.join(root_dir, 'build', 'config', "target_config", "ws63"))
from target_config.ws63.build_ws63_update import upg_base_info

def exec_shell(cmd, cwd=None, cmd_dump=False):
    cmd_list = cmd
    if isinstance(cmd, str):
        cmd_list = cmd.split(' ')

    if cmd_dump:
        print(str(cmd_list))

    subp = subprocess.Popen(cmd_list, shell=False, cwd=cwd)
    subp.wait()
    code = subp.returncode

    return code

def copy_files_to_path(src_files):
    cnts = 0
    for file in src_files:
        if not os.path.isfile(file):
            print(f'[!]File `{file}` not found, skip!')
            continue
        try:
            shutil.copy(file, src_files[file])
        except BaseException as e:
            print(e)
            exit(-1)
        cnts += 1
    return cnts

def build_update_package(update_bin_path, temp_dir, output_path):
    with open(update_bin_path, "rb") as fp:
        update_pkg = fp.read()
    common_head_sh = hashlib.sha256()
    common_head_sh.update(update_pkg)
    common_head_hash = common_head_sh.digest().hex().upper()
    print("build update package cal sha256", common_head_hash)
    filelist = {
        "image2_all_ota1.bin": {
            "sha256": common_head_hash,
        },
        "image2_all_ota2.bin": {
            "sha256": common_head_hash,
        },
    }

    pkg_dir = os.path.join(temp_dir, "package")
    if os.path.exists(pkg_dir):
        shutil.rmtree(pkg_dir)
    os.makedirs(pkg_dir)
    pkg_inner_dir = os.path.join(pkg_dir, "package")
    os.makedirs(pkg_inner_dir)

    with open(os.path.join(pkg_inner_dir, "filelist.json"), "w", encoding="utf-8") as fp:
        fp.write(json.dumps(filelist, indent=4))
    try:
        shutil.copy(update_bin_path, os.path.join(pkg_inner_dir, "image2_all_ota1.bin"))
        shutil.copy(update_bin_path, os.path.join(pkg_inner_dir, "image2_all_ota2.bin"))
    except BaseException as e:
        shutil.rmtree(pkg_dir)
        print(e)
        exit(-1)

    pkg_file = os.path.join(output_path, "package.zip")
    if os.path.exists(pkg_file):
        print("pkg file exist, delete")
        os.unlink(pkg_file)
    errcode = exec_shell(["zip", "-r", pkg_file, "."], pkg_dir, True)
    if errcode != 0:
        shutil.rmtree(pkg_dir)
        print("make pkg shell cmd err:", errcode)
        exit(-1)

    shutil.rmtree(pkg_dir)

def pack_fwpkg():
    packet_script_path = os.path.join(root_path, "tools", "pkg", "packet.py")
    print("packet ing...")
    errcode = exec_shell([sys.executable, packet_script_path, "ws63", "ws63-liteos-app-iot", ""], None, True)
    if errcode != 0:
        print("packet error!")
        sys.exit(-1)
    print("packet succ")

def all_build():
    builder = CMakeBuilder(["-c", "ws63-liteos-hilink"])
    builder.build()
    builder = CMakeBuilder(["-c", "ws63-liteos-app-iot"])
    builder.build()
    build_update_script_path = os.path.join(root_path, "build", "config", "target_config", "ws63", \
                                            "build_ws63_update.py")
    errcode = exec_shell([sys.executable, build_update_script_path, "--pkt=app_iot"], None, True)
    if errcode != 0:
        print("build ws63 update error!")
        exit(-1)

    info = upg_base_info()

    temp_dir = os.path.join(output_root, "ws63", "temp")
    if os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
    os.makedirs(temp_dir)

    build_update_package(os.path.join(info.upg_output, "update.fwpkg"), temp_dir, info.upg_output)

    fwpkg_src_file = os.path.join(info.output, "fwpkg", "ws63-liteos-app-iot", "ws63-liteos-app-iot_all.fwpkg")
    copy_files = {
        info.hilink_bin: temp_dir,
        info.hilink_check: temp_dir,
        info.app_iot_bin: temp_dir,
        info.app_iot_check: temp_dir,
        fwpkg_src_file: temp_dir,
        os.path.join(info.upg_output, "package.zip"): temp_dir,
    }

    if copy_files_to_path(copy_files) != len(copy_files):
        shutil.rmtree(temp_dir)
        exit(-1)

    total_file = os.path.join(output_root, "ws63", "ws63-liteos_all.zip")
    if os.path.exists(total_file):
        print("total_file exist, delete")
        os.unlink(total_file)
    errcode = exec_shell(["zip", "-rj", total_file, temp_dir], None, True)
    if errcode != 0:
        shutil.rmtree(temp_dir)
        print("make full pkg err:", errcode)
        exit(-1)
    shutil.rmtree(temp_dir)
    print("gen package.zip")
    file_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "package.py")
    os.system("cd " + file_dir + " && python3 package.py")

def sdk_build():
    builder = CMakeBuilder(["-c", "ws63-liteos-hilink"])
    builder.build()

    pack_fwpkg()

    build_update_script_path = os.path.join(root_path, "build", "config", "target_config", "ws63", \
                                            "build_ws63_update.py")
    errcode = exec_shell([sys.executable, build_update_script_path, "--pkt=app_iot"], None, True)
    if errcode != 0:
        print("build ws63 update error!")
        exit(-1)

    info = upg_base_info()

    temp_dir = os.path.join(output_root, "ws63", "temp")
    if os.path.exists(temp_dir):
        shutil.rmtree(temp_dir)
    os.makedirs(temp_dir)

    build_update_package(os.path.join(info.upg_output, "update.fwpkg"), temp_dir, info.upg_output)

    shutil.rmtree(temp_dir)

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("param num err")
    if sys.argv[1] == "sdk":
        sdk_build()
    elif sys.argv[1] == "all":
        all_build()
    else:
        print("param err, should be: 'sdk' or 'all'")