import io
import re
import hashlib
import tarfile
import json
import logging
import asyncio
import concurrent.futures
from functools import cache
from pymoku.sdk import BuildStep
from pymoku.plugins import PluginCache
from pymoku.parts import SourceFile


_cache = PluginCache()
log = logging.getLogger("pymokusdk")


@cache
def get_platform(platform_name):
    """ Find a Platform checkpoint
    """
    manifests = _cache.checkpoint_manifests(_id=f'pymoku:{platform_name}')

    if not manifests:
        _cache.store_manifests(_cache.cloud_manifests(_type='checkpoint'))
    manifests = _cache.checkpoint_manifests(_id=f'pymoku:{platform_name}')

    if not manifests:
        log.error(f"No checkpoints found for {platform_name}")
        return None

    log.debug(f"Using platform {platform_name} from {manifests[0]['location']}")

    return CheckpointPlatform(platform_name, manifests)


class CheckpointPlatform(BuildStep):
    """ This mimics the vivado.Platform interface but loads the checkpoint from
        a barfile rather than trying to build it.

        This step downloads the barfile to the local pymoku cache; the
        CheckpointShell unpacks each DCP into the build directory.
    """
    def __init__(self, name, manifests):
        super().__init__(name=name)
        # there is no "platform" checkpoint, only the individual cells
        self.manifests = manifests

        self.hwver = self.slot_shells()[0].hwver
        self.part = self.slot_shells()[0].part
        self.arch = self.slot_shells()[0].arch
        # TODO these should not be platform properties
        self.slot_ains = self.slot_shells()[0].num_ain
        self.slot_aouts = self.slot_shells()[0].num_aout
        self.slot_cbufs = self.slot_shells()[0].num_cbufs
        self.sha = manifests[0]['plat_sha']

    @cache
    def slot_shells(self):
        shells = [CheckpointShell(self, m) for m in self.manifests]
        return sorted([s for s in shells if s.slot is not None], key=lambda s: s.slot)

    async def do_step(self):
        # manifest may be in one shared barfile, reduce to unique barfiles
        barfiles = []
        to_fetch = []
        for m in self.manifests:
            if m['location'] != 'local' and m['barfile'] not in barfiles:
                to_fetch.append(m)
                barfiles.append(m['barfile'])

        # split downloads across mutliple threads (and unblock other BuildSteps)
        tasks = []
        loop = asyncio.get_running_loop()
        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as pool:
            for m in to_fetch:
                tasks.append(loop.run_in_executor(pool, _cache.download_barfile, m))
            await asyncio.gather(*tasks)


class CheckpointShell(BuildStep):
    def __init__(self, platform, manifest):
        self.platform = platform
        self.manifest = manifest
        self.hwver = manifest['extra']['hwver']
        self.cell_name = manifest['extra']['cell_name']
        self.pblock_name = manifest['extra']['pblock_name']
        self.part = manifest['extra']['part']
        self.arch = manifest['extra']['arch']
        self.num_ain = manifest['extra']['num_ain']
        self.num_aout = manifest['extra']['num_aout']
        self.num_cbufs = manifest['extra']['num_cbufs']

        name = f'{platform.name}-{self.safe_cell_name}'
        super().__init__(name=name)
        self.depends_on(platform)

        self.workdir = self.build_path / 'checkpoints' / platform.name
        self.shell_dcp = SourceFile(self.workdir / f'{name}.dcp')
        self.creates(self.shell_dcp)

    @property
    def plat_name(self):
        return self.platform.name

    @property
    def path(self):
        return next(iter(self.outputs())).path

    @property
    def safe_cell_name(self):
        name = self.cell_name.translate(
            str.maketrans({"[": "_", "]": None, "*": "_"}))
        name = name.replace('_INSTR', '')
        return name

    @property
    def slot(self):
        """ This makes some big assumptions about cell naming conventions"""
        if (m := re.search(r'INSTRS\[(\d+)\]', self.cell_name)) is None:
            return None
        return int(m.group(1))

    def config_state(self):
        sha = hashlib.sha256(json.dumps(self.manifest['extra']).encode())
        sha.update(self.manifest['sha256'].encode())
        sha.update(self.manifest['plat_sha'].encode())
        return sha.digest()

    async def do_step(self):
        barfile_data = _cache.download_barfile(self.manifest)
        barfile_data = io.BytesIO(barfile_data)

        with tarfile.open(fileobj=barfile_data) as arc:
            member = arc.getmember(self.manifest['filename'])
            assert member.isfile()
            data = arc.extractfile(member).read()

        sha = hashlib.sha256(data)

        assert sha.hexdigest() == self.manifest['sha256']

        self.shell_dcp.path.parent.mkdir(exist_ok=True, parents=True)
        self.shell_dcp.path.write_bytes(data)
