import logging
import zipfile
from io import BufferedWriter, BytesIO
from typing import Optional, Union, Dict, Tuple, Generator, List

from qm.StreamMetadata import (
    StreamMetadata,
    StreamMetadataError,
)
from qm.utils.async_utils import run_async
from qm.utils.general_utils import run_until_with_timeout
from qm.api.job_result_api import JobResultServiceApi
from qm.api.models.capabilities import ServerCapabilities
from qm.persistence import BaseStore
from qm.results.base_streaming_result_fetcher import (
    BaseStreamingResultFetcher,
    _parse_dtype,
    JobResultItemSchema,
    JobResultSchema,
)
from qm.results.multiple_streaming_result_fetcher import MultipleStreamingResultFetcher
from qm.results.single_streaming_result_fetcher import SingleStreamingResultFetcher

logger = logging.getLogger(__name__)


class StreamingResultFetcher:
    """Access to the results of a QmJob

    This object is created by calling [QmJob.result_handles][qm.jobs.running_qm_job.RunningQmJob.result_handles]

    Assuming you have an instance of StreamingResultFetcher:
    ```python
        job_results: StreamingResultFetcher
    ```
    This object is iterable:

    ```python
        for name, handle in job_results:
            print(name)
    ```

    Can detect if a name exists:

    ```python
    if "somename" in job_results:
        print("somename exists!")
        handle = job_results.get("somename")
    ```
    """

    def __init__(
        self,
        job_id: str,
        service: JobResultServiceApi,
        store: BaseStore,
        capabilities: ServerCapabilities,
    ) -> None:
        self._job_id = job_id
        self._service = service
        self._store = store
        self._schema = None
        self._capabilities = capabilities

        self._all_results: Dict[str, BaseStreamingResultFetcher] = {}
        self._add_job_results()

    def _add_job_results(self):
        self._schema = StreamingResultFetcher._load_schema(self._job_id, self._service)
        stream_metadata_errors, stream_metadata_dict = self._get_stream_metadata()
        for (name, item_schema) in self._schema.items.items():
            stream_metadata = stream_metadata_dict.get(name)
            if item_schema.is_single:
                result = SingleStreamingResultFetcher(
                    job_id=self._job_id,
                    schema=item_schema,
                    service=self._service,
                    store=self._store,
                    stream_metadata_errors=stream_metadata_errors,
                    stream_metadata=stream_metadata,
                    capabilities=self._capabilities,
                )
            else:
                result = MultipleStreamingResultFetcher(
                    job_results=self,
                    job_id=self._job_id,
                    schema=item_schema,
                    service=self._service,
                    store=self._store,
                    stream_metadata_errors=stream_metadata_errors,
                    stream_metadata=stream_metadata,
                    capabilities=self._capabilities,
                )
            self._all_results[name] = result

    def __getitem__(self, item: str):
        return self.get(item)

    def __getattr__(self, item: str):
        if item == "shape" or item == "__len__":
            return  # this is here because of a bug in pycharm debugger: ver: 2022.3.2 build #PY-223.8617.48 (24/1/23)
        return self.get(item)

    def _get_stream_metadata(
        self,
    ) -> Optional[Tuple[List[StreamMetadataError], Dict[str, StreamMetadata]]]:
        return self._service.get_program_metadata(job_id=self._job_id)

    def __iter__(self) -> Generator[Tuple[str, BaseStreamingResultFetcher], None, None]:
        for item in self._schema.items.values():
            yield item.name, self.get(item.name)

    def is_processing(self) -> bool:
        """Check if the job is still processing results

        Returns:
            True if results are still being processed, False otherwise
        """
        key = list(self._all_results.keys())[0]
        return self._all_results[key].is_processing()

    def save_to_store(
        self,
        writer: Optional[Union[BufferedWriter, BytesIO, str]] = None,
        flat_struct: bool = False,
    ):
        """Save all results to store (file system by default) in a single NPZ file

        Args:
            writer: An optional writer to be used instead of the pre-
                populated store passed to [qm.QuantumMachinesManager.QuantumMachinesManager][]
            flat_struct: results will have a flat structure - dimensions
                will be part of the shape and not of the type

        """
        own_writer = False
        if writer is None:
            own_writer = True
            writer = self._store.all_job_results(self._job_id).for_writing()
        zipf = None
        try:
            zipf = zipfile.ZipFile(writer, allowZip64=True, mode="w", compression=zipfile.ZIP_DEFLATED)
            for (name, result) in self:
                with zipf.open(f"{name}.npy", "w") as entry:
                    result.save_to_store(entry, flat_struct)
                pass
        finally:
            zipf.close()
            if own_writer:
                writer.close()

    @staticmethod
    def _load_schema(job_id: str, service: JobResultServiceApi) -> JobResultSchema:
        response = service.get_job_result_schema(job_id)
        return JobResultSchema(
            {
                item.name: JobResultItemSchema(
                    item.name,
                    _parse_dtype(item.simple_d_type),
                    tuple(item.shape),
                    item.is_single,
                    item.expected_count,
                )
                for item in response.items
            }
        )

    def get(self, name: str) -> Optional[Union[MultipleStreamingResultFetcher, SingleStreamingResultFetcher]]:
        """Get a handle to a named result from [stream_processing][qm.qua._dsl.stream_processing]

        Args:
            name: The named result using in [stream_processing][qm.qua._dsl.stream_processing]


        Returns:
            A handle object to the results `MultipleNamedJobResult` or `SingleNamedJobResult` or None if the named results in unknown
        """
        return self._all_results[name] if name in self._all_results else None

    def __contains__(self, name: str):
        return name in self._all_results

    def wait_for_all_values(self, timeout: Optional[float] = None) -> bool:
        """Wait until we know all values were processed for all named results

        Args:
            timeout: Timeout for waiting in seconds

        Returns:
            True if all finished successfully, False if any result was closed before done
        """

        def on_iteration() -> bool:
            all_job_states = [result.get_job_state() for result in self._all_results.values()]
            all_done = all(state.done for state in all_job_states)
            any_closed = any(state.closed for state in all_job_states)
            return all_done or any_closed

        def on_complete() -> bool:
            return all(result.get_job_state().done for result in self._all_results.values())

        return run_until_with_timeout(
            on_iteration_callback=on_iteration,
            on_complete_callback=on_complete,
            timeout=timeout if timeout else float("infinity"),
            timeout_message="Job was not done in time",
        )

    def get_debug_data(self, writer: Optional[Union[BufferedWriter, BytesIO, str]] = None):
        """
        Returns:
            debugging data to report to QM
        """
        if writer is None:
            writer = f"./{self._job_id}-DebugData.zip"

        owning_writer = False
        if type(writer) is str:
            writer = open(writer, "wb+")
            owning_writer = True

        try:
            run_async(self._fetch_all_job_debug_data(writer))

        finally:
            if owning_writer:
                writer.close()

    async def _fetch_all_job_debug_data(self, writer):
        async for result in self._service.get_job_debug_data(self._job_id):
            writer.write(result)
