import json
import hashlib
import tarfile
import io
import functools
import concurrent.futures
import importlib
import sys
import warnings
from tinydb import TinyDB, where
from tinydb.storages import MemoryStorage
from pymoku import log
from pymoku.progress import progress
from pathlib import Path
from zipimport import zipimporter
from tempfile import TemporaryDirectory
from packaging.version import Version, InvalidVersion

try:
    import requests
except ImportError:
    warnings.warn("'requests' not installed.  Cloud plugins unavailable")

BARFILE_REPO = 'https://pymoku-barfile-repo-237683438523.australia-southeast1.run.app'
LOCAL_CACHE = Path.home() / '.pymoku/barfile_cache'
DEV_CACHE = None


def catch(throws=Exception, returns=None):
    def wrapper(func):
        @functools.wraps(func)
        def inner(*args, **kwargs):
            try:
                return func(*args, **kwargs)
            except throws as exp:
                log.warning(f'Caught: {exp}')
                return returns
        return inner
    return wrapper


def _try_version(manifest: str):
    try:
        return Version(manifest['version'])
    except (InvalidVersion, TypeError, KeyError):
        return Version('0.0dev0')


def _preferred_location(a, b):
    if a is None:
        return b

    if b is None:
        return a

    assert a['sha256'] == b['sha256'], 'a and b are not the same'

    if a['location'] == 'device':
        return a
    if b['location'] == 'device':
        return b
    if a['location'] == 'local':
        return a
    if b['location'] == 'local':
        return b
    return a


def _preferred(a, b):
    """ Choose between two similar manifest items by version and locality

    Greater versions are prioritised, followed by a cache locality closer
    to the device.

    If those two things are equal then return a.

    Args:
        a, b: manifest dictionaries to compare

    Return:
        The preferred item
    """
    if a is None:
        return b

    if b is None:
        return a

    assert a['_id'] == b['_id'], 'items are not the same _id'
    assert a['_type'] == b['_type'], 'items are not the same _type'

    a_ver = _try_version(a)
    b_ver = _try_version(b)

    if a_ver.is_devrelease and a['location'] != 'local':
        return b

    if b_ver.is_devrelease and b['location'] != 'local':
        return a

    if a_ver > b_ver:
        return a

    if a_ver < b_ver:
        return b

    if a['location'] == 'device':
        return a
    if b['location'] == 'device':
        return b
    if a['location'] == 'local':
        return a
    if b['location'] == 'local':
        return b
    return a


def _reduce(manifests):
    # assumes manifests are already for an equivalent function (apples to apples)
    return functools.reduce(_preferred, manifests, None)


def _reduce_by_id(manifests):
    """
    Return the newest version, and closest stored of each manifest, grouped by
    id.

    Any compatibility requirements should be part of the query creating the
    manifests input.
    """
    keys = {m['_id'] for m in manifests}
    return [_reduce([n for n in manifests if n['_id'] == k]) for k in keys]


def _reduce_by_location(manifests):
    cache = {}
    for m in manifests:
        cache[m['sha256']] = _preferred_location(cache.get(m['sha256']), m)
    return list(cache.values())


class DefaultPlugin(object):
    def get_name(self, _id):
        return _id.replace('pymoku:', '')

    def attach(self, moku, manifest):
        from pymoku import Slot
        return Slot(moku, manifest['slot'])

    def applet_frame(self):
        return None


def _download_from_cloud(manifest):
    with progress() as prog:
        prog.set_message(f"Downloading {manifest['_id']}")
        prog.set_progress(0.0)
        r = requests.get(manifest['barfile'], stream=True)
        total_size = int(r.headers.get("content-length", 0))
        i = 0
        data = bytearray(total_size)
        for chunk in r.iter_content(0x40000):
            data[i:i + len(chunk)] = chunk
            i += len(chunk)
            prog.set_progress(i / total_size)
        return data


