#!/usr/bin/env python3

# -*- coding: utf-8 -*-



# Copyright (c) 2021 Huawei Device Co., Ltd.

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import bisect

import copy

import os

import struct

import tempfile

from hashlib import sha256



from log_exception import UPDATE_LOGGER

from blocks_manager import BlocksManager

from utils import OPTIONS_MANAGER

from utils import EXTEND_VALUE

from utils import FILE_MAP_ZERO_KEY

from utils import FILE_MAP_NONZERO_KEY

from utils import FILE_MAP_COPY_KEY

from utils import MAX_BLOCKS_PER_GROUP

from utils import UPDATE_BIN_FILE_NAME

from utils import FORBIDEN_UPDATE_IMAGE_SET





class FullUpdateImage:

    """

    Full image processing class

    """



    def __init__(self, target_package_images_dir,

                 full_img_list, full_img_name_list,

                 verse_script, full_image_path_list,

                 no_zip=False):

        self.target_package_images_dir = target_package_images_dir

        self.full_img_list = full_img_list

        self.full_img_name_list = full_img_name_list

        self.verse_script = verse_script

        self.full_image_path_list = full_image_path_list

        self.no_zip = no_zip



    def update_full_image(self):

        """

        Processing of the full image

        :return full_image_content_len_list: full image content length list

        :return full_image_file_obj_list: full image temporary file list

        """

        full_image_file_obj_list = []

        full_image_content_len_list = []

        for idx, each_name in enumerate(self.full_img_list):

            full_image_content = self.get_full_image_content(

                self.full_image_path_list[idx])

            img_name = self.full_img_name_list[idx][:-4]

            if full_image_content is False:

                UPDATE_LOGGER.print_log(

                    "Get full image content failed!",

                    log_type=UPDATE_LOGGER.ERROR_LOG)

                return False, False

            each_img = tempfile.NamedTemporaryFile(

                dir=self.target_package_images_dir,

                prefix="full_image%s" % img_name, mode='wb')

            each_img.write(full_image_content)

            each_img.seek(0)

            full_image_content_len_list.append(len(full_image_content))

            full_image_file_obj_list.append(each_img)

            # 全量流式升级处理

            OPTIONS_MANAGER.full_image_new_data[each_name] = full_image_content

            

            UPDATE_LOGGER.print_log(

                "Image %s full processing completed" % img_name)

        update_image_set = set(self.full_img_list) - FORBIDEN_UPDATE_IMAGE_SET

        if not self.no_zip and len(update_image_set) != 0:

            # No zip mode (no script command)

            image_write_cmd = self.verse_script.full_image_update(UPDATE_BIN_FILE_NAME)

            cmd = '%s_WRITE_FLAG%s' % (UPDATE_BIN_FILE_NAME, image_write_cmd)

            if each_name not in FORBIDEN_UPDATE_IMAGE_SET:

                self.verse_script.add_command(cmd=cmd)



        UPDATE_LOGGER.print_log(

            "All full image processing completed! image count: %d" %

            len(self.full_img_list))

        return full_image_content_len_list, full_image_file_obj_list



    @staticmethod

    def get_full_image_content(each_name):

        """

        Obtain the full image content.

        :param each_name: image name

        :return content: full image content if available; false otherwise

        """

        each_image_path = each_name

        if not os.path.exists(each_image_path):

            UPDATE_LOGGER.print_log(

                "The file is missing "

                "from the target package, "

                "the component: %s cannot be full update processed. " %

                each_image_path)

            return False

        with open(each_image_path, 'rb') as f_r:

            content = f_r.read()

        return content





