"""This module contains the Image class."""
import configparser
import json
import pathlib
import random
from collections import defaultdict
from operator import itemgetter

from multimethod import multimethod

import woodblock.datagen
import woodblock.file
from woodblock.errors import ImageConfigError
from woodblock.scenario import Scenario


class Image:
    """The Image class represents a carving test image.

    An image contains a sequence of `Scenario` instances. An image has a fixed block size. All blocks smaller than the
    block size will be padded with data generated by a configurable data generator.

    Args:
        block_size: Block size of the image.
        padding_generator: A data generator used to generate padding.
    """

    def __init__(self, block_size: int = 512, padding_generator=woodblock.datagen.Random()):
        self._block_size = block_size
        self._scenarios = list()
        self._generate_padding = padding_generator

    def add(self, scenario):
        """Add a ``Scenario`` to the image."""
        self._scenarios.append(scenario)

    @staticmethod
    def from_config(path: pathlib.Path):
        """Create an Image instance based on a configuration file.

        Args:
            path: The path to the configuration file.
        """
        config = configparser.RawConfigParser()
        with path.open('r') as config_handle:
            try:
                config.read_file(config_handle)
            except configparser.Error as err:
                raise ImageConfigError(err.message)
        general = _parse_general_section(config)
        woodblock.file.corpus(path.absolute().parent / general['corpus'])
        if 'seed' in general:
            woodblock.random.seed(general['seed'])
        num_filler_blocks = (general['min filler blocks'], general['max filler blocks'])
        image = Image()
        for section in config.sections():
            if section != 'general':
                scenario = _parse_scenario_section(section_name=section, section=config[section],
                                                   num_filler_blocks=num_filler_blocks,
                                                   block_size=general['block size'])
                image.add(scenario)
        return image

    @multimethod
    def write(self, path: pathlib.Path):
        """Write the image to disk.

        This method write the image to the specified ``path``. Moreover, it also writes the image metadata to disk. The
        metadata file will be ``path`` with the ".json" extension. E.g. if ``path`` is "test-image.dd" then the actual
        image will be in "test-image.dd" and the metadata will be in "test-image.dd.json".

        Args:
            path: The output path.
        """
        with path.open('wb') as file_handle:
            self.write(file_handle)
        self._write_metadata(path)

    @multimethod
    def write(self, image_handle):  # pylint: disable=function-redefined
        """Write the image to disk.

        This method write the image to the specified ``path``. Note that no metadata is written. You have to call
        metadata yourself if you need it.

        Args:
            image_handle: A ``.write()``-supporting file-like object.
        """
        for scenario in self._scenarios:
            self._write_scenario(image_handle, scenario)

    @property
    def metadata(self):
        """Return the image metadata."""
        meta = {'block_size': self._block_size, 'seed': woodblock.random.get_seed(),
                'corpus': str(woodblock.file.get_corpus()),
                'scenarios': [s.metadata for s in self._scenarios]}
        self._update_metadata_with_image_offsets(meta, self._compute_image_offsets())
        return meta

    def _write_metadata(self, image_path: pathlib.Path):
        metadata_path = pathlib.Path('.'.join((str(image_path.absolute()), 'json')))
        with metadata_path.open('w') as file_handle:
            json.dump(self.metadata, file_handle)

    def _write_scenario(self, image, scenario):
        for fragment in scenario:
            for chunk in fragment:
                image.write(chunk)
            self._pad_fragment(image, fragment.size)

    def _pad_fragment(self, image, fragment_size):
        remainder = fragment_size % self._block_size
        if remainder != 0:
            padding_size = self._block_size - remainder
            image.write(self._generate_padding(padding_size))

    def _compute_image_offsets(self):
        image_offsets = defaultdict(dict)
        current_offset = 0
        for scenario in self._scenarios:
            for frag in scenario:
                frag_meta = frag.metadata
                frag_size = frag_meta['fragment']['size']
                end_offset = current_offset + frag_size
                image_offsets[frag_meta['file']['id']][frag_meta['fragment']['number']] = {'start': current_offset,
                                                                                           'end': end_offset}
                current_offset += frag_size
                if current_offset % self._block_size != 0:
                    current_offset += self._block_size - (current_offset % self._block_size)
        return image_offsets

    @staticmethod
    def _update_metadata_with_image_offsets(meta, image_offsets):  # pylint: disable=invalid-name
        for scenario_meta in meta['scenarios']:
            for file_meta in scenario_meta['files']:
                file_id = file_meta['original']['id']
                for frag_meta in file_meta['fragments']:
                    frag_meta['image_offsets'] = image_offsets[file_id][frag_meta['number']]