class PluginCache:
    """ Allows searching, fetching and storing Plugins
    """
    def __init__(self):
        # optional constructor for standalone uses.  Mixin init via
        # load_plugins_for_moku()
        self._db = TinyDB(storage=MemoryStorage)
        self.store_manifests(self.local_manifests())

    def functions_for_platform(self, platform_sha):
        """
        Returns:
            A list of functions available for the specified platform.
        """
        return _reduce_by_id(
            self._db.search((where('_type') == 'instrument')
                            & (where('plat_sha') == platform_sha)))

    def function_ids_for_platform(self, platform_sha):
        return {x['_id'] for x in self.functions_for_platform(platform_sha)}

    def function_for_slot(self, _id, platform_sha, slot):
        fs = _reduce_by_id(
            self._db.search((where('_type') == 'instrument')
                            & (where('_id') == _id)
                            & (where('plat_sha') == platform_sha)
                            & (where('slot') == slot)))
        fs = sorted(fs, key=lambda f: _try_version(f))
        if not fs:
            log.warning(f'no bitstream available {_id}:{slot}:{platform_sha}')
            return None
        return fs[-1]  # TODO version filter/check sorting

    def platform_manifest(self, platform_sha):
        """
        Returns:
            The manifest dictionary for the specified platform.
        """
        return _reduce(
            self._db.search((where('_type') == 'platform')
                            & (where('sha256') == platform_sha)))

    def function_manifest(self, function_id, function_sha):
        """
        Returns:
            The manifest dictionary for the specified function.
        """
        return _reduce(
            self._db.search((where('_type') == 'instrument')
                            & (where('_id') == function_id)
                            & (where('sha256') == function_sha)))

    def plugin_manifest(self, function_id):
        _id = ":".join(function_id.split(':')[:2])
        return _reduce(
            self._db.search((where('_type') == 'python')
                            & (where('_id') == _id)))

    def checkpoint_manifests(self, _id):
        # TODO this will return all the versions
        return self._db.search((where('_type') == 'checkpoint')
                               & (where('_id') == _id))

    def plugin_for_function(self, function_id):
        """ Get a plugin module defined from a loaded barfile for the given id

        Args:
            function_id: Full id string for the function being loaded.  The
                trailing component will be stripped before lookup
                example: 'pymoku:oscilloscope:4x' will be resolved as
                'pymoku:oscilloscope' as the final component will identify a
                specific configuration for the given class

        Returns:
            Plugin module for given id
        """
        if (manifest := self.plugin_manifest(function_id)) is None:
            log.warning(f"No plugin for {function_id}")
            return DefaultPlugin()

        if (plugin := self._plugin_cache.get(function_id)) is None \
                or _try_version(manifest).is_devrelease:
            log.debug(f"Load plugin {manifest['_id']}-{manifest['version']}")
            plugin = self.load_plugin(manifest)
            self._plugin_cache[function_id] = plugin

        return plugin

    def store_manifests(self, manifests):
        for manifest in manifests:
            # update matching records with potentially different locations
            self._db.upsert(manifest, (
                (where('_id') == manifest['_id'])
                & (where('_type') == manifest['_type'])
                & (where('sha256') == manifest['sha256'])
                & (where('version') == manifest['version'])
                & (where('location') == manifest['location'])))

    def local_manifests(self, path=LOCAL_CACHE):
        if path.is_file():
            paths = [path]
        else:
            paths = Path(path).rglob('*.bar')

        for path in paths:
            for manifest in self.barfile_manifest(path):
                manifest['location'] = 'local'
                manifest['barfile'] = path
                yield manifest

    @catch(returns=[])
    def cloud_manifests(self, **params):
        r = requests.get(BARFILE_REPO + '/barfiles/', params=params)
        manifests = r.json()
        for manifest in manifests:
            manifest['location'] = 'cloud'
        return manifests

    @catch(returns=[])
    def barfile_manifest(self, path):
        # extract tarfile
        with tarfile.open(path, mode='r') as tar:
            manifest = tar.extractfile('MANIFEST')
            manifest = json.loads(manifest.read().decode())
            return manifest['items']

    def download_barfile(self, manifest):
        """ Download the barfile from path specified in the manifest.

        Args:
            manifest: Manifest item to load

            Expected to have the location and barfile key added

        Returns:
            The requested barfile as a bytearray
        """
        if manifest['location'] == 'local':
            return Path(manifest['barfile']).read_bytes()

        # check if we already have it cached locally
        if (local := self._db.get(
                (where('_id') == manifest['_id'])
                & (where('_type') == manifest['_type'])
                & (where('sha256') == manifest['sha256'])
                & (where('version') == manifest['version'])
                & (where('location') == 'local'))):
            return Path(local['barfile']).read_bytes()

        # anything that isn't local we write to local cache, so after refresh
        # the local version will be preferred
        elif manifest['location'] == 'device':
            # we can assume self is a moku
            data = self._fileserv.get('b', Path(manifest['barfile']).name)
        elif manifest['location'] == 'cloud':
            data = _download_from_cloud(manifest)
        else:
            raise Exception(f"Can't retrieve barfile: {manifest['barfile']} {manifest['location']}")

        # Don't cache dev versions
        if _try_version(manifest).is_devrelease:
            return data

        # a barfile may contain many more items than just the requested manifest
        # so reloading the cache afterwards is recommended

        # best name I can think of is hash of the manifest
        with tarfile.TarFile.open(fileobj=io.BytesIO(data), mode='r') as tar:
            manifest = tar.extractfile(member='MANIFEST')
            barfile_name = hashlib.md5(manifest.read()).hexdigest()

        # write to local cache, and reload manifests to the _db
        LOCAL_CACHE.mkdir(parents=True, exist_ok=True)
        target_barfile = (LOCAL_CACHE / barfile_name).with_suffix('.bar')
        target_barfile.write_bytes(data)
        self.store_manifests(self.local_manifests(target_barfile))

        return data


