import re
import sys
import struct
import json


# kunlun1,kunlun2,kunlun3的img类型
# FRW是指未压缩的fw CFRW是指压缩之后的fw
_ImgType={
  0x0000:"imgROM", # Maybe,ROM doesn't need a imgType */
  0xFC00:"imgSP", # sp */
  0xFC01:"imgSBL", # sbl */
  0xFC02:"imgFRW", # kl1: 1MB fw, uncompressed for sta, run in flash; compressed for cco, run in psram */
  0xFC03:"imgCFG", # The first pib file. */
  0xFC04:"imgEFUSE", # efuse */
  0xFC05:"imgOEM", # oem */
  0xFC06:"imgFRW2", # The second fw, maybe the backup one. not used. */
  0xFC07:"imgCFG2", # The second pib file. */
  0xFC08:"imgCFRW1", # kl1: 2MB fw, compressed for sta or cco, run in flash; kl2: 4M flash fw */
  0xFC09:"imgCFRW2", # The second fw, maybe the backup one. not used. */
  0xFC0A:"imgCFRW", # kl1: 2MB fw, compressed for sta or cco, run in psram; kl2: 4M psram fw */
  0xFC0B:"imgREF", # ref */
  0xFC0C:"imgCUS", # cus */
  0xFC0D:"imgCFRW4F", # kl1: 4MB fw, compressed for sta or cco, run in flash */
  0xFC0E:"imgCFRW4P", # kl1: 4MB fw, compressed for sta or cco, run in psram */
  0xFFFF:"imgMAX", # invalid */
  0x00:"imgV1ROM", # Maybe,ROM doesn't need a imgType */
  0x01:"imgV1SBL", # sbl */
  0x02:"imgV1FRWPLC", # FW-PLC */
  0x03:"imgV1CFG", # The first pib file. */
  0x05:"imgV1OEM", # oem */
  0x06:"imgV1FRWCUS", # FW-CUS */
  0xFF:"imgV1MAX"  # invalid */
}



_DevType={
  0xDC01:"devkunlun1",
  0xDC02:"devkunlun2",
  0xDC03:"devkunlun3",
  0xFFFF:"devMax",
  0x01:"devV1kunlun1",
  0x02:"devV1kunlun2",
  0x03:"devV1kunlun3",
  0xFF:"devV1Max"
}


_PKT_HEADER_MAGIC_NO = 0x00005566
_IMG_HEADER_MAGIC_NO = 0xa4e49a17

def load_flash_info(file_name:str):
  flash_info_list:list[dict]=[]
  with open(file_name,mode='r',encoding='utf-8') as f:
    lines=f.readlines()
  for line in lines:
    # 如果开头两个字符是数字
    info = re.match(r"^\d{2}",line)
    if(info):
      # print(info.group(0))
      sp=line.split()
      flash_info={}
      flash_info["DevType"]=_DevType[int(sp[1],base=16)]
      flash_info["ImgType"]=_ImgType[int(sp[2],base=16)]
      flash_info["Offset"]=int(sp[3],base=16)
      flash_info["Size"]=int(sp[4])
      flash_info_list.append(flash_info)
  return flash_info_list


# 查看是不是数据头,不是返回 "",是返回解析值
def tran_pkg_header(data:bytearray):
  size=len(data)
  ret={}
  if(size>=96):
    # 大写为无符号 B 单字节 H 双字节 I 4字节 L 4字节 Q 8字节
    # f 单精度浮点 d 双精度浮点 s 字符串
    a=struct.unpack('>BBHIIIIIIIIIIIIIIIIIIIIIIHBB',data)
    if(a[5]==_PKT_HEADER_MAGIC_NO):
      ret["enctype"]=str(a[0])
      ret["pattern"]=hex(a[1])
      ret["dev_type"]=hex(a[2])
      ret["file_crc"]=hex(a[3])
      ret["file_len"]=str(a[4])
      ret["magic_no"]=hex(a[5])
      ret["version"]=hex(a[6])
      ret["sp_start"]=hex(a[7])
      ret["sp_len"]=str(a[8])
      ret["sbl_start"]=hex(a[9]) # 32bytes
      ret["sbl_len"]=str(a[10])
      ret["oem_start"]=hex(a[11])
      ret["oem_len"]=str(a[12])
      ret["pib_start"]=hex(a[13])
      ret["pib_len"]=str(a[14])
      ret["fw_start"]=hex(a[15])
      ret["fw_len"]=str(a[16])
      ret["efuse_start"]=hex(a[17]) # 64bytes
      ret["efuse_len"]=str(a[18])
      ret["fw2_start"]=hex(a[19])
      ret["fw2_len"]=str(a[20])
      ret["ref_start"]=hex(a[21])
      ret["ref_len"]=str(a[22])
      ret["cus_start"]=hex(a[23])
      ret["cus_len"]=str(a[24])
      ret["vendor_id"]=hex(a[25])
      ret["img_type"]=hex(a[26])
      ret["img_flag"]=hex(a[27])
      t=str(ret)
      t=t.replace(',','\n')
      # t=json.dumps(ret, sort_keys=True, indent=2, separators=(',', ': '))
      return t
  return ""