def _parse_general_section(config: dict) -> dict:
    if 'general' not in config:
        raise ImageConfigError('Mandatory "general" section is not present.')
    section = config['general']
    if 'corpus' not in section:
        raise ImageConfigError('Mandatory "corpus" key in "general" section is not present.')
    general = dict()
    general['corpus'] = section['corpus']
    if 'block size' in section:
        general['block size'] = int(section['block size'])
    else:
        general['block size'] = 512
    if 'seed' in section:
        general['seed'] = int(section['seed'])
    general['min filler blocks'] = _get_min_filler_fragment_blocks(section) or 1
    general['max filler blocks'] = _get_max_filler_fragment_blocks(section) or 10
    return general


def _get_min_filler_fragment_blocks(section):
    return _get_number_of_blocks(section, 'min')


def _get_max_filler_fragment_blocks(section):
    return _get_number_of_blocks(section, 'max')


def _get_number_of_blocks(section: dict, min_or_max: str):
    if '{} filler blocks'.format(min_or_max) in section:
        blocks = int(section['{} filler blocks'.format(min_or_max)])
        if blocks < 1:
            raise ImageConfigError('Invalid value for "{} filler blocks": {}.'.format(min_or_max, blocks))
        return blocks
    return None


def _parse_scenario_section(section_name: str, section: dict, num_filler_blocks: tuple, block_size: int):
    try:
        layout = _get_layout_type(section)
    except KeyError:
        raise ImageConfigError('Section [{}] contains no "layout" key.'.format(section_name))

    if layout == 'fragment-sequence':
        return _parse_fragment_sequence_layout(section_name, section, num_filler_blocks, block_size)
    if layout == 'intertwine':
        return _parse_intertwine_layout(section_name, section, block_size)
    raise ImageConfigError('Unsupported layout type: ')


def _get_layout_type(section: dict) -> str:
    layout = section['layout'].lower()
    if layout == 'intertwine':
        return 'intertwine'
    return 'fragment-sequence'


def _parse_intertwine_layout(section_name: str, section: dict, block_size: int):
    scenario = Scenario(section_name)
    try:
        num_files = section['num files']
    except KeyError:
        raise ImageConfigError(
            'Section [{}] has an intertwined layout but does not specify "num files".'.format(section_name))
    num_files = int(num_files)
    min_frags, max_frags = _parse_frags_nums(section_name, section)
    scenario.add(woodblock.file.intertwine_randomly(number_of_files=num_files, block_size=block_size,
                                                    min_fragments=min_frags, max_fragments=max_frags))
    return scenario


def _parse_fragment_sequence_layout(section_name: str, section: dict, num_filler_blocks: tuple, block_size: int):
    scenario = Scenario(section_name)
    files = dict()
    for key, value in section.items():
        if key.startswith('file'):
            file_number = int(key[4:])
            if file_number in files:
                files[file_number]['path'] = value
            else:
                files[file_number] = {'frags': None, 'path': value}
        elif key.startswith('frags_file') or key.startswith('frags file'):
            file_number = int(key[10:])
            if file_number in files:
                files[file_number]['frags'] = int(value)
            else:
                files[file_number] = {'frags': int(value), 'path': None}
    _assert_each_file_has_a_defined_num_of_frags(files, section_name)
    file_fragments = _create_file_fragments(files)
    layout = _parse_layout_line(section['layout'])

    min_filler_blocks = _get_min_filler_fragment_blocks(section) or num_filler_blocks[0]
    max_filler_blocks = _get_max_filler_fragment_blocks(section) or num_filler_blocks[1]
    if min_filler_blocks < 1 or max_filler_blocks < 1 or min_filler_blocks > max_filler_blocks:
        raise ImageConfigError(
            'Invalid min/max number of fillers blocks: min={}, max={}'.format(min_filler_blocks, max_filler_blocks))
    for fragment in layout:
        if fragment['type'] == 'file':
            scenario.add(file_fragments[fragment['file_num']][fragment['frag_num']])
        elif fragment['type'] == 'random':
            size = _get_filler_fragment_size(min_filler_blocks, max_filler_blocks, block_size)
            scenario.add(woodblock.fragments.RandomDataFragment(size))
        elif fragment['type'] == 'zeroes':
            size = _get_filler_fragment_size(min_filler_blocks, max_filler_blocks, block_size)
            scenario.add(woodblock.fragments.ZeroesFragment(size))
    return scenario


