#!/usr/bin/env python3
# -*- coding: utf-8 -*-

#
# avrdude - A Downloader/Uploader for AVR device programmers
# Copyright (C) 2024 MX682X
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

# this is a small python skript that tries to find the PICkit5_TP folder
# created by MPLAB X that contains the scripts.xml file.
# Works on Windows and Linux (MacOS untested), as long as the default installation
# folder was not changed. If it was changed, and the program is unable to locate
# the scripts.xml file, you should be promted to enter a path. You can
# either provide the path to the file, or to the directory providing the file.
#
# The idea behind this program is to extract the sub-programs defined in the
# script.xml. The original file has a size of about 300MB, containing a lot of
# redundant information (like the XML tags) as well as sub-programs for chips 
# avrdude doesn't support, like ARM MCUs.
# This python script goes through all functions, removing identical ones,
# and indexes those. The index is then used to connect the MCUs with those functions
# so that the correct array pointers can be loaded.
# Warning: If you run this python program, the previously generated files will be
# overwritten without warning.
# 

import os, fnmatch, mmap
from pathlib import Path


# The list of functions, as a Python Dictionary, that will be used by avr-dude
c_func_list = [
    # Started with UPDI
    "EnterProgMode",
    "EnterProgModeHvSp",       # High Voltage Pulse on UPDI line
    "EnterProgModeHvSpRst",    # High Voltage Pulse on Reset Pin
    "EnterProgModeHvUpt",      # User Power Toggle
    "ExitProgMode",
    "SetSpeed",
    "GetDeviceID",
    "EraseChip",
    "WriteProgmem",
    "ReadProgmem",
    "WriteDataEEmem",
    "ReadDataEEmem",
    "WriteCSreg",
    "ReadCSreg",
    "WriteMem8",
    "ReadMem8",
    "WriteConfigmem",
    "WriteConfigmemFuse",
    "WriteConfigmemLock",
    "ReadConfigmem",
    "ReadConfigmemFuse",
    "ReadConfigmemLock",
    "WriteIDmem",
    "ReadIDmem",
    "ReadSIB",

    # Added from dW
    "switchtoISP",
    #"ReadMemIO",
    #"WriteMemIO",

    # Added from ISP
    "ReadCalibrationByte",

    # Added from JTAG/PDI
    #"WriteSRAM",   # Is a duplicate of WriteMem8
    #"ReadSRAM",
    "WriteBootMem",
    "ReadBootMem",
]

# List of MCUs Names that are not supported by avrdude
mcu_to_exclude = [
    "ATA5700M322", "ATA5702M322", "ATA5782", "ATA5787", "ATA5831", "ATA5835", "ATA8210", "ATA8510", 
    "ATtiny416auto", "AVR16DV14", "AVR16DV20"
]


import platform

work_dir = os.path.abspath(os.getcwd())
cache_dir = os.path.join(work_dir, "scripts_cache")

print(work_dir)
print(cache_dir)



# Beginning of C and H Files
common_header = \
'''\
/* This file was auto-generated by scripts_decoder.py. 
 * Any changes will be overwritten on regeneration 
 */

/*
 * avrdude - A Downloader/Uploader for AVR device programmers
 * Copyright (C) 2024 MX682X
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */

'''

# generates the h-file. generates the struct definition
def generate_h_file(c_funcs, file_dir):
    h_header = \
'''
#ifndef pickit5_lut_h
#define pickit5_lut_h

#ifdef __cplusplus
  extern "C" {
#endif

struct avr_script_lut {
'''

    h_trailer = \
