#-------------------------------------------------------------------------------
# Name:       trex_ext
# Purpose:    External module
#
# Author:     SUPER_紫电  QQ,911344755
#
# Created:    15/09/2022
# Copyright:  (c) SUPER_紫电 2022 All rights reserved.
# Licence:    <Modified BSD License>
#-------------------------------------------------------------------------------

from trex import *

module_ext = [
    'funs_ext',
    'xxx',
    'xxx',
]
for name in module_ext:
    if common.import_module(name) != None:
        print('Import module: ' + name)

def extfun():
    actions = {
        't': {'fun': pst,         'param': (),   'menu': 'PST测试模式   '},
        'f': {'fun': clr_res_cyl, 'param': (0,), 'menu': '完全清零保留区'},
        'q': {'fun': clr_res_cyl, 'param': (9,), 'menu': '快速清零保留区'},
    }
    title = '\n' + (win_size.X//2) * '='
    title+= '\n' + 'External Actions'.center(win_size.X//2,' ')
    title+= '\n' + (win_size.X//2) * '='
    actions_pool(title, actions)

def pst():
    actions = {
        't1': {'fun': seam_cal,         'param': (common.workdir,), 'menu': 'PST 0xB1 0x0160'},
        't4': {'fun': mini_cal,         'param': (common.workdir,), 'menu': 'PST 0xB1 0x04E0'},
        '44': {'fun': mini_arco_test,   'param': (common.workdir,), 'menu': 'PST 0x44       '},
        't0': {'fun': mini_wrro,        'param': (common.workdir,), 'menu': 'PST 0xB2 0x0000'},
        't2': {'fun': mini_wrro2,       'param': (common.workdir,), 'menu': 'PST 0xB2 0x0002'},
        'df': {'fun': mini_ibi,         'param': (common.workdir,), 'menu': 'PST 0xDF 0x0001'},
        'dp': {'fun': download_packets, 'param': (common.workdir,), 'menu': '下载固件包     '},
        'vm': {'fun': poll_vscstat,     'param': (),                'menu': 'VSC状态监视器  '},
        'ir': {'fun': spin_down_reset,  'param': (),                'menu': '停止电机并复位 '},

    }
    title = '\n' + (win_size.X//2) * '='
    title+= '\n' + 'PST Actions'.center(win_size.X//2,' ')
    title+= '\n' + (win_size.X//2) * '='
    actions_pool(title, actions)

def flash():
    actions = {
        '0': {'fun': flash_read,       'param': (common.workdir,), 'menu': '读取Flash    '},
        '1': {'fun': flash_file_check, 'param': (),                'menu': '检查Flash文件'},
        '2': {'fun': flash_file_read,  'param': (common.workdir,), 'menu': '读取Flash文件'},
        'w': {'fun': flash_write,      'param': (common.workdir,), 'menu': '写入Flash    '},
        'f': {'fun': flash_file_write, 'param': (common.workdir,), 'menu': '写入Flash文件'},
        'i': {'fun': init_region,      'param': (),                'menu': '初始化区域表 '},
    }
    title = '\n' + (win_size.X//2) * '='
    title+= '\n' + 'Flash Actions'.center(win_size.X//2,' ')
    title+= '\n' + (win_size.X//2) * '='
    actions_pool(title, actions)

def modules():
    actions = {
        '0': {'fun': res_read,      'param': (common.workdir,), 'menu': '读取保留区文件'},
        '1': {'fun': res_check,     'param': (),                'menu': '检查保留区文件'},
        '2': {'fun': res_read35,    'param': (common.workdir,), 'menu': '读取文件 0x35 '},
        'w': {'fun': res_write,     'param': (common.workdir,), 'menu': '写入保留区文件'},
        'o': {'fun': load_ovl_wrfl, 'param': (common.workdir,), 'menu': '加载Ovl写文件 '},
        'l': {'fun': load_ovl,      'param': (common.workdir,), 'menu': '加载Overlay   '},
        'f': {'fun': clr_res_cyl,   'param': (0,),              'menu': '完全清零保留区'},
        'q': {'fun': clr_res_cyl,   'param': (9,),              'menu': '快速清零保留区'},
        'i': {'fun': init_fw,       'param': (),                'menu': '初始化固件    '},
    }
    title = '\n' + (win_size.X//2) * '='
    title+= '\n' + 'Reserved File'.center(win_size.X//2,' ')
    title+= '\n' + (win_size.X//2) * '='
    actions_pool(title, actions)

def start():
    actions = {
        '0': {'fun': backup_firmware, 'param': (), 'menu': '备份固件          '},
        '1': {'fun': get_ident,       'param': (), 'menu': '驱动器Identify数据'},
        '2': {'fun': vscid,           'param': (), 'menu': 'vscid和工作目录   '},
        '3': {'fun': flash,           'param': (), 'menu': 'Flash操作         '},
        '4': {'fun': modules,         'param': (), 'menu': '保留区文件        '},
        'r': {'fun': spin_down_reset, 'param': (), 'menu': '停止电机并复位    '},
        'd': {'fun': servo_spin_down, 'param': (), 'menu': '私服电机停止      '},
        'u': {'fun': servo_spin_up,   'param': (), 'menu': '私服电机启动      '},
        's': {'fun': soft_reset,      'param': (), 'menu': '软复位            '},
        'a': {'fun': set_ata_port,    'param': (), 'menu': '设置ATA端口       '},
        'w': {'fun': set_workdir,     'param': (), 'menu': '设置工作目录      '},
        'v': {'fun': get_vscstat,     'param': (), 'menu': '显示VSC状态       '},
        'm': {'fun': poll_vscstat,    'param': (), 'menu': 'VSC状态监视器     '},
        'e': {'fun': extfun,          'param': (), 'menu': '运行外部函数      '},
    }
    title = '\n' + (win_size.X-1) * '='
    title+= '\n' + 'Task File Register Execute Tool'.center(win_size.X-1,' ')
    title+= '\n' + (win_size.X-1) * '='
    actions_pool(title, actions)

def beep():
    freq = 1193180000 / 440000 # 440Hz
    outp(0x43, 0xB6)
    outp(0x42, int(freq) & 0xFF)
    outp(0x42, int(freq) >> 9)
    sleep(0.1)
    outp(0x61, inp(0x61) | 0x03)
    sleep(1)
    outp(0x61, inp(0x61) & 0xFC)

def poll_vscstat():
    if not get_vsc_status():
        print('Error executing command')
        return
    drv_state1 = {
        0: 'Idle',
        1: 'Standby',
        2: 'Sleep',
        3: 'Background DST',
        4: 'Background Data Lifeguard',
        5: 'Background VSC Command',
        6: 'PST test in progress',
    }
    vscstatus = vsctype.vsc_status
    cur_state = vscstatus.CurState
    print('Current Drive State 1 = %d %s' % (cur_state, drv_state1[cur_state]))
    if cur_state != 6:
        return
    start = time()
    while cur_state == 6:
        cls()
        runtime = time()-start
        if runtime//3600>24:
            runtime = '%d-%02d:%02d:%02d' % (runtime//3600//24,runtime//3600%24,runtime%3600//60,runtime%60)
        else:
            runtime = '%02d:%02d:%02d' % (runtime//3600, runtime%3600//60,runtime%60)
        result  = 'PST test in progress - %s' % runtime
        result += '\n----------------------------------------------------------'
        result += '\nPTM File Id              = 0x%02X' % vscstatus.PTMId
        result += '\nPTM Test Id              = 0x%02X' % vscstatus.TestId
        result += '\nCurrent VSC Command      = 0x%02X' % vscstatus.ActionCode
        result += '\nExtended Error Status    = 0x%04X' % vscstatus.ExtErr
        result += '\nExtended Error Code      = 0x%04X' % vscstatus.ExtErrCode
        if vscstatus.PTMId == 0xB9 or vscstatus.PTMId == 0xBA:
            buf = vscstatus.CurLBA
            lba = 0
            for i in range(6):
                lba += buf[i]<<(i*8)
            result += '\nCurrent LBA              = %d' % lba
            progress = vscstatus.BGProgress
            result += '\nProgress Indicator       =  %d.%02d%%' % (progress//10, progress%10)
        result += '\nCurrent Virtual Cylinder = %d' % vscstatus.CurVirCyl
        result += '\nCurrent Virtual Head     = %d' % vscstatus.CurVirHd
        result += '\nCurrent Data Zone        = %d' % vscstatus.CurZone
        dtemp = vscstatus.DTemp
        ctemp = (dtemp-32)//1.8
        result += '\nDrive Temperature (avg)  = %d (%d)' % (dtemp, ctemp)
        result += '\n----------------------------------------------------------'
        if not get_vsc_status(): break
        vscstatus = vsctype.vsc_status
        cur_state = vscstatus.CurState
        ptm_id = vscstatus.PTMId
        print(result)
        if ptm_id == 0xFF22 or ptm_id == 0xFFF1:
            break
        sleep(0.5)
    if vsctype.ext_err == 0:
        print('Completed Successfully')
    else:
        errinfo(dec_err_code())
    beep()

def load_ovl(dir_name):
    filename = dir_name + '\\11.rpm'
    if not os.path.exists(filename):
        errinfo('No File: %s' % filename)
        return False
    actions = [
        {'file': '10', 'fun': ld10, 'des': 'Load Cache Overlay(0x10)'},
        {'file': '5c', 'fun': ld5c, 'des': 'Load Servo Overlay(0x5C)'},
        {'file': '11', 'fun': ld11, 'des': 'Load Permanent Overlay(0x11)'},
        {'file': '13', 'fun': ld13, 'des': 'Load Cache Overlay(0x13)'},
        {'file': '80', 'fun': ld80, 'des': 'Load Overlay File(0x80)'},
        {'file': '1b', 'fun': ld1b, 'des': 'Load Transient Overlay(0x1B)'},
    ]
    result = False
    for i in range(len(actions)):
        file_id = actions[i]['file']
        if os.path.exists(dir_name + '\\' + file_id + '.rpm'):
            print(actions[i]['des'] + ' - ', end='')
            result = actions[i]['fun'](dir_name)
            if result:
                print('%s' % dec_err_code())
            else:
                get_err_code()
                errinfo('%s' % dec_err_code())
                break
    return result

def clr_res_cyl(blk=0):
    if not load_ovl(common.workdir):
        return
    vir_hd_cnt = get_vir_hd()
    if vir_hd_cnt < 1:
        return
    #pys_zone_tbl_hdr, pys_zone_tbl = get_phys_zone_tbl(0)
    pys_zone_tbl = get_phys_zone_tbl(0)[1]
    if pys_zone_tbl == []:
        return
    first_cyl_zn = pys_zone_tbl[0].first_cyl_zn
    last_cyl_zn  = pys_zone_tbl[0].last_cyl_zn if blk == 0 else blk
    spt_hd0_zn   = pys_zone_tbl[0].spt_hd0_zn
    for c in range(last_cyl_zn, first_cyl_zn-1, -1):
        clrcyl(c, 0)
    if vir_hd_cnt == 1:
        return
    for c in range(last_cyl_zn, first_cyl_zn-1, -1):
        clrcyl(c, 1)

def mini_test(file_path, file_id, test_id):
    if not load_ovl(file_path):
        return False
    if not wbx_pst(file_path, file_id, 0):
        return False
    if xpst(file_id, test_id) != 0:
        return False
    return rstat_bitclr(0)

def mini_cal(file_path):
    if mini_test(file_path, 0xB1, 0x4E0):
        poll_vscstat()
    else:
        get_err_code()
        errinfo('%s' % dec_err_code())

def seam_cal(file_path):
    if mini_test(file_path, 0xB1, 0x160):
        poll_vscstat()
    else:
        get_err_code()
        errinfo('%s' % dec_err_code())

def mini_wrro(file_path, test_id=0):
    if mini_test(file_path, 0xB2, test_id):
        poll_vscstat()
    else:
        get_err_code()
        errinfo('%s' % dec_err_code())

def mini_wrro2(file_path):
    if mini_wrro(file_path, 2):
        poll_vscstat()
    else:
        get_err_code()
        errinfo('%s' % dec_err_code())

def mini_ibi(file_path):
    drv_data_tbl = Drv_Data_Tbl()
    if not get_drv_tbl_data(drv_data_tbl):
        return
    if mini_test(file_path, 0xDF, 1):
        if drv_data_tbl.ProdID == 0x1694: # VIKING
            wait_not_busy(drv_data_tbl.VirHdCnt*60)
        else:
            poll_vscstat()
    else:
        get_err_code()
        errinfo('%s' % dec_err_code())
    #return mini_test(file_path, 0xDF, 1)

def mini_arco(file_path , head_dcm, media_dcm):
    if not load_ovl(file_path):
        return False
    file_id = 0x2401 if os.path.exists(file_path + '\\2401.rpm') else 0xC4
    if not wbx_pst(file_path, file_id, 0):
        return False
    filename = file_path + '\\C5.rpm'
    if not ldbin(filename):
        return False
    b = vsctype.super_buffer[0x24]
    start_offset = 512
    if b == 0: start_offset = 256
    if b == 2: start_offset = 512
    if b == 3: start_offset = 640
    if b == 4:
        b = vsctype.super_buffer[0x23]
        if b == 0x00: start_offset = 0x300
        if b == 0x50: start_offset = 808
        if b == 0xD0: start_offset = 872
        if b == 0xE0: start_offset = 880
        if b == 0xF0: start_offset = 888
        if b == 0xF8: start_offset = 892
    if b == 5:
        b = vsctype.super_buffer[0x23]
        if b == 0x84: start_offset = 962
        if b == 0xC0: start_offset = 992
        if b == 0xE0: start_offset = 1008
    if not wbx_pst(file_path, 0xC5, start_offset):
        return False
    if xpst(file_id, 1, [0x44, 0x19, head_dcm, media_dcm]) != 0:
        return False
    return rstat_bitclr(0)

def mini_arco_test(file_path):
    head_dcm, media_dcm = 0, 0
    if rdfl(0xA):
        hdr_size  = vsctype.super_buffer[6]
        head_dcm  = vsctype.super_buffer[hdr_size+13]
        media_dcm = vsctype.super_buffer[hdr_size+12]
    if head_dcm == 0 or media_dcm == 0:
        if ldbin(file_path + '\\28.rpm'):
            head_dcm  = vsctype.super_buffer[0x44]
            media_dcm = vsctype.super_buffer[0x48]
        else:
            dcm_code  = input('输入磁头和盘体DCM代码(0-9,A-Z)(例如:UZ):')
            head_dcm  = ord(dcm_code[0].upper())
            media_dcm = ord(dcm_code[1].upper())
    print('Head  DCM: 0x%02X(%c) \nMedia DCM: 0x%02X(%c)' % (head_dcm,head_dcm,media_dcm,media_dcm))
    if mini_arco(file_path, head_dcm, media_dcm):
        poll_vscstat()
    else:
        get_err_code()
        errinfo('%s' % dec_err_code())

# Download Microcode (PTM)
def download_ptm(filename):
    if not ldbin(filename):
        return False
    file_size = len(vsctype.super_buffer)
    if init_download_ptm(0x108, file_size) != 0:
        return False
    if not smart_write_data(vsctype.super_buffer, 0, file_size//512):
        return False
    return rstat_bitclr(0)

def download_packets(file_path):
    cf1 = ''
    packet_type = ['cf', 'ce', 'ib', 'rc', 'rd', 'se', 'sl', '115']
    packet_list = []
    for file in os.listdir(file_path):
        if file.endswith('.cf1'):
            cf1 = file
        else:
            if file[-3:-1] in packet_type or file.endswith('.115'):
                packet_list.append(file)
    if cf1 == '':
        return False
    if not download_ptm(cf1):
        return False
    if not spindown_reset():
        return False
    if not wait_not_busy(30):
        return False
    #init2(False)
    zero_size(2)
    if not load_ovl(file_path):
        return False
    for packet_name in packet_list:
        if not download_ptm(packet_name):
            return False
    #init2(True)
    fix_size(2)
    return rstat_bitclr(0)