def _parse_frags_nums(section_name: str, section: dict) -> tuple:
    min_frags = 1
    max_frags = 4
    if 'min frags' in section:
        try:
            min_frags = int(section['min frags'])
        except ValueError:
            raise ImageConfigError('"min frags" has to be an integer > 0 in section [{}]'.format(section_name))
        else:
            if min_frags < 1:
                raise ImageConfigError('"min frags" has to be an integer > 0 in section [{}]'.format(section_name))
    if 'max frags' in section:
        try:
            max_frags = int(section['max frags'])
        except ValueError:
            raise ImageConfigError('"max frags" has to be an integer > 0 in section [{}]'.format(section_name))
        else:
            if max_frags < 1:
                raise ImageConfigError('"max frags" has to be an integer > 0 in section [{}]'.format(section_name))
    if min_frags > max_frags:
        raise ImageConfigError('"min frags" has to be <= "max frags" in section [{}]'.format(section_name))
    return min_frags, max_frags


def _get_filler_fragment_size(min_blocks: int, max_blocks: int, block_size: int) -> int:
    return random.randint(min_blocks, max_blocks) * block_size  # nosec


def _create_file_fragments(files):
    fragments = defaultdict(dict)
    for file_num in files:
        num_frags = files[file_num]['frags']
        file_path = files[file_num]['path']
        if file_path is not None:
            if (woodblock.file.get_corpus() / file_path).is_dir():
                frags = woodblock.file.draw_fragmented_files(file_path, number_of_files=1, min_fragments=num_frags,
                                                             max_fragments=num_frags)[0]
            else:
                frags = woodblock.file.File(file_path).fragment_randomly(num_fragments=num_frags)
        else:
            frags = woodblock.file.draw_fragmented_files(number_of_files=1, min_fragments=num_frags,
                                                         max_fragments=num_frags)[0]
        fragments[file_num] = {i + 1: f for i, f in enumerate(frags)}
    return fragments


def _assert_each_file_has_a_defined_num_of_frags(files, section):  # pylint: disable=invalid-name
    for num, file in files.items():
        if file['frags'] is None:
            raise ImageConfigError('No "frags file{} definition is missing in section "[{}]".'.format(num, section))


def _parse_layout_line(line: str) -> list:
    layout = list()
    for token in (x.strip() for x in line.split(',')):
        token = token.lower()
        if token == 'r':  # nosec
            layout.append({'type': 'random'})
        elif token == 'z':  # nosec
            layout.append({'type': 'zeroes'})
        else:
            file_num, frag_num = _parse_file_definition(token)
            layout.append({'type': 'file', 'file_num': int(file_num), 'frag_num': int(frag_num)})
    return layout


def _parse_file_definition(token: str):
    if '.' in token:
        sep = '.'
    elif '-' in token:
        sep = '-'
    else:
        raise ImageConfigError('No valid file/fragment separator found. Use "-" or ".".')
    return token.split(sep)


class ImageLogParser:
    def __init__(self, fp):
        self._log = json.load(fp)

    def get_fragment_order(self):
        fragments = list()
        for scenario in self._log['scenarios']:
            for file in scenario['files']:
                fragments.extend((scenario['name'], file['original']['id'], file['original']['path'],
                                  frag['number'], frag['file_offsets']['start'], frag['file_offsets']['end'],
                                  frag['image_offsets']['start'], frag['image_offsets']['end'])
                                 for frag in file['fragments'])
        return tuple(sorted(fragments, key=itemgetter(6)))