'''
};

typedef struct avr_script_lut SCRIPT;
const unsigned char * get_devid_script_by_nvm_ver(unsigned char version);
int get_pickit_dw_script(SCRIPT *scr, const char *partdesc);
int get_pickit_isp_script(SCRIPT *scr, const char *partdesc);
int get_pickit_jtag_script(SCRIPT *scr, const char *partdesc);
int get_pickit_updi_script(SCRIPT *scr, const char *partdesc);
int get_pickit_pdi_script(SCRIPT *scr, const char *partdesc);
int get_pickit_tpi_script(SCRIPT *scr, const char *partdesc);

#ifdef __cplusplus
}
#endif

#endif // pickit5_lut_h
'''
    global common_header
    if file_dir is None:
        return

    h_lut_path = os.path.join(file_dir, "pickit5_lut.h")  # first - handle defining the structure
    if (os.path.exists(h_lut_path)):
        os.remove(h_lut_path)

    with open(h_lut_path, 'w') as h_file:
        h_file.write(common_header)
        h_file.write(h_header)
        for func_name in c_funcs:
            h_file.write("  const unsigned char *{0};\n         unsigned int  {0}_len;\n".format(func_name))
        h_file.write(h_trailer)
        print("h-File generated")
    #EOF




# Tries to locate the xml file in a known location
def find_xml():
    home_dir = str(Path.home())
    print("Home Path: {0}".format(home_dir))
    if home_dir == None:
        return
    home_dir = os.path.join(home_dir, ".mchp_packs", "Microchip")
    home_dir_A = os.path.join(home_dir, "PICkit5_TP")
    result = []
    for root, dirs, files in os.walk(home_dir_A):
        for name in files:
            if fnmatch.fnmatch(name, "scripts.xml"):
                file_path = os.path.join(root, name)
                result.append((os.path.getctime(file_path), file_path))

    print("List of scripts.xml files:")
    print(result)
    time, path = 0, ""
    for t, p in result:     # find the most recent scripts file in out list
        if t > time:
            time, path = t, p
    return path
    # EOF


# Example of a function declaration in the XML-File
#  <script>
#    <function>SetSpeedFromDevice</function>
#    <processor>ATSAM4LC2A</processor>
#    <ri4command>0x00001504</ri4command>
#    <scrbytes>
#      <byte>0x3B</byte>
#      <byte>0x02</byte>
#      <byte>0x00</byte>
#      <byte>0x00</byte>
#      <byte>0x00</byte>
#    </scrbytes>
#  </script>



