import sys
import os
import re
import subprocess, shlex



################################################## download ##################################################
def get_sdk_root_path_by_project_path(ProjectPath):
    Pattern = r'(.*)/boards/'
    SearchResult = re.search(Pattern, ProjectPath, re.S)
    if SearchResult != None:
        return SearchResult.group(1)
    else:
        print("Error: correct ssdk root path was not found")
        sys.exit(1)

def add_double_quotes(Strings):
    return '"' + Strings + '"'

def jlink_exe_file(FilePathName):
    JlinkCmd = JlinkExe + add_double_quotes(FilePathName)
    args = shlex.split(JlinkCmd)
    # subprocess.run(args, check=True, stdout=subprocess.DEVNULL)
    subprocess.run(args, check=True)

def download_file(File, Addr):
    print("loadfile: %s %s" %(File, Addr))
    CmdFilePathName = ProjectPath + "/cmd.txt"
    LoadFileString = "LoadFile " + add_double_quotes(File) + " " + Addr + " noreset"
    ExitString = "q"
    with open(CmdFilePathName, "w") as f:
        f.write("h" + "\n")
        f.write(LoadFileString + "\n")
        f.write(ExitString)
    try:
        jlink_exe_file(CmdFilePathName)
        os.remove(CmdFilePathName)
    except Exception as e:
        os.remove(CmdFilePathName)
        raise Exception("download error:\n    %s" %e)

def get_download_info_by_project_file(ProjectPathName, FlashType):
    with open(ProjectPathName, "r", encoding="utf-8") as f:
        Content = f.read()
    Pattern = r'(<configuration\s*Name="Download_{0}"(?:(?!<|hidden).)*/>)'.format(FlashType)
    SearchResult = re.search(Pattern, Content, re.S)
    if SearchResult != None:
        return SearchResult.group(1)
    else:
        print("Error: no download info find")
        sys.exit(1)

def get_multicore_download_info_by_project_file(ProjectPathName, FlashType, Part):
    with open(ProjectPathName, "r", encoding="utf-8") as f:
        Content = f.read()
    Pattern = r'(<configuration\s*Name="(?:Download|Sign)_{0}_Multicore_{1}"(?:(?!<|hidden).)*/>)'.format(FlashType, Part)
    SearchResultMulticore = re.search(Pattern, Content, re.S)
    if SearchResultMulticore != None:
        return SearchResultMulticore.group(1)
    else:
        return None

def get_str_from_download_info(DownloadInfo, FindString):
    Pattern = FindString + r'="(.*?)"'
    SearchResult = re.search(Pattern, DownloadInfo, re.S)
    if SearchResult != None:
        return SearchResult.group(1)
    else:
        return None

def is_sfs_file(FileName):
    Pattern = r'sfs_.*'
    SearchResult = re.search(Pattern, FileName, re.S)
    if SearchResult != None:
        return True
    else:
        return False

def is_need_sign(FileName):
    if is_sfs_file(FileName) == True:
        return False
    elif FileName == "sf_signed.bin" or FileName == "sp0_signed.bin" or FileName == "sp1_signed.bin" or FileName == "sx0_signed.bin" or FileName == "sx1_signed.bin":
        return True
    elif FileName == "sf.bin" or FileName == "sp0.bin" or FileName == "sp1.bin" or FileName == "sx0.bin" or FileName == "sx1.bin":
        return False
    else:
        print("Error: %s is error file name" %FileName)
        sys.exit(1)

def is_need_download(FileName):
    if is_sfs_file(FileName) == True:
        return True
    elif FileName == "sf_signed.bin" or FileName == "sf.bin" or FileName == "sp0.bin" or FileName == "sp1.bin" or FileName == "sx0.bin" or FileName == "sx1.bin":
        return True
    elif FileName == "sp0_signed.bin" or FileName == "sp1_signed.bin" or FileName == "sx0_signed.bin" or FileName == "sx1_signed.bin":
        return False
    else:
        print("Error: %s is error file name" %FileName)
        sys.exit(1)

def get_unsigned_bin_name(FileName):
    DictTmp = {'sf_signed.bin':'sf.bin', 'sp0_signed.bin':'sp0.bin', 'sp1_signed.bin':'sp1.bin', 'sx0_signed.bin':'sx0.bin', 'sx1_signed.bin':'sx1.bin'}
    return  DictTmp.get(FileName)

