# -*- coding: utf-8 -*-
"""
    Implementation of Flash Programmer service

Copyright 2018-2022, 2024-2025 NXP.

SPDX-License-Identifier: BSD-3-Clause

"""
import gdb
from flash.scripts.elfparser import SectionHeader
from errors import Result, Ok, Error, ErrorCode, safecall, exceptcall
import zlib
import struct

class FlashAppletBuffer(object):
    def __init__(self, prop_bag):
        self.address = prop_bag['ws_address']
        self.size = prop_bag['ws_size']

class FlashSector(object):
    """Flash Sector: represents sector inside flash. 
    Implements read-modify-write-verify strategy"""
    def __init__(self, start, end, number, verbose=False, nt_limit=0x1000, region="program"):
        self.start = start
        self.end = end
        self.size = end - start
        self.data = None
        self.modified = False
        self.sector_num = number
        self.compare_mask = None
        self.verbose = verbose
        self.nt_limit = nt_limit
        self.region = region

    def merge(self, addr, size, data, f_base):
        """
        This function will apply data that fit in to the sector
        :param addr: int
        :param data: []
        """

        if (addr + size <= self.start) or (addr >= self.end):  # nothing to do
            return
            
        if self.data is None:
            self.data = bytearray([0xFF]*self.size)
            self.compare_mask = bytearray([0x0]*self.size)
        #we need two indexes what to get and were to write
        #in offsets from the beginning
        src_ofst = 0 if addr >= self.start else self.start - addr
        dst_ofst = 0 if addr <= self.start else addr - self.start
        copy_size = min(addr + size, self.end) - max(addr, self.start)
        if copy_size <= 0:
            Result(Error(ErrorCode.FlashDataError, "F-ERROR: Can't prepare data from section address=0x%X size=0x%X" % (addr, size))).unwrap()
            

        #put all data into the one array
        d1 = self.data[:dst_ofst]
        d2 = data[src_ofst:src_ofst + copy_size]
        d3 = self.data[dst_ofst + copy_size:]
        
        self.compare_mask[dst_ofst:dst_ofst + copy_size] = bytearray([0x1]*copy_size)
        self.data = d1 + d2 + d3
        self.modified = True
        sector_data = self.data
        lst = [x for x,y,z in zip(self.data, sector_data, self.compare_mask) if z != 0 and x != y ]
        self.data = sector_data


class FlashDesc(object):
    """
        FlashDescriptor: Container for FlashSector objects and contain information about flash.
    """
    def __init__(self, prop_bag, cmd_interface, verbose=False, nt_limit=0x20000, restricted_addressess=None, is_crc_implemented=False):
        """
        :param prop_bag: flash desc
        :type prop_bag:dict
        """
        self.sector_count = int(prop_bag['sector'], 16)
        self.sector_size = int(prop_bag['sector_size'], 16)
        self.base_address = int(prop_bag['address'], 16)
        self.size = int(prop_bag['size'], 16)
        self.flash_cmd = cmd_interface
        self.sectors = []
        self.program_faddress = int(prop_bag['program_flash_address'], 16)
        self.data_faddress = int(prop_bag['data_flash_address'], 16)
        self.data_fsectors = int(prop_bag['data_sectors'], 16)
        self.data_fsectorssize = int(prop_bag['data_sect_size'], 16)
        self.max_byte_buffer= int(prop_bag['max_byte_buffer'][2:])
        self.verbose = verbose
        self.restricted_zones = None
        self.is_crc_implemented = is_crc_implemented

        self.restrictedAddr = restricted_addressess
        if self.restrictedAddr != None: # Check if restricted feature    
            restricted_addr_int = [int(addr, 16) for addr in self.restrictedAddr]
            self.restricted_zones = self.get_restricted_zones(restricted_addr_int)

        self.restrictedAddr = restricted_addressess
        if(prop_bag['data_flash_address'] == "0xdeadbeef"):
            self.append_sectors(self.base_address, "program", nt_limit)
        else:
            self.append_sectors(self.program_faddress, "program", nt_limit)
            self.append_sectors(self.data_faddress, "data", nt_limit)

    def get_restricted_zones(self, restricted_addressess):
        return [(restricted_addressess[i], restricted_addressess[i+1]) for i in range(0,len(restricted_addressess),2)]
    
    def append_sectors(self, flash_start_addr, region_type, nt_limit):
        i = 0
        total_size = 0
        sector_size = 0
        start_addr = 0
        end_addr = 0
        if region_type == "program":
            total_size = self.size
            sector_size = self.sector_size
        elif region_type == "data":
            total_size = self.data_fsectors * self.data_fsectorssize
            sector_size = self.data_fsectorssize

        while end_addr < flash_start_addr + total_size:
            start_addr = flash_start_addr + (i * sector_size)
            end_addr = flash_start_addr + ((i + 1) * sector_size)
            if end_addr > flash_start_addr + total_size:
                end_addr = flash_start_addr + total_size
            self.sectors.append(FlashSector(start_addr, end_addr, i, self.verbose, nt_limit, region_type))
            i += 1