def convert_xml(xml_path, c_funcs):
    if xml_path == None:
        print("No Path to XML file provided")
        return

    program_iface = {
        "UPDI": dict(),
        "PDI": dict(),
        "dW": dict(),
        "ISP": dict(),
        "TPI": dict(),
        "JTAG": dict()
    }
    function_dict = {
        "UPDI": dict(),
        "PDI": dict(),
        "dW": dict(),
        "ISP": dict(),
        "TPI": dict(),
        "JTAG": dict()
    }

    # Prepare directories
    parent_dir = os.getcwd()
    src_dir = os.path.join(parent_dir, "src")
    print("Opening file {0}".format(xml_path))
    print("Parent Dir: {0}".format(parent_dir))
    print("Src directory: {0}".format(src_dir))

     
    # create h-File (make sure to provide all function definitions)
    generate_h_file(c_funcs, src_dir)


    with open(xml_path, "r") as xml_script:
        print ("XML File opened")
        scr_bytes_buffer = bytearray(2048)   # allocate 2kB of memory in advance, avoids memory managment
        while True:
            line = xml_script.readline()    # go line by line, hopefully reducing memory usage compared to readlines()
            if line == "":
                break       # exit when End of file

            if line.startswith("    <function>"):
                function = line[14:]                        # remove "    <function>"
                function = function.split("<", 2)[0]        # remove trailing element

                scr_header = xml_script.readlines(3)        # read 3 more lines
                processor = scr_header[0][15:]              # remove "    <processor>"
                chip_name = processor.split("<", 2)[0]      # remove trailing element
            
                try:
                    function_name, programming_mode = function.split('_')
                except:
                    continue    # If the function did not contain any '_', continue to next line

                if programming_mode not in ["UPDI", "PDI", "dW", "ISP", "TPI", "JTAG"]:
                    continue    # Filters out "FPGA" and other edge cases

                if chip_name in mcu_to_exclude:
                    continue    # don't handle chips avrdude doesn't know anyway

                if function_name not in c_funcs:
                    continue    # Filter out debug Functions

                func_bytes = None
                counter = 0
                while True:
                    byte = xml_script.readline()
                    if (byte.startswith("      ")):         # 6 spaces is already unique enough
                        scr_bytes_buffer[counter] = int(byte[14:16], 16) # only handle the value
                        counter += 1
                    elif (byte.startswith("    </scrbytes>")):  # done with the list
                        func_bytes = bytes(scr_bytes_buffer[:counter])  # create an immutable bytes array
                        break

                if func_bytes == None or len(func_bytes) == 0:
                    continue    # continue with next chip if somethin went wrong or is empty (SetSpeed_dw)

                if function_name not in function_dict[programming_mode].keys():
                    function_dict[programming_mode][function_name] = []

                if func_bytes not in function_dict[programming_mode][function_name]:
                    function_dict[programming_mode][function_name].append(func_bytes)

                index = function_dict[programming_mode][function_name].index(func_bytes)

                #function_dict = {
                #  "UPDI": {
                #    "EnterProgMode" : [bytes_0, bytes_1]},
                #    "SetSpeed"      : [bytes_2, bytes_3]}
                #  }
                #}


                if chip_name not in program_iface[programming_mode]:
                    program_iface[programming_mode][chip_name] = [(function_name, index)]
                    #print("Added to " + programming_mode + ": " + chip_name)   # Debugging
                else:
                    program_iface[programming_mode][chip_name].append((function_name, index))
                #program_iface = {
                #    "UPDI": {
                #        "Attiny1614": [
                #            ("EnterProgMode", 0),
                #            ("ExitProgMode", 0),
                #            ("...", 1),
                #        ]
                #    }
                #}
            # /if starts with
        # /while True
    # /with open

    print("XML File processed")

    # create c-File
    global common_header

    for prog_iface, prog_mcu_list in program_iface.items():
        lower_prog_iface = prog_iface.lower()
        c_lut_path = os.path.join(src_dir, "pickit5_lut_" + lower_prog_iface + ".c")
        if (os.path.exists(c_lut_path)):
            os.remove(c_lut_path)
        with open(c_lut_path, 'w') as c_file:
            c_file.write(common_header)
            c_file.write("#include <ac_cfg.h>\n")
            c_file.write("#include <stddef.h>\n")
            c_file.write("#include <string.h>\n")
            c_file.write("#include \"pickit5_lut.h\"\n\n\n")

            struct_init_func = ""
            struct_init_len = ""
            common_func = []        # List of Functions that exist once

            for (func_name, func_array_bytes) in function_dict[prog_iface].items():
                for (array_iter, func_bytes) in enumerate(func_array_bytes):
                    func_length = len(func_bytes)
                    c_file.write("const unsigned char {0}_{1}_{2}[{3}]".format(
                        func_name, lower_prog_iface, array_iter, func_length) + " = {")
                    num_line = "  "
                    for (iter, byte) in enumerate(func_bytes):      # go through every byte
                        if (iter % 16 == 0):
                            c_file.write(num_line)                  # new line after 16 bytes
                            num_line = "\n  "
                        num_line += "0x{0:02x}, ".format(byte)      # and generate String
                    c_file.write(num_line + "\n};\n\n")             # complete array

                if len(func_array_bytes) == 1:     # look for common function
                    if (prog_iface == "JTAG"):                      # This handles the edge case in JTAG where only the
                        if (func_name == "ReadConfigmem") or (func_name == "WriteConfigmem"):
                            continue                                # XMEGA has the functions, but not the old JTAG
                    common_func.append(func_name)                   
                    struct_init_func += "  scr->{0} = {0}_{1}_0;\n".format(func_name, lower_prog_iface)
                    struct_init_len  += "  scr->{0}_len = sizeof({0}_{1}_0);\n".format(func_name, lower_prog_iface)

                # EOFL   
                    
            
            c_file.write("\n\n\nstatic void pickit_{0}_script_init(SCRIPT *scr);\n".format(lower_prog_iface))   # declaration
            c_file.write("static void pickit_{0}_script_init(SCRIPT *scr)".format(lower_prog_iface) + " {\n")   # definition
            c_file.write("  memset(scr, 0x00, sizeof(SCRIPT));  // Make sure everything is NULL\n\n")
            c_file.write(struct_init_func)
            c_file.write("\n")              # improve readability
            c_file.write(struct_init_len)
            c_file.write("}\n\n\n")

            c_file.write("const char * const pickit5_{0}_chip_lut[]".format(lower_prog_iface) + " = {")
            chip_line = "  "
            for (iter, chip_name) in enumerate(prog_mcu_list.keys()):     # go through every chip
                if (iter % 8 == 0):
                    c_file.write(chip_line)                     # new line after 8 Chips
                    chip_line = "\n  "
                chip_line += "{0:>17},".format( '"' + chip_name + '"')      # and generate String
            c_file.write(chip_line + "\n};\n\n")                # complete array


            if (prog_iface == "UPDI"):
                c_file.write("const unsigned char * get_devid_script_by_nvm_ver(unsigned char version) {\n")
                c_file.write("  if (version >= '0') version -= '0'; // allow chars\n")
                c_file.write("  if (version > 9) return NULL;       // Not a valid number\n")
                c_file.write("  if (version <= 3)                   // tiny, mega, DA, DB, DD, EA\n")
                c_file.write("    return GetDeviceID_updi_0;\n")
                c_file.write("  else                                // DU, EB\n")
                c_file.write("    return GetDeviceID_updi_1;\n}\n\n")

            c_file.write("int get_pickit_{0}_script(SCRIPT *scr, const char* partdesc)".format(lower_prog_iface) + " {\n")
            c_file.write("  if ((scr == NULL) || (partdesc == NULL)) {\n    return -1;\n  }\n")
            c_file.write("  int namepos = -1;\n")
            c_file.write("  for (int i = 0; i < {0}; i++)".format(len(prog_mcu_list.keys())) + " {\n")
            c_file.write("    if (strcmp(pickit5_{0}_chip_lut[i], partdesc) == 0)".format(lower_prog_iface) + " {\n")
            c_file.write("      namepos = i;\n      break;\n    }\n  }\n")
            c_file.write("  if (namepos == -1) {\n    return -2;\n  }\n\n")
            c_file.write("  pickit_{0}_script_init(scr);   // load common functions\n\n".format(lower_prog_iface))

            case_list = []
            case_func_list = []
            #print(common_func)
            c_file.write("  switch (namepos) {\n")
            for (switch_iterator, (chip_name, functions)) in enumerate(prog_mcu_list.items()):
                new_case = "    case {0}:  /* {1} */\n".format(switch_iterator, chip_name)
                new_func_str = ""
                
                for func_name, func_num in functions:       # generate list of unique function assignments
                    if func_name in common_func:
                        continue        # skip common functions, they were set in _init()
                    new_func_str += "{0}_{1}_{2}\n".format(func_name, lower_prog_iface, func_num)

                
                if new_func_str not in case_func_list:      # Check if there is an already existing function set
                    case_list.append(new_case)
                    case_func_list.append(new_func_str)
                else:                                       # if it exists, figure out the index at which to add the case
                    case_list[case_func_list.index(new_func_str)] += new_case
            
            for (case_str, func_list) in zip(case_list, case_func_list):
                c_file.write(case_str)
                func_list = func_list.split("\n")[:-1]      # Remove last Element that will be an empty string
                for func in func_list:
                    func_name = func.split("_")[0]
                    c_file.write("      scr->{0} = {1};\n".format(func_name, func))
                    c_file.write("      scr->{0}_len = sizeof({1});\n".format(func_name, func))
                c_file.write("      break;\n")

            c_file.write("  }\n  return namepos;\n}")
            # End of switch case
            print("finished " + prog_iface)

    print("c-File generated")



xml_path = find_xml()
if xml_path == None:
    print("Unable to find scripts.xml in the default location.")
    print("Please Enter a Path to the File or Directory:")
    xml_path = input(">")
    if (os.path.isdir(xml_path)):
        os.path.join(xml_path, "scripts.xml")
    if (os.path.exists(xml_path) == False):
        print("File not found, exiting")
        quit()
convert_xml(xml_path, c_func_list)   
quit()