def tran_img_headerv0(data:bytearray):
  size=len(data)
  ret={}
  if(size>=64):
    # 大写为无符号 B 单字节 H 双字节 I 4字节 L 4字节 Q 8字节
    # f 单精度浮点 d 双精度浮点 s 字符串
    a=struct.unpack('>HHIHBBIIIBBBBI',data[0:32])
    if(a[8]==_IMG_HEADER_MAGIC_NO):
      ret["devType"]=hex(a[0])+f"({_DevType[a[0]]})"
      ret["imgType"]=hex(a[1])+f"({_ImgType[a[1]]})"
      ret["imgSize"]=str(a[2])
      ret["imgVer"]=hex(a[3])
      ret["psramSize"]=str(a[4])
      ret["hdrVer"]=hex(a[5])
      ret["fwSize"]=str(a[6])
      ret["imgCRC"]=hex(a[7])
      ret["guard"]=hex(a[8])
      ret["layoutIdx"]=hex(a[9])
      # reserved 3bytes
      ret["fwCRC"]=hex(a[13])
      ret["sha256"]=data[32:].hex()
      t=str(ret)
      t=t.replace(',','\n')
      return t
  return ""

def tran_img_headerv1(data:bytearray):
  size=len(data)
  ret={}
  if(size>=64):
    # 大写为无符号 B 单字节 H 双字节 I 4字节 L 4字节 Q 8字节
    # f 单精度浮点 d 双精度浮点 s 字符串
    a=struct.unpack('>BBBBIBBBBIIIBBBBI',data[0:32])
    if(a[11]==_IMG_HEADER_MAGIC_NO):
      ret["devType"]=hex(a[0])+f"({_DevType[a[0]]})"
      ret["imgType"]=hex(a[1])+f"({_ImgType[a[1]]})"
      ret["encType"]=hex(a[2])
      ret["cfg"]=hex(a[3])
      ret["imgSize"]=str(a[4])
      ret["zipType"]=hex(a[5])
      ret["flashSize"]=str(a[6])
      ret["psramSize"]=hex(a[7])
      ret["hdrVer"]=hex(a[8])
      ret["imgVer"]=hex(a[9])
      ret["imgCRC"]=hex(a[10])
      ret["guard"]=hex(a[11])
      ret["layoutIdx"]=hex(a[12])
      # reserved 3bytes
      ret["runAddr"]=hex(a[16])
      ret["sha256"]=data[32:].hex()
      t=str(ret)
      t=t.replace(',','\n')
      return t
  return ""


# 判断是不是pkt_header
def pkt_header_check(data:bytearray):
  magic=(data[12]<<24)|(data[13]<<16)|(data[14]<<8)|(data[15])
  if (magic==_PKT_HEADER_MAGIC_NO):
    return True
  print(f"magic={hex(magic)}")
  return False

# 判断是不是img_header
def img_header_check(data:bytearray):
  if(len(data)<32):
    return ''
  magic=(data[20]<<24)|(data[21]<<16)|(data[22]<<8)|(data[23])
  if (magic==_IMG_HEADER_MAGIC_NO):
    imghdr_v=data[11]
    print(f"magic={hex(magic)}, imghdr_v={imghdr_v}")
    # 返回使用的结构体类型
    if(imghdr_v==0x10):
      return 'V1'
    else:
      return "V0"
  return ""


def bin_to_hex(bin:bytearray,flash_info:list[dict]=[]):
  all_size=len(bin)
  pack_size=32
  turned=0
  turned_old=0
  out_text=''
  index=0
  while turned < all_size:
    if(turned+pack_size<=all_size):
      data=bin[turned:turned+pack_size]
      turned_old=turned
      turned+=pack_size
    else:
      data=bin[turned:]
      turned=all_size
    if(index<len(flash_info)):
      if(turned_old==(flash_info[index])["Offset"]):
        out_text+=f"{(flash_info[index])['ImgType']}----------------------\n"
        index+=1
    out_text+=f"[{hex(turned_old)}] {data.hex(' ')}\n"
  return out_text

# 这个是从固件里直接读出数据头
def bin_to_hex2(bin:bytearray):
  all_size=len(bin)
  pack_size=32
  turned=0
  turned_old=0
  out_text=''
  if(pkt_header_check(bin)):
    out_text+=tran_pkg_header(bin[0:96])+'\n'
  while turned < all_size:
    if(turned+pack_size<=all_size):
      data=bin[turned:turned+pack_size]
      turned_old=turned
      turned+=pack_size
    else:
      data=bin[turned:]
      turned=all_size
    if(all_size-turned_old>=64):
      hdr_data=bin[turned_old:turned_old+64]
    else:
      hdr_data=bytearray()
    hdr_v=img_header_check(hdr_data)
    if(hdr_v=='V0'):
      out_text+=tran_img_headerv0(hdr_data)+'\n'
    elif(hdr_v=='V1'):
      out_text+=tran_img_headerv1(hdr_data)+'\n'
    out_text+=f"[{hex(turned_old)}] {data.hex(' ')}\n"
  return out_text



def bin_to_hex_file(bin_file_name:str,hex_file_name:str,info_file_name:str=None):
  with open(bin_file_name,mode='rb') as f:
    bin=f.read()
  if info_file_name is None:
    info=[]
  else:
    info=load_flash_info(info_file_name)
  text=bin_to_hex2(bin)
  with open(hex_file_name,mode='w+',encoding="utf-8") as f:
    f.write(text)

# bin_to_hex.py input_file output_file
if __name__ == "__main__":
  # bin_to_hex_file(sys.argv[1],sys.argv[2],sys.argv[3])
  # load_flash_info("work/upload-test.info")
  bin_to_hex_file(sys.argv[1],sys.argv[2])
