import pickle

from pathlib import Path
from zipfile import ZipFile, ZIP_DEFLATED
from itertools import islice, repeat, chain
from typing import Union, Sequence, Mapping, Iterable, Any

from coba.context import CobaContext
from coba.primitives import Source, Sink, Filter, Environment, Interaction

class ObjectsToZipMember(Sink[Iterable[Sequence[object]]]):
    def __init__(self, zip:str):
        self._zip = zip
        self._start = 0
        if Path(self._zip).exists():
            for n in ZipFile(self._zip).namelist():
                try:
                    self._start = max(self._start,int(n)+1)
                except ValueError: #pragma: no cover
                    pass

    def write(self, envs: Iterable[Sequence[object]]) -> None:
        for i,env in enumerate(envs,self._start):
            with CobaContext.logger.time(f"Writing environment {i}..."):
                # Dump before opening the file so we
                # can keep it open for as short a time
                # as possible.
                dumps = list(map(pickle.dumps,env))
                with ZipFile(self._zip,mode='a',compression=ZIP_DEFLATED) as zip:
                    with zip.open(str(i),mode='w') as f:
                        f.writelines(dumps)

class ZipMemberToObjects(Source[Iterable[object]]):
    def __init__(self, zip:str, member:str):
        self._zip    = zip
        self._member = member

    def read(self) -> Iterable[object]:
        try:
            with ZipFile(self._zip) as z:
                with z.open(self._member) as m:
                    yield from map(pickle.load, repeat(m))
        except EOFError:
            pass

class EnvironmentsToObjects(Filter[Environment, Iterable]):
    def filter(self, env: Union[Environment,Sequence[Environment]]) -> Iterable[Iterable[object]]:
        with CobaContext.logger.time("Materializing environment..."):
            yield list(self._env_to_objects(env))

    def _env_to_objects(self,env):
        from coba import __version__ #imported here to avoid circular dependency
        yield {"version":2,"coba_version":__version__}
        yield env.params
        I = iter(env.read())
        batch = list(islice(I,1000))
        while batch:
            yield batch
            batch = list(islice(I,1000))

class EnvironmentFromObjects(Environment):
    def __init__(self, source: Source[Iterable[object]]) -> None:
        self._source = source

    @property
    def params(self) -> Mapping[str,Any]:
        return list(islice(self._source.read(),1,2))[0]

    def read(self) -> Iterable[Interaction]:

        I = iter(self._source.read())
        version_data = list(islice(I,2))[0]

        if version_data['version'] == 1: #pragma: no cover
            yield from I
        else:
            yield from chain.from_iterable(I)