class MokuPlugins(PluginCache):
    """ A mixin specifically for a Moku to manipulate plugins (slots/platforms)

    Uses a local TinyDB (NoSQL/Document DB) as a cache for what is available
    locally, on the device as well as cloud sources.
    """
    def install_barfile(self, manifest):
        if manifest is None:
            return

        if manifest.get('location') == 'device':
            # already on the device
            return

        if self._db.search((where('_id') == manifest['_id'])
                           & (where('sha256') == manifest['sha256'])
                           & (where('location') == 'device')):
            # already on the device
            return

        try:
            data = self.download_barfile(manifest)
            md5 = hashlib.md5(data)
            # need the remote name to be unique but detirministic
            fname = f'{md5.hexdigest()}.bar'
            fname = fname.replace(":", "-").replace("+", "_")
            self._fileserv.send('b', fname, data)
            self.store_manifests(self._local_manifests(LOCAL_CACHE))  # update cache
            self.store_manifests(self._device_manifests())
        except Exception as e:
            log.warning(f"Can't upload barfile for {manifest['_id']}")
            log.warning(str(e))

    def install_empties_for_platform(self):
        empties = [self.function_for_slot('pymoku:empty', self.platform_sha, slot=i)
                   for i in range(self.num_slots)]
        # filter for unique barfiles (might need to check locations)
        empties = dict.fromkeys({x['barfile']: x for x in empties if x}).values()
        with concurrent.futures.ThreadPoolExecutor() as pool:
            pool.map(self.install_barfile, empties)

    def load_plugins_for_moku(self):
        """ Find, download and load all plugin objects from Moku self.

        Caches the results on self to be accessed by other MokuPlugin methods.
        """
        with progress() as prog:
            prog.set_message("Reload plugins")
            prog.set_progress(0.0)
            # just start a new db each time
            self._db = TinyDB(storage=MemoryStorage)
            self._plugin_cache = dict()
            # these are in order of priority
            self.store_manifests(self._cloud_platforms())
            prog.set_progress(0.2)
            self.store_manifests(self._cloud_functions())
            prog.set_progress(0.4)
            self.store_manifests(self._cloud_plugins())
            prog.set_progress(0.6)
            self.store_manifests(self._device_manifests())
            prog.set_progress(0.7)
            self.store_manifests(self._local_manifests(LOCAL_CACHE))
            prog.set_progress(0.8)
            if DEV_CACHE:
                log.warning(f"Loading dev barfiles from {DEV_CACHE}")
                self.store_manifests(self._local_manifests(DEV_CACHE))

    def platforms_for_moku(self):
        """
        Returns:
            A list of platform objects that are available to deploy.
        """
        # return _reduce_by_id(self._db.search(where('_type') == 'platform'))
        return _reduce_by_location(self._db.search((where('_type') == 'platform')
                                                   & (where('hwver') == self.hw_ver)))

    def load_functions_for_platform(self):
        self.store_manifests(self._cloud_functions())

    def _cloud_platforms(self):
        return self.cloud_manifests(_type='platform', hwver=self.hw_ver)

    def _cloud_functions(self):
        if self.platform_sha is None:
            return []
        return self.cloud_manifests(_type='instrument', plat_sha=self.platform_sha)

    def _cloud_plugins(self):
        return self.cloud_manifests(_type='python')

    def _local_manifests(self, path):
        paths = Path(path).glob('*.bar')
        for path in paths:
            for manifest in self.barfile_manifest(path):
                # other sources come pre filtered
                if manifest['_type'] == 'platform' and manifest['hwver'] != self.hw_ver:
                    continue
                manifest['location'] = 'local'
                manifest['barfile'] = path
                yield manifest

    @catch(returns=[])
    def _device_manifests(self):
        manifests = self.bitstreams()
        for manifest in manifests:
            manifest['location'] = 'device'
        return manifests

    def load_plugin(self, manifest):
        log.info(f"Loading plugin for {manifest['_id']}")
        assert manifest['_type'] == 'python'

        try:
            with TemporaryDirectory() as tmp:
                barfile = Path(tmp) / 'barfile.bar'
                barfile.write_bytes(self.download_barfile(manifest))

                with tarfile.open(name=barfile, mode='r') as tar:
                    member = manifest['filename']
                    zipfile = Path(tmp) / member
                    tar.extract(member, path=zipfile.parent)
                    imp = zipimporter(str(zipfile))
                    imp.invalidate_caches()
                    spec = imp.find_spec('__pymoku_plugin')
                    mod = importlib.util.module_from_spec(spec)
                    sys.modules['__pymoku_plugin'] = mod
                    imp.exec_module(mod)
                    p = mod.Plugin()
                    p.manifest = manifest
                    return p
        except Exception:
            log.exception(f"Couldn't load plugin {manifest['_id']}")
            return DefaultPlugin()
        finally:
            # remove all plugin modules
            # There is never a need to import from __pymoku_plugin.  When
            # hot-reloading, this forces all import within a plugin to reload
            # in the correct order.
            # It also prevents namespace conflicts between plugins
            # Only caveat is that deferred imports within a plugin won't work.
            ks = [k for k in sys.modules.keys() if k.startswith('__pymoku')]
            for k in ks:
                del sys.modules[k]
