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

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

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():
                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

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)
            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'))
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(
    {
        'map/definition.csv': 'map_definition.patch',
        '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' ]
    }