#!/usr/bin/env python3
from pathlib import Path
from doit import get_var

DOIT_CONFIG = {
    'default_tasks': ['build'],
    'action_string_formatting': 'both'
}


def files(p):
    for i in p:
        if i.is_file():
            yield i


class RawFiles(list):
    def __init__(self, *args, **kwargs):
        self.src_path = Path(kwargs.pop('src_path'))
        self.image_path = Path(kwargs.pop('image_path'))
        super().__init__(*args, **kwargs)

    def sources(self):
        for i in self:
            yield self.src_path/i

    def targets(self):
        for i in self:
            yield self.image_path/i


class BasePatchwork(dict):
    patches_path = None
    patchwork_path = None
    vanilla_path = None

    def __init__(self, *arkg, **kwargs):
        self.patches_path = Path(kwargs.pop('patches_path'))
        self.patchwork_path = Path(kwargs.pop('patchwork_path'))
        self.vanilla_path = Path(kwargs.pop('vanilla_path'))
        super().__init__(*arkg, **kwargs)

    def targets(self):
        for i in self.keys():
            yield self.patchwork_path/i

    def key_vanilla(self, orig_key):
        return self.vanilla_path/orig_key

    def key_target(self, orig_key):
        return self.patchwork_path/orig_key

    def tasks(self, actions):
        def repack_action(action, **kwargs):
            if callable(action):
                return (action, [], kwargs)
            elif isinstance(action, str):
                return action.format(**kwargs)
            else:
                return action
        for k, v in self.items():
            original = self.vanilla_path/k
            patch = self.patches_path/v
            target = self.patchwork_path/k
            yield {
                'name': k,
                'file_dep': [original, patch],
                'targets': [target],
                'actions': [repack_action(action, original=original, patch=patch, target=target)
                            for action in actions],
                'clean': True
            }


class Patchwork(BasePatchwork):
    def tasks_(self, actions):
        for k, v in self.items():
            yield {
                'name': k,
                'file_dep': [self.patches_path/v, self.vanilla_path/k],
                'targets': [self.patchwork_path/k],
                'actions': [
                    f"mkdir -p {self.patchwork_path}/{Path(k).parent}",
                    f"patch -u --binary -N -o '{self.patchwork_path}/{k}' "
                    f"'{self.vanilla_path}/{k}' '{self.patches_path}/{v}'"
                ],  # action ( mkdir, patch )
                'clean': True
            }


class ScriptedPatchwork(BasePatchwork):
    def __init__(self, *args, **kwargs):
        from yaml import safe_load
        file = kwargs.pop('source_file')
        super().__init__(*args, **kwargs)
        with open(file) as s:
            for key, data in safe_load(s).items():
                try:
                    discovered_source = next(
                        (Path(self.vanilla_path)/'history'/'states').glob(f"{key}*"))
                    relative_path = discovered_source.relative_to(
                        self.vanilla_path)
                    self[relative_path] = data
                except StopIteration:
                    print(f"{key}: not found")


class RasterPatchwork(BasePatchwork):
    @staticmethod
    def patch_action(target, original, patch):
        from PIL import Image
        with Image.open(original) as base, Image.open(patch) as patch:
            def mask():
                for bg, p in zip(base.tobytes(), patch.tobytes()):
                    if p == 35:
                        yield bg
                    else:
                        yield p

            output = Image.frombytes(base.mode, base.size, bytes(mask()))
            output.putpalette(base.palette)
            target.parent.mkdir(exist_ok=True, parents=True)
            output.save(target)
            output.close()


class Patchset(list):
    def __init__(self, *args, **kwargs):
        self.modimage_path = Path(kwargs.pop('modimage_path'))
        super().__init__(*args, **kwargs)

    def image_files(self):
        for patches in self:
            if isinstance(patches, BasePatchwork):
                for target in patches:
                    yield self.modimage_path/target
            elif isinstance(patches, RawFiles):
                for target in patches:
                    yield self.modimage_path/target.relative_to(patches.src_path)

    def image_sources(self):
        for patches in self:
            if isinstance(patches, BasePatchwork):
                for target in patches:
                    yield patches.patchwork_path/target
            elif isinstance(patches, RawFiles):
                for target in patches:
                    yield target


mod_name = 'randchgs'
mod_install_path = get_var(
    'descriptor_mod_path', 'C:/Users/User/Documents/Paradox Interactive/Hearts of Iron IV/mod/randchgs')
