import json
import os
import stat
import tarfile

from occystrap import common
from occystrap import constants
from occystrap import util
from occystrap.outputs.base import ImageOutput
from occystrap.util import safe_path_join
from shakenfist_utilities import logs


LOG = logs.setup_console(__name__)


class MountWriter(ImageOutput):
    def __init__(self, image, tag, image_path):
        super().__init__()

        self.image = image
        self.tag = tag
        self.image_path = image_path

        self.tar_manifest = [{
            'Layers': [],
            'RepoTags': [
                '%s:%s' % (self.image.split('/')[-1],
                           self.tag)]
        }]

        self.bundle = {}

        # For out-of-order layer delivery
        self._indexed_layers = []

        if not os.path.exists(self.image_path):
            os.makedirs(self.image_path)

    @property
    def requires_ordered_layers(self):
        return False

    def _manifest_filename(self):
        return 'manifest'

    def fetch_callback(self, digest):
        layer_file_in_dir = safe_path_join(
            self.image_path, digest, 'layer.tar')
        LOG.debug('Layer file is %s' % layer_file_in_dir)
        return not os.path.exists(layer_file_in_dir)

    def process_image_element(self, element):
        if element.element_type == constants.CONFIG_FILE:
            config_data = element.data.read()
            with open(safe_path_join(
                    self.image_path,
                    element.name), 'wb') as f:
                d = json.loads(config_data)
                f.write(json.dumps(
                    d, indent=4,
                    sort_keys=True).encode('ascii'))
            self.tar_manifest[0]['Config'] = \
                element.name
            self._track_element(
                element.element_type, len(config_data))

        elif element.element_type == \
                constants.IMAGE_LAYER:
            layer_dir = safe_path_join(
                self.image_path, element.name)
            if not os.path.exists(layer_dir):
                os.makedirs(layer_dir)

            layer_file = os.path.join(
                element.name, 'layer.tar')

            if element.layer_index is not None:
                self._indexed_layers.append(
                    (element.layer_index, layer_file))
            else:
                self.tar_manifest[0][
                    'Layers'].append(layer_file)

            layer_file_in_dir = safe_path_join(
                self.image_path, layer_file)
            layer_size = 0
            if os.path.exists(layer_file_in_dir):
                LOG.debug(
                    'Skipping layer already in'
                    ' output directory')
                layer_size = os.path.getsize(
                    layer_file_in_dir)
            else:
                with open(
                        layer_file_in_dir, 'wb') as f:
                    d = element.data.read(102400)
                    while d:
                        layer_size += len(d)
                        f.write(d)
                        d = element.data.read(102400)

                layer_dir_in_dir = safe_path_join(
                    self.image_path, element.name,
                    'layer')
                os.makedirs(layer_dir_in_dir)
                with tarfile.open(
                        layer_file_in_dir) as layer:
                    for mem in layer.getmembers():
                        dirname, filename = \
                            os.path.split(mem.name)

                        # Some light reading on how this
                        # works...
                        # https://www.madebymikal.com/interpreting-whiteout-files-in-docker-image-layers/
                        # https://github.com/opencontainers/image-spec/blob/main/layer.md#opaque-whiteout
                        if filename == '.wh..wh..opq':
                            # A deleted directory, but
                            # only for layers below this
                            # one.
                            os.setxattr(
                                os.path.join(
                                    layer_dir_in_dir,
                                    dirname),
                                'trusted.overlay.opaque',
                                b'y')

                        elif filename.startswith(
                                '.wh.'):
                            # A single deleted element,
                            # which might not be a file.
                            os.mknod(
                                os.path.join(
                                    layer_dir_in_dir,
                                    mem.name[4:]),
                                mode=stat.S_IFCHR,
                                device=0)

                        else:
                            path = mem.name
                            layer.extract(
                                path,
                                path=layer_dir_in_dir)

            self._track_element(
                element.element_type, layer_size)

    def finalize(self):
        # Reconstruct layer order from indices
        if self._indexed_layers:
            self._indexed_layers.sort(
                key=lambda x: x[0])
            self.tar_manifest[0]['Layers'] = [
                name for _, name
                in self._indexed_layers]

        manifest_filename = self._manifest_filename() + '.json'
        manifest_path = safe_path_join(
            self.image_path, manifest_filename)
        with open(manifest_path, 'wb') as f:
            f.write(json.dumps(self.tar_manifest, indent=4,
                               sort_keys=True).encode('ascii'))

        c = {}
        catalog_path = os.path.join(self.image_path, 'catalog.json')
        if os.path.exists(catalog_path):
            with open(catalog_path, 'r') as f:
                c = json.loads(f.read())

        c.setdefault(self.image, {})
        c[self.image][self.tag] = manifest_filename
        with open(catalog_path, 'w') as f:
            f.write(json.dumps(c, indent=4, sort_keys=True))

        self._log_summary()

    def write_bundle(self, container_template=constants.RUNC_SPEC_TEMPLATE,
                     container_values=None):
        if not container_values:
            container_values = {}

        rootfs_path = os.path.join(self.image_path, 'rootfs')
        if not os.path.exists(rootfs_path):
            os.makedirs(rootfs_path)
        LOG.info('Writing image bundle to %s' % rootfs_path)

        working_path = os.path.join(self.image_path, 'working')
        if not os.path.exists(working_path):
            os.makedirs(working_path)

        delta_path = os.path.join(self.image_path, 'delta')
        if not os.path.exists(delta_path):
            os.makedirs(delta_path)

        # The newest layer is listed first in the mount command
        layer_dirs = []
        self.tar_manifest[0]['Layers'].reverse()
        for layer in self.tar_manifest[0]['Layers']:
            layer_dirs.append(safe_path_join(
                self.image_path,
                layer.replace('.tar', '')))

        # Extract the rootfs as overlay mounts
        util.execute('mount -t overlay overlay -o lowerdir=%(layers)s,'
                     'upperdir=%(upper)s,workdir=%(working)s %(rootfs)s'
                     % {
                         'layers': ':'.join(layer_dirs),
                         'upper': delta_path,
                         'working': working_path,
                         'rootfs': rootfs_path
                     })

        # Rename the container configuration to a well known location. This is
        # not part of the OCI specification, but is convenient for now.
        container_config_filename = os.path.join(self.image_path,
                                                 'container-config.json')
        runtime_config_filename = os.path.join(self.image_path, 'config.json')
        os.rename(safe_path_join(
                      self.image_path,
                      self.tar_manifest[0]['Config']),
                  container_config_filename)

        common.write_container_config(container_config_filename,
                                      runtime_config_filename,
                                      container_template=container_template,
                                      container_values=container_values)
