#!/usr/bin/env python3

"""
rom.py: Module that handles hex I/O to files and rom management
"""

import os
import mmap
import pathlib
import binascii
from pktools.lib.messages import PKTMessage as Messages

class Rom:
    """Class Rom represents and manipulates rom files."""

    def __init__(self, rfile, fbyte = b'\00'):
        """Constructor for Rom class. Initializes file to be used by class.

        Args:
            rfile (str): Path to the rom file.
            fbyte (byte): The empty byte value

        Raises:
            ValueError: If given path is not a valid file
        """
        if not pathlib.Path(rfile).is_file():
            raise ValueError(Messages.FIOE001.print(rfile))
        self.romfile = rfile

        self.free_byte = fbyte

    def size(self):
        """Retrieves the file size of the rom, in bytes.

        Returns:
            int: Size of the file in bytes
        """
        return os.path.getsize(self.romfile)

    def kb_size(self):
        """Retrieves the file size of the rom, in kilobytes.

        Returns:
            int: Size of the file in kilobytes
        """
        return os.path.getsize(self.romfile) / 1024

    def read(self, nbytes, offset=0):
        """Retrieves a list of bytes from the rom.

        Args:
            nbytes (int): The number of bytes to read
            offset (int): The offset at which to start reading values

        Returns:
            list: A list of bytes
        """
        with open(self.romfile, 'rb') as f:
            f.seek(offset)
            data = f.read(nbytes)
        return data

    def read_hex(self, nbytes, offset=0):
        """Retrieves a list of bytes from the rom in hex form.

        Args:
            nbytes (int): The number of bytes to read
            offset (int): The offset at which to start reading values

        Returns:
            list: A list of bytes
        """
        with open(self.romfile, 'rb') as f:
            f.seek(offset)
            hexdata = ["{:02x}".format(a) for a in f.read(nbytes)]
        return hexdata

    def write(self, data, offset=0, overflow=False):
        """Writes a list of bytes to the rom. Overwrites values.

        Args:
            data (list): The bytes to write
            offset (int): The offset at which to start writing values

        Raises:
            ValueError: If data to write exceeds the length of the file
        """
        if (offset + len(data)) > self.size():
            raise ValueError(Messages.FIOE002.print())

        with open(self.romfile, 'r+b') as f:
            f.seek(offset)
            f.write(data)

    def write_hex(self, hexlist, offset=0, truncate=False):
        """Writes a list of bytes to the rom in hex. Overwrites values.

        Args:
            hexlist (list): The hex values to write
            offset (int): The offset at which to start writing values

        Raises:
            ValueError: If data to write exceeds the length of the file
        """
        data = binascii.unhexlify(''.join(hexlist))

        if (offset + len(data)) > self.size():
            raise ValueError(Messages.FIOE002.print())

        with open(self.romfile, 'r+b') as f:
            f.seek(offset)
            f.write(data)

    def find(self, data, offset=0):
        """Search for a bytearray in the rom.

        Args:
            data (list): The bytearray to search for
            offset (int): The offset at which to start searching for values

        Returns:
            int: The location of the first occurrence after the offset
        """
        with open(self.romfile, 'rb') as f:
            mm = mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ)
            mm.seek(offset)
            found = mm.find(data)
            if found != -1:
                return found

    def findall(self, data):
        """Search for all occurrences of a bytearray in the rom.

        Args:
            data (list): The bytearray to search for

        Returns:
            int: A list of locations to the bytearray
        """
        ret =[]
        offset = 0

        with open(self.romfile, 'rb') as f:
            mm = mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ)
            while True:
                mm.seek(offset)
                found = mm.find(data)
                if found == -1:
                    break
                ret.append(found)
                offset = found + 1

        return ret

    def find_freespace(self, numbytes, offset=0):
        """Iterates through the freespace blocks required in a rom

        Args:
            numbytes (int): The number of free bytes to look for
            offset (int): The location at which to start looking for

        Yields:
            int: The next valid location for a freespace block
        """
        cur_offset = offset
        freespace_block = self.free_byte * numbytes

        with open(self.romfile, 'rb') as f:
            mm = mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ)
            while cur_offset + numbytes <= self.size():
                mm.seek(cur_offset)
                found = mm.find(freespace_block)
                if found != -1:
                    yield found
                cur_offset = found + 1
                if (cur_offset + numbytes) >= self.size():
                    break;
                mm.seek(cur_offset + numbytes)
                while mm.read_byte() == self.free_byte:
                    yield cur_offset
                    cur_offset = cur_offset + 1
                    if (cur_offset + numbytes) >= self.size():
                        break;