dir_vanilla = Path(get_var('vanilla_path', './vanilla/current'))
if not dir_vanilla.is_dir():
    raise EnvironmentError(dir_vanilla)
dir_image = Path('build/image')
dir_modimage = dir_image/mod_name
dir_src_raw = Path('src/raw')
dir_src_raw_files = RawFiles(
    files(Path(dir_src_raw).rglob("*")),
    src_path=dir_src_raw,
    image_path=dir_image
)
dir_patches_path = Path('src/patches')
dir_patchwork_path = Path('build/patched')

patchwork = Patchwork(
    {
        'history/countries/ISR - Israel.txt': 'history_countries_ISR.patch',
        'history/states/454-Israel.txt': 'history_states_454.patch',
    },
    patches_path=dir_patches_path,
    patchwork_path=dir_patchwork_path,
    vanilla_path=dir_vanilla
)

patchwork_scripted = ScriptedPatchwork(
    source_file=dir_patches_path/'extrapoint.yaml',
    patches_path='.',
    patchwork_path=dir_patchwork_path,
    vanilla_path=dir_vanilla
)

patchwork_raster = RasterPatchwork(
    {
        'map/terrain.bmp': 'terrain.bmp',
    },
    patches_path=dir_patches_path,
    patchwork_path=dir_patchwork_path,
    vanilla_path=dir_vanilla
)

patch_set = Patchset(
    (
        dir_src_raw_files,
        patchwork,
        patchwork_scripted,
        patchwork_raster,
    ),
    modimage_path=dir_modimage
)


def task_patch_history_states_scripted():
    def history_state_patch(task, provinces):
        source = Path(next(iter(task.file_dep)))
        target = Path(task.targets[0])
        target.parent.mkdir(parents=True, exist_ok=True)
        seen_history = False
        wrote_history = False
        with source.open(newline='') as infile, target.open('w') as outfile:
            for line_in in infile:
                outfile.write(line_in)
                seen_history = seen_history or ('history' in line_in)
                if not wrote_history and seen_history and '{' in line_in:
                    for province_id, province_info in provinces.items():
                        outfile.write(
                            "\t\tvictory_points = {\r\n"
                            f"\t\t\t{province_id} {
                                province_info['points']}\r\n"
                            "\t\t}\r\n"
                        )
                    wrote_history = True
    for state_src, state_data in patchwork_scripted.items():
        yield {
            'name': state_src,
            'file_dep': [patchwork_scripted.key_vanilla(state_src)],
            'targets': [patchwork_scripted.key_target(state_src)],
            'actions': [(history_state_patch, [], {'provinces': state_data})],
            'clean': True
        }


def task_patch_raster():
    yield from patchwork_raster.tasks([patchwork_raster.patch_action])


def task_patch():
    def mkdir(target, original, patch):
        Path(target).parent.mkdir(parents=True, exist_ok=True)
    yield from patchwork.tasks([
        mkdir,
        "patch -u --binary -N -o '{target}' '{original}' '{patch}'"
    ])


def task_image():
    return {
        'file_dep': list(patch_set.image_sources()),
        'targets': list(patch_set.image_files()),
        'actions': [f"mkdir -p {dir_modimage}",
                    f"rsync -rv {dir_src_raw}/ {patchwork.patchwork_path}/"
                    f" {dir_modimage}/"],
        'clean': True
    }


def task_image_mod():
    def prepare_modname_mod(task):
        source = Path(next(iter(task.file_dep)))
        target = Path(task.targets[0])
        target.parent.mkdir(parents=True, exist_ok=True)
        with source.open() as s:
            with target.open('w') as o:
                o.write(s.read())
                o.write(f"\npath=\"{mod_install_path}\"\n")
    return {
        'file_dep': ['src/raw/descriptor.mod'],
        'targets': [f"{dir_image}/{mod_name}.mod"],
        'actions': [prepare_modname_mod],
        'clean': True
    }


def task_build():
    return {
        'file_dep': [f"{dir_image}/{mod_name}.mod"] + list(patch_set.image_files()),
        'targets': [f"{dir_image}/{mod_name}.zip"],
        'actions': ["rm -f %(targets)s",
                    f"cd {dir_image} && zip -r {mod_name}.zip {mod_name}",
                    f"cd {dir_image} && zip -r {mod_name}.zip {mod_name}.mod"],
        'clean': ['rm -rf build']
    }