class IncUpdateImage:

    """

    Increment update image class

    """



    def __init__(self, image_path, map_path):

        """

        Initialize the inc image.

        :param image_path: img file path

        :param map_path: map file path

        """

        self.image_path = image_path

        self.map_path = map_path

        self.offset_value_list = []

        self.care_block_range = None

        self.extended_range = None

        self.reserved_blocks = BlocksManager("0")

        self.file_map = []

        self.offset_index = []

        self.block_size = None

        self.total_blocks = None

        self.parse_raw_image_file(image_path, map_path)



    def parse_raw_image_file(self, image_path, map_path):

        """

        Parse the .img file.

        :param image_path: img file path

        :param map_path: map file path

        """

        self.block_size = block_size = 4096

        self.total_blocks = total_blocks = \

            os.path.getsize(self.image_path) // self.block_size

        reference = b'\0' * self.block_size

        with open(image_path, 'rb') as f_r:

            care_value_list, offset_value_list = [], []

            nonzero_blocks = []

            for i in range(self.total_blocks):

                blocks_data = f_r.read(self.block_size)

                if blocks_data != reference:

                    nonzero_blocks.append(i)

                    nonzero_blocks.append(i + 1)

            self.care_block_range = BlocksManager(nonzero_blocks)

            care_value_list = list(self.care_block_range.range_data)

            for idx, value in enumerate(care_value_list):

                if idx != 0 and (idx + 1) % 2 == 0:

                    be_value = int(care_value_list[idx - 1])

                    af_value = int(care_value_list[idx])

                    file_tell = be_value * block_size

                    offset_value_list.append(

                        (be_value, af_value - be_value,

                         file_tell, None))



            self.offset_index = [i[0] for i in offset_value_list]

            self.offset_value_list = offset_value_list

            extended_range = \

                self.care_block_range.extend_value_to_blocks(EXTEND_VALUE)

            all_blocks = BlocksManager(range_data=(0, total_blocks))

            self.extended_range = \

                extended_range.get_intersect_with_other(all_blocks). \

                get_subtract_with_other(self.care_block_range)

            self.parse_block_map_file(map_path, f_r)



    def parse_block_map_file(self, map_path, image_file_r):

        """

        Parses the map file for blocks where files are contained in the image.

        :param map_path: map file path

        :param image_file_r: file reading object

        :return:

        """

        remain_range = self.care_block_range

        temp_file_map = {}



        with open(map_path, 'r') as f_r:

            # Read the .map file and process each line.

            for each_line in f_r.readlines():

                each_map_path, ranges_value = each_line.split(None, 1)

                each_range = BlocksManager(ranges_value)

                each_range = each_range.get_subtract_with_other(BlocksManager("0"))

                # each_range may not contained in the remain range.

                intersect_range = each_range.get_intersect_with_other(remain_range)

                if each_range.size() != intersect_range.size():

                    each_range = intersect_range

                temp_file_map[each_map_path] = each_range

                # After the processing is complete,

                # remove each_range from remain_range.

                remain_range = remain_range.get_subtract_with_other(each_range)

        reserved_blocks = self.reserved_blocks

        # Remove reserved blocks from all blocks.

        remain_range = remain_range.get_subtract_with_other(reserved_blocks)



        # Divide all blocks into zero_blocks

        # (if there are many) and nonzero_blocks.

        zero_blocks_list = []

        nonzero_blocks_list = []

        nonzero_groups_list = []

        default_zero_block = ('\0' * self.block_size).encode()



        nonzero_blocks_list, nonzero_groups_list, zero_blocks_list = \

            self.apply_remain_range(

                default_zero_block, image_file_r, nonzero_blocks_list,

                nonzero_groups_list, remain_range, zero_blocks_list)



        temp_file_map = self.get_file_map(

            nonzero_blocks_list, nonzero_groups_list,

            reserved_blocks, temp_file_map, zero_blocks_list)

        self.file_map = temp_file_map



    def apply_remain_range(self, *args):

        """

        Implement traversal processing of remain_range.

        """

        default_zero_block, image_file_r, \

            nonzero_blocks_list, nonzero_groups_list, \

            remain_range, zero_blocks_list = args

        for start_value, end_value in remain_range:

            for each_value in range(start_value, end_value):

                # bisect 二分查找，b在self.offset_index中的位置

                idx = bisect.bisect_right(self.offset_index, each_value) - 1

                chunk_start, _, file_pos, fill_data = \

                    self.offset_value_list[idx]

                data = self.get_file_data(self.block_size, chunk_start,

                                          default_zero_block, each_value,

                                          file_pos, fill_data, image_file_r)



                zero_blocks_list, nonzero_blocks_list, nonzero_groups_list = \

                    self.get_zero_nonzero_blocks_list(

                        data, default_zero_block, each_value,

                        nonzero_blocks_list, nonzero_groups_list,

                        zero_blocks_list)

        return nonzero_blocks_list, nonzero_groups_list, zero_blocks_list



    @staticmethod

    def get_file_map(*args):

        """

        Obtain the file map.

        nonzero_blocks_list nonzero blocks list,

        nonzero_groups_list nonzero groups list,

        reserved_blocks reserved blocks ,

        temp_file_map temporary file map,

        zero_blocks_list zero block list

        :return temp_file_map file map

        """

        nonzero_blocks_list, nonzero_groups_list, \

            reserved_blocks, temp_file_map, zero_blocks_list = args

        if nonzero_blocks_list:

            nonzero_groups_list.append(nonzero_blocks_list)

        if zero_blocks_list:

            temp_file_map[FILE_MAP_ZERO_KEY] = \

                BlocksManager(range_data=zero_blocks_list)

        if nonzero_groups_list:

            for i, blocks in enumerate(nonzero_groups_list):

                temp_file_map["%s-%d" % (FILE_MAP_NONZERO_KEY, i)] = \

                    BlocksManager(range_data=blocks)

        if reserved_blocks:

            temp_file_map[FILE_MAP_COPY_KEY] = reserved_blocks

        return temp_file_map



    @staticmethod

    def get_zero_nonzero_blocks_list(*args):

        """

        Get zero_blocks_list, nonzero_blocks_list, and nonzero_groups_list.

        data: block data,

        default_zero_block: default to zero block,

        each_value: each value,

        nonzero_blocks_list: nonzero_blocks_list,

        nonzero_groups_list: nonzero_groups_list,

        zero_blocks_list: zero_blocks_list,

        :return new_zero_blocks_list: new zero blocks list,

        :return new_nonzero_blocks_list: new nonzero blocks list,

        :return new_nonzero_groups_list: new nonzero groups list.

        """

        data, default_zero_block, each_value, \

            nonzero_blocks_list, nonzero_groups_list, \

            zero_blocks_list = args

        # Check whether the data block is equal to the default zero_blocks.

        if data == default_zero_block:

            zero_blocks_list.append(each_value)

            zero_blocks_list.append(each_value + 1)

        else:

            nonzero_blocks_list.append(each_value)

            nonzero_blocks_list.append(each_value + 1)

            # The number of nonzero_blocks is greater than

            # or equal to the upper limit.

            if len(nonzero_blocks_list) >= MAX_BLOCKS_PER_GROUP:

                nonzero_groups_list.append(nonzero_blocks_list)

                nonzero_blocks_list = []

        new_zero_blocks_list, new_nonzero_blocks_list, \

            new_nonzero_groups_list = \

            copy.copy(zero_blocks_list), \

            copy.copy(nonzero_blocks_list),\

            copy.copy(nonzero_groups_list)

        return new_zero_blocks_list, new_nonzero_blocks_list, \

            new_nonzero_groups_list



    @staticmethod

    def get_file_data(*args):

        """

        Get the file data.

        block_size: blocksize,

        chunk_start: the start position of chunk,

        default_zero_block: default to zero blocks,

        each_value: each_value,

        file_pos: file position,

        fill_data: data,

        image_file_r: read file object,

        :return data: Get the file data.

        """

        block_size, chunk_start, default_zero_block, each_value, \

            file_pos, fill_data, image_file_r = args

        if file_pos is not None:

            file_pos += (each_value - chunk_start) * block_size

            image_file_r.seek(file_pos, os.SEEK_SET)

            data = image_file_r.read(block_size)

        else:

            if fill_data == default_zero_block[:4]:

                data = default_zero_block

            else:

                data = None

        return data



    def range_sha256(self, ranges):

        """

        range sha256 hash content

        :param ranges: ranges value

        :return:

        """

        hash_obj = sha256()

        for data in self.__get_blocks_set_data(ranges):

            hash_obj.update(data)

        return hash_obj.hexdigest()



    def write_range_data_2_fd(self, ranges, file_obj):

        """

        write range data to fd

        :param ranges: ranges obj

        :param file_obj: file obj

        :return:

        """

        for data in self.__get_blocks_set_data(ranges):

            file_obj.write(data)



    def get_ranges(self, ranges):

        """

        get ranges value

        :param ranges: ranges

        :return: ranges value

        """

        return [each_data for each_data in self.__get_blocks_set_data(ranges)]



    def __get_blocks_set_data(self, blocks_set_data):

        """

        Get the range data.

        """

        with open(self.image_path, 'rb') as f_r:

            for start, end in blocks_set_data:

                diff_value = end - start

                idx = bisect.bisect_right(self.offset_index, start) - 1

                chunk_start, chunk_len, file_pos, fill_data = \

                    self.offset_value_list[idx]



                remain = chunk_len - (start - chunk_start)

                this_read = min(remain, diff_value)

                if file_pos is not None:

                    pos = file_pos + ((start - chunk_start) * self.block_size)

                    f_r.seek(pos, os.SEEK_SET)

                    yield f_r.read(this_read * self.block_size)

                else:

                    yield fill_data * (this_read * (self.block_size >> 2))

                diff_value -= this_read



                while diff_value > 0:

                    idx += 1

                    chunk_start, chunk_len, file_pos, fill_data = \

                        self.offset_value_list[idx]

                    this_read = min(chunk_len, diff_value)

                    if file_pos is not None:

                        f_r.seek(file_pos, os.SEEK_SET)

                        yield f_r.read(this_read * self.block_size)

                    else:

                        yield fill_data * (this_read * (self.block_size >> 2))

                    diff_value -= this_read