import itertools
import tarfile
import json
import io
import hashlib
from pathlib import Path
from pymoku.parts import BuildStep, SourceFile
from pymoku.sdk import SCM_VERSION, log, keys


KEY_DIR = Path.home() / '.config/pymoku/keys'


class BarFile(BuildStep):
    VERSION = 0

    def __init__(self, barfile, *manifests):
        super().__init__(name=barfile)

        # list of manifests or list of lists of manifests
        try:
            self.manifests = list(itertools.chain.from_iterable(manifests))
        except TypeError:
            self.manifests = manifests

        self.barfile = self.build_path / barfile
        for m in self.manifests:
            self.depends_on(m.bitstreams())
        self.creates(SourceFile(self.barfile))
        self._bs_arcnames = {}

    def arcname_for(self, bitstream):
        # generate a unique filename for bitstreams in the archive
        if bitstream not in self._bs_arcnames:
            self._bs_arcnames[bitstream] = f'{len(self._bs_arcnames):03d}'
        return self._bs_arcnames[bitstream]

    def config_state(self):
        # Manifests aren't build steps, and we don't know when they change
        # until their own dependencies are built.
        sha = hashlib.sha256(SCM_VERSION.encode())
        sha.update(json.dumps(dict(keys.list_public_keys())).encode())
        return sha.digest()

    async def do_step(self):
        manifests = []
        for m in self.manifests:
            manifests.extend(m(self.arcname_for))

        self.update_progress(total=2 + len(self._bs_arcnames))

        manifest_bytes = self.create_manifest([m for m, bs in manifests])
        self.update_progress(advance=1)

        with tarfile.open(self.barfile, mode='w') as tf:
            m_tinfo = tarfile.TarInfo(name='MANIFEST')
            m_tinfo.size = len(manifest_bytes)
            tf.addfile(m_tinfo, io.BytesIO(manifest_bytes))

            # sign with all the keys TODO maybe specify which ones?
            for keyname, private_key in keys.list_private_keys():
                sig_bytes = private_key.sign(manifest_bytes)
                sig_tinfo = tarfile.TarInfo(name=f'MANIFEST.{keyname}.sig')
                sig_tinfo.size = len(sig_bytes)
                tf.addfile(sig_tinfo, io.BytesIO(sig_bytes))

            for bs, fname in self._bs_arcnames.items():
                bs_info = tf.gettarinfo(bs.path, arcname=fname)
                with open(bs.path, 'rb') as f:
                    tf.addfile(bs_info, fileobj=f)
                self.update_progress(advance=1)

        # 32 MB is the upload limit through the web interface, but also probably
        # a reasonable limit in general
        size = self.barfile.stat().st_size
        if size > 1024 * 2**15:
            log.warning(f'Barfile {self.barfile} exceeds 32MB ({size:d})')

    def create_manifest(self, items):
        manifest = dict(version=self.VERSION, items=items)
        return json.dumps(manifest, indent=2).encode()