def gen_info_from_download_info(DownloadInfo, FileString, AddrString):
    if DownloadInfo == None:
        return
    FilePathName = get_str_from_download_info(DownloadInfo, FileString)
    if FilePathName != None and FilePathName != "":
        FilePathName =  os.path.abspath(FilePathName.replace("$(ProjectDir)", ProjectPath)).replace('\\', '/')
        FilePath = os.path.dirname(FilePathName)
        FileName = os.path.basename(FilePathName)
        if is_need_sign(FileName):
            # 签名
            UnsignedBinPathName = FilePath + "/" + get_unsigned_bin_name(FileName)
            Addr = get_str_from_download_info(DownloadInfo, AddrString)
            if Addr != None and Addr != "":
                UnsignedBinAddr = Addr
                SignInfoDict.update({UnsignedBinPathName:UnsignedBinAddr})
        if is_need_download(FileName):
            Addr = get_str_from_download_info(DownloadInfo, AddrString)
            if Addr != None and Addr != "":
                DownloadInfoDict.update({FilePathName:Addr})

def get_sign_and_download_info_by_project_file(ProjectPathName, FlashType, Part):
    DownloadInfo = get_download_info_by_project_file(ProjectPathName, FlashType)
    MulticoreDownloadInfo = get_multicore_download_info_by_project_file(ProjectPathName, FlashType, Part)
    FileString = "debug_additional_load_file"
    AddrStr = "debug_additional_load_file_address"
    gen_info_from_download_info(DownloadInfo, FileString, AddrStr)
    gen_info_from_download_info(MulticoreDownloadInfo, FileString, AddrStr)
    FileString = "debug_additional_load_file1"
    AddrStr = "debug_additional_load_file_address1"
    gen_info_from_download_info(DownloadInfo, FileString, AddrStr)
    gen_info_from_download_info(MulticoreDownloadInfo, FileString, AddrStr)
    FileString = "debug_additional_load_file2"
    AddrStr = "debug_additional_load_file_address2"
    gen_info_from_download_info(DownloadInfo, FileString, AddrStr)
    gen_info_from_download_info(MulticoreDownloadInfo, FileString, AddrStr)
    FileString = "debug_additional_load_file3"
    AddrStr = "debug_additional_load_file_address3"
    gen_info_from_download_info(DownloadInfo, FileString, AddrStr)
    gen_info_from_download_info(MulticoreDownloadInfo, FileString, AddrStr)
    FileString = "external_build_file_name"
    AddrStr = "external_load_address"
    gen_info_from_download_info(DownloadInfo, FileString, AddrStr)
    gen_info_from_download_info(MulticoreDownloadInfo, FileString, AddrStr)

def download_file_by_download_info(**DownloadInfoDict):
    for key, value in DownloadInfoDict.items():
        if not os.path.exists(key):
            print("Error: %s does not exist" %key)
            sys.exit(1)
        download_file(key, value)



################################################## sign ##################################################
def get_core_by_pathname(FilePathName):
    FileName = os.path.basename(FilePathName)
    DictTmp = {'sf.bin':'0', 'sp0.bin':'2', 'sp1.bin':'3', 'sx0.bin':'4', 'sx1.bin':'5'}
    return  DictTmp.get(FileName)

def get_entry_by_pathname(FilePathName, PartId):
    FileName = os.path.basename(FilePathName)
    if PartId == "E3104" or PartId == "E3106" or PartId == "E3205" or PartId == "E3206":
        DictTmp = {'sf.bin':'0x504000'}
    elif PartId == "E3110" or PartId == "E3210" or PartId == "E3340" or PartId == "D3246" or PartId == "D3248":
        DictTmp = {'sf.bin':'0x404000'}
    elif PartId == "E3420":
        DictTmp = {'sf.bin':'0x404000', 'sx0.bin':'0x500000', 'sx1.bin':'0x580000'}
    elif PartId == "E3430":
        DictTmp = {'sf.bin':'0x404000', 'sx0.bin':'0x600000', 'sx1.bin':'0x680000'}
    elif PartId == "E3640" or PartId == "E3648":
        DictTmp = {'sf.bin':'0x404000', 'sp0.bin':'0x700000', 'sp1.bin':'0x780000', 'sx0.bin':'0x600000', 'sx1.bin':'0x680000'}
    else:
        print("Error: %s does not exist" %PartId)
        sys.exit(1)
    if DictTmp.get(FileName) == None:
        print("Error: %s does not exist" %FileName)
        sys.exit(1)
    return  DictTmp.get(FileName)