class FlashWriter(object):
    def __init__(self, flash_desc, is_elf):
        self.desc = flash_desc
        self._offset = None
        self._data = None
        self._total_size = 0
        self._sum = 0
        self.elf_file = is_elf

    def load_data(self, offset=None, data=None, size=None, section=None, f_base=None):
        """
        Loads Data into class attribute
        Does handling of offset, data, and size with verbose logging and restricted zone checks
        Needs to be called before execute
        """
        if self.elf_file:
            #extract data
            section_size = section.get_size()
            section_addr = section.get_offset(f_base) + self.desc.base_address

            if self.desc.verbose:
                    gdb.write("found section at address=0x%X, size=%s\n" % (section_addr, section_size))

            data = section.section.read_data_chunk(0, section_size) #take a data from elf
            for sector in self.desc.sectors :
                addr = section_addr + self.desc.base_address
                sector.merge(addr, section_size, data, self.desc.base_address)  
            self._data = data
            self._offset = section_addr
        else:
            if size is None:
                size = len(data)
            if data == None:
                Result(Error(ErrorCode.FlashWriteError, "F-ERROR: Cannot load data. Aborting Command")).unwrap()

            if size > len(data):
                # limit to data length if requested size exceeds available data
                self._data = data
                if self.desc.verbose:
                    gdb.write("Size set larger than data length")
            elif size < len(data):
                # Ignore some data if size is set by user is smaller
                self._data = data[:size] 
            else:
                self._data = data

            self._total_size = size
            self._offset = offset
            if self.desc.restricted_zones !=None:
                if self.check_restricted_zones(offset, size, self.desc.restricted_zones):
                    Result(Error(ErrorCode.FlashWriteError, "F-ERROR: Restricted area detected. Aborting Command"))
            for sector in self.desc.sectors:
                sector.merge(offset, size, self._data, self.desc.base_address)

    def get_modified_sectors(self):
        sectors = []
        [sectors.append(idx) for idx in self.desc.sectors if idx.data != None]
        return sectors

    def batch_sectors(self, sectors, PackageSize):
        sectors.sort(key=lambda s: s.start)
        batches = []
        current_batch = []
        current_size = 0
        expected_next_start = None

        for sec in sectors:
            if expected_next_start is None:
                current_batch = [sec]
                current_size = sec.size
                expected_next_start = sec.end
            elif sec.start == expected_next_start and current_size + sec.size <= PackageSize:
                current_batch.append(sec)
                current_size += sec.size
                expected_next_start = sec.end
            else:
                combined_data = bytearray()
                for s in current_batch:
                    combined_data.extend(s.data)
                batches.append((current_batch[0].start, combined_data))
                current_batch = [sec]
                current_size = sec.size
                expected_next_start = sec.end
        if current_batch:
            combined_data = bytearray()
            for s in current_batch:
                combined_data.extend(s.data)
            batches.append((current_batch[0].start, combined_data))

        return batches
    
    def group_sectors_by_region_and_buffer(self, sectors):
        """Group modified sectors by region and buffer size using sector.region."""
        # Filter modified sectors with valid data
        sectors = self.get_modified_sectors()
        if len(sectors) == 0:
            Result(Error(ErrorCode.FlashWriteError, "F-ERROR: No modified sectors detected, please check offset. Aborting Command")).unwrap()
            return
        
        PackageSize = int(self.desc.max_byte_buffer) 
        # Separate by region using sector.region
        program_sectors = [s for s in sectors if s.region == "program"]
        data_sectors = [s for s in sectors if s.region == "data"]

        return {
            "program_flash": self.batch_sectors(program_sectors, PackageSize),
            "data_flash": self.batch_sectors(data_sectors, PackageSize)
        }
    
    def execute(self, erase=False, verify=False, is_MRAM=False):
        """
        Executes monitor commands for writing data to flash memory with optional erase and verification
        Checks if desired flash memory algorithms has CRC verify implemented
        """
        
        if self._data is None or self._offset is None:
            Result(Error(ErrorCode.FlashWriteError, "F-ERROR: No data loaded, call load_data() first. Aborting Command")).unwrap()

        if is_MRAM:
            if self.desc.is_crc_implemented:
                self._execute_with_CRC(erase, verify)
            else:
                self._execute_with_READ(erase, verify)
        else:
                self._execute_with_sector_batches(erase, verify)


    def _execute_with_sector_batches(self, erase=False, verify=False):
        
        batches = self.group_sectors_by_region_and_buffer(self.desc.sectors)
        
        for region, batch_list in batches.items():
            for addr, data in batch_list:
                
                skip_rewrite = 0
                size = len(data)
                addr = addr - self.desc.base_address
                if not self.elf_file:
                    self._sum += size

                if self.desc.is_crc_implemented:
                    skip_rewrite = self._check_CRC_before_write(addr, data, erase)
                else:
                    read_buff = 0
                    read_buff = self.desc.flash_cmd.read(addr, size)
                    skip_rewrite = self._check_buffers_before_write(data, read_buff, erase)

                self._info_print(skip_rewrite, addr, erase)

                if skip_rewrite == 1: # Same data detected, skip erase, skip write, continue to next batch
                    continue
                
                # Step 1: erase if needed
                if skip_rewrite == 2: # Data detected, if erase flag is not set, continue to next batch
                    if not erase:
                        continue
                    self.desc.flash_cmd.erase(addr, size)
                    
                # Step 2: write                
                self._progress_print(addr)
                self.desc.flash_cmd.write(addr, data)
                # Step 3: verify
                if verify:
                    if self.desc.is_crc_implemented:
                        crc_buf = zlib.crc32(data) & 0xFFFFFFFF
                        target_crc = self.desc.flash_cmd.verify(addr, size, crc_buf)
                        if self.desc.verbose:
                            target_crc_int = struct.unpack('<I', target_crc)[0]
                            gdb.write("CRC verify was done on target... new target_crc = 0x{:X}\n".format(target_crc_int))
                    else:
                        read_buff = 0
                        read_buff = self.desc.flash_cmd.read(addr, size)
                        if data != read_buff:
                            Result(Error(ErrorCode.FlashVerifyError, "F-ERROR: Verification failed. Data mismatch."))
                            return

    def _execute_with_CRC(self, erase=False, verify=False):

        if self.desc.verbose:
                gdb.write("CRC verify detected. Executing write with CRC verify...\n")

        total_len = len(self._data)
        PackageSize = int(self.desc.max_byte_buffer) 
        self._offset = self._offset - self.desc.base_address

        skip_rewrite = self._check_CRC_before_write(self._offset, self._data, erase)

        self._info_print(skip_rewrite, self._offset, erase)

        if skip_rewrite == 1:
            return
        
        # Step 1: erase if needed
        if skip_rewrite == 2: # Data detected, if erase flag is not set, continue to next batch
            if not erase:
                return
            self.desc.flash_cmd.erase(self._offset, total_len)
        
        # Step 2: write
        index = 0
        while index < total_len:
            take = min(PackageSize, total_len - index)
            self._sum += take
            self._progress_print(self._offset + index)
            self.desc.flash_cmd.write(self._offset + index, self._data[index:index+take])
            index += take

        # Step 3: verify
        if verify and total_len > 0:
            crc_buf = zlib.crc32(self._data) & 0xFFFFFFFF
            target_crc = self.desc.flash_cmd.verify(self._offset, total_len, crc_buf)
            if self.desc.verbose:
                target_crc_int = struct.unpack('<I', target_crc)[0]
                gdb.write("CRC verify was done on target... new target_crc = 0x{:X}\n".format(target_crc_int))

    def _execute_with_READ(self, erase=False, verify=False):

        if self.desc.verbose:
                gdb.write("CRC verify not detected. Executing write with read back verify...\n")
        
        total_len = len(self._data)
        PackageSize = int(self.desc.max_byte_buffer) 
        index = 0
        self._offset = self._offset - self.desc.base_address

        if erase and total_len > 0:
            if self.desc.verbose:
                gdb.write("Erase argument detected. Executing erase before write...\n")
            self.desc.flash_cmd.erase(self._offset, total_len) #TODO: not sure if total len can be erased

        while index < total_len:
            
            take = min(PackageSize, total_len - index)
            if self.desc.verbose:
                gdb.write("Packet Size = 0x{:X}\n".format(take))

            # Step 1: read data (overwrite protection)
            if erase: # erase done above, no need to check for data existing
                pass
            else:  
                read_buff = 0
                read_buff = self.desc.flash_cmd.read(self._offset + index, take)
                all_FF = read_buff == b"\xFF" * len(read_buff)
                if not all_FF:
                    gdb.write("F-Info: Data detected. Skip write operation.")
                    return
                
            # Step 2: write
            self._progress_print(self._offset + index)
            self._sum += take
            self.desc.flash_cmd.write(self._offset + index, self._data[index:index+take])

            # Step 3: verify
            if verify:
                read_buff = 0
                read_buff = self.desc.flash_cmd.read(self._offset + index, take)
                if self._data[index:index+take] != read_buff:
                    Result(Error(ErrorCode.FlashVerifyError, "F-ERROR: Verification failed. Data mismatch."))
                    return
            index += take

    '''
        return values for check:
        0 - ALL FF's, skip erase but do write
        1 - Same Data, skip erase, skip write
        2 - Data detected, check if erase flag, erase and write
    '''
    def _check_buffers_before_write(self, write_bytearray, read_bytearray, erase):

        if all(byte == 0xFF for byte in read_bytearray):
            return 0
        if write_bytearray == read_bytearray:
            return 1
        return 2

    
    def _check_CRC_before_write(self, offset, data, erase):
        
        if data is None:
            return 0
        
        if self.desc.verbose:
            gdb.write("Checking CRC before starting write commands...\n")
        
        local_crc = zlib.crc32(data) & 0xFFFFFFFF
        if self.desc.verbose:
                gdb.write("Local CRC calculated using zlib.crc32 (ITU-T V.42)... local_crc = 0x{:X}\n".format(local_crc))

        target_crc = None
        try:
            target_crc = self.desc.flash_cmd.verify(offset, len(data), 0xFFFFFFFF)
        except:
            return 1
     
        if target_crc is None:
            gdb.write("F-Warning: No returned target CRC. Skip write command.\n")
            return 1
        
        target_crc_int = struct.unpack('<I', target_crc)[0]

        all_FF = bytearray([0xFF] * len(data))
        all_FF_crc = zlib.crc32(all_FF) & 0xFFFFFFFF

        if self.desc.verbose:
                gdb.write("Target CRC calculated using C algorithm (ITU-T V.42)... target_crc = 0x{:X}\n".format(target_crc_int))

        if target_crc_int == local_crc:
            return 1

        if target_crc_int != all_FF_crc:
            return 2

        return 0

    def in_restricted_zone(self, addr, zone):
        return zone[0] <= addr <= zone[1]


    def check_restricted_zones(self, offset, size, zones):
        end_addr = offset + size
        for z in zones:
            if self.in_restricted_zone(offset, z):
                return True
            elif self.in_restricted_zone(end_addr, z):
                return True
        return False 
    
    def _progress_print(self, addr):
        if self._sum > self._total_size:
            self._sum = self._total_size
        if self.elf_file:
            gdb.write("F-WRITE: .elf sections at ADDRESS: " + str(hex(addr))  + "\n")
        else:
            percentage = (self._sum / self._total_size) * 100
            gdb.write("F-WRITE:" + str(self._sum) +"/" + str(self._total_size) +" BYTES | " + str("%.2f" % percentage) +"\n")

    def _info_print(self, skip_rewrite, addr = 0, erase = False):
        if self._sum > self._total_size:
            self._sum = self._total_size
        if self.elf_file:
            if skip_rewrite == 0:
                return
            elif skip_rewrite == 1:
                gdb.write("F-Info: Same Data detected. Skip .elf sections at ADDRESS: " + str(hex(addr))  + "\n")
            elif skip_rewrite == 2:
                if not erase:
                    gdb.write("F-Info: Data detected. Skip .elf sections at ADDRESS: " + str(hex(addr))  + "\n")
        else:
            if skip_rewrite == 0:
                return
            elif skip_rewrite == 1:
                gdb.write("F-Info: Same Data detected. Skip " + str(self._sum) +"/" + str(self._total_size) + " BYTES\n")
            elif skip_rewrite == 2:
                if not erase:
                    gdb.write("F-Info: Data detected. Skip " + str(self._sum) +"/" + str(self._total_size) + " BYTES\n")