import os
import sys
import shutil
import json



def get_path(dir_item:str,sub:str):
  with open('ohos_config.json') as f:
    j=json.loads(f.read())
    return os.path.join(j[dir_item],sub)


_kunlun_path="/mnt/disk/work/kunlun/iotelic/kunlun/Mainline/tools/gen_img"
_harmony_path=get_path("out_path","")


def copy_cco():
  global _kunlun_path
  _kunlun_path=_kunlun_path+'_cco'
  shutil.copy(f"{_harmony_path}/OHOS_Image.bin",
              f"{_kunlun_path}/iot_ht_3p.bin")

  shutil.copy(f"{_kunlun_path}/tmp_plc_cco.lds",
              f"{_kunlun_path}/tmp_plc.lds")


  _customer_path=f"{_kunlun_path}/customer/4M"
  if(os.path.exists(_customer_path)):
    shutil.rmtree(_customer_path)

  _cmd=f"{_kunlun_path}/gen_img_cco.sh  -d HZ -c htzd3201cco03 -i cco -r psram -a sg -v debug -m lzma -s 4m -b 3 -p 8 -t plc"
  with os.popen(_cmd) as f:
    ret=f.readlines()
  if(os.path.exists(_customer_path)):
    print("cco create success.")
    return True
  else:
    print("cco create failed.")
    return False


def copy_sta(sta_model:str):
  global _kunlun_path
  _kunlun_path=_kunlun_path+'_sta'
  shutil.copy(f"{_harmony_path}/OHOS_Image.bin",
              f"{_kunlun_path}/../../ap/.output/plc/lib/ht_sta.bin")

  shutil.copy(f"{_kunlun_path}/tmp_plc_sta.lds",
              f"{_kunlun_path}/tmp_plc.lds")

  _customer_path=f"{_kunlun_path}/customer/4M"
  if(os.path.exists(_customer_path)):
    shutil.rmtree(_customer_path)

  cmd_dict={
    "sta":f"{_kunlun_path}/gen_img_sta.sh -d HZ -c htzd3211sta03 -i sta -r flash -a sg -v debug -m lzma -s 4m -b 3 -p 0 -t plc ",
    "3ps":f"{_kunlun_path}/gen_img_sta.sh -d HZ -c htzd32113ps03 -i sta -r flash -a sg -v debug -m lzma -s 4m -b 3 -p 0 -t plc ",
    "iic":f"{_kunlun_path}/gen_img_sta.sh  -d HZ -c htzd3211iic03 -i sta -r flash -a sg -v debug -m lzma -s 4m -b 3 -p 0 -t plc ",
  }
  _cmd=cmd_dict[sta_model]
  with os.popen(_cmd) as f:
    ret=f.readlines()
  # os.system(_cmd)
  if(os.path.exists(_customer_path)):
    print(f"{sta_model} create success.")
    return True
  else:
    print(f"{sta_model} create failed.")
    return False




def main(par:str):
  if(par in ["cco", "sta", "3ps", "iic"]):
    ack=False
    if(par=="cco"):
      ack=copy_cco()
      src_type='cco'
    else:
      ack=copy_sta(par)
      src_type='sta'
    os.system(f"rm {_harmony_path}/bin/HZ-*.bin")
    if(ack):
      # 把打包好的镜像拷贝回来
      src_path=f"{_kunlun_path}/customer/4M/htzd/{src_type}/bin"
      bin_list=os.listdir(src_path)
      for item in bin_list:
        src=os.path.join(src_path,item)
        dst=f"{_harmony_path}/bin/{item}"
        # print(f"copy {src} to {dst}")
        shutil.copy(src,dst)
        # 只需要复制其中一个
        print("copy end.")
        break

  else:
    print("error.")




if __name__ == "__main__":
  if(len(sys.argv)>=2):
    main(sys.argv[1])
  else:
    print("please set par.")