def get_bpt_base_by_part(FlashType):
    if FlashType == "HYPERFLASH":
        return 0x100C0000
    elif FlashType == "NORFLASH":
        return 0x10007000
    else:
        print("Error: FlashType error")
        sys.exit(1)

def get_dlp_by_addr(Addr, FlashType):
    return int((int(Addr, 16) - get_bpt_base_by_part(FlashType)) / 512)

def get_iib_info_by_sign_info(**SignInfoDict):
    global SignedBinPathName
    IIBInfo = ""
    for key, value in SignInfoDict.items():
        if not os.path.exists(key):
            print("Error: %s does not exist" %key)
            sys.exit(1)
        if os.path.basename(key) == "sf.bin":
            SignedBinPathName = add_double_quotes(os.path.dirname(key) + "/sf_signed.bin")
        IIBInfo += "--iib core=%s type=0x0 image=%s dlp=0x%x to=%s entry=%s\n" % (get_core_by_pathname(key), add_double_quotes(key), get_dlp_by_addr(value, FlashType), get_entry_by_pathname(key, Part), get_entry_by_pathname(key, Part))
    return IIBInfo

def sign_file_by_sign_info(**SignInfoDict):
    SecVer = 0
    KeyName = "TestRSA2048_ossl.pem"
    PSN = "0x100"
    IIBInfo = get_iib_info_by_sign_info(**SignInfoDict)
    ATB_SIGNER = "./sign_tool/windows/atb_signer.exe"
    os.chdir(os.path.join(SdkRootPath, "tools", "sdtools"))
    SignCmd = ATB_SIGNER + " sign --v 2 --sec_ver %s --dgst sha256 --rcp key=sign_tool/keys/%s\n%s--psn %s --of %s" %(SecVer, KeyName, IIBInfo, PSN, SignedBinPathName)
    print("SignCmd: %s" %SignCmd)
    print("sign file: %s" %SignedBinPathName)
    try:
        args = shlex.split(SignCmd)
        # subprocess.run(args, check=True, stdout=subprocess.DEVNULL)
        subprocess.run(args, check=True)
    except Exception as e:
        raise Exception("sign error:\n    %s" %e)



################################################## main ##################################################
if __name__ == "__main__":
    # 处理入参
    JlinkPath = os.path.abspath(sys.argv[1]).replace('\\', '/')
    ProjectPathName = os.path.abspath(sys.argv[2]).replace('\\', '/')
    Part = sys.argv[3]
    FlashType = sys.argv[4]
    OperationType = sys.argv[5]

    # 生成需要的变量
    ProjectPath = os.path.dirname(ProjectPathName)
    SdkRootPath = get_sdk_root_path_by_project_path(ProjectPath)

    # 如果是ide，进行签名后退出；如果是cmd，签名后继续下载
    if OperationType == "$(OPS)":
        OnlySignBin = True
    elif OperationType == "download":
        OnlySignBin = False
    else:
        print("Error: %s is unsupport operation type" %OperationType)
        sys.exit(1)



    SignInfoDict = {}
    DownloadInfoDict = {}
    get_sign_and_download_info_by_project_file(ProjectPathName, FlashType, Part)
    print(SignInfoDict)
    print(DownloadInfoDict)

    if not SignInfoDict:
        print("Error: no signed files info found")
        sys.exit(1)
    sign_file_by_sign_info(**SignInfoDict)

    # 初始化TCM，擦除Flash
    if OnlySignBin == False:
        if not DownloadInfoDict:
            print("Error: no download files info found")
            sys.exit(1)
        DeviceName = Part + "_" + FlashType
        JlinkPathName = add_double_quotes(JlinkPath + '/JLink.exe')
        JlinkExe = JlinkPathName + " -device " + DeviceName + " -if SWD -speed 4000 -autoconnect 1 -CommandFile "

        if FlashType == "HYPERFLASH":
            InitFilePathName = os.path.join(SdkRootPath, "devices", Part, "flashloader", "init_before_download_hyperflash.txt").replace('\\', '/')
        elif FlashType == "NORFLASH":
            InitFilePathName = os.path.join(SdkRootPath, "devices", Part, "flashloader", "init_before_download_norflash.txt").replace('\\', '/')
        else:
            print("Error: %s is not right type" %FlashType)
            sys.exit(1)

        print("init tcm and erase flash: %s" %InitFilePathName)
        try:
            jlink_exe_file(InitFilePathName)
        except Exception as e:
            raise Exception("download error:\n    %s" %e)

        download_file_by_download_info(**DownloadInfoDict)