#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import threading
import time
import uuid
from numbers import Integral

import numpy as np

from .core import ENTITY_TYPE, get_tiled
from .core.operand import Fetch
from .context import get_context, LocalContext
from .executor import Executor
from .config import options
from .utils import classproperty, calc_nsplits, serialize_graph
try:
    from .resource import cpu_count, cuda_count
except ImportError:  # pragma: no cover
    from multiprocessing import cpu_count
    from .lib.nvutils import get_device_count as cuda_count


class LocalSession(object):
    def __init__(self, **kwargs):
        engine = kwargs.pop('engine', None)
        self._endpoint = None
        self._session_id = uuid.uuid4()
        self._context = LocalContext(self)
        self._executor = Executor(engine=engine, storage=self._context)

        self._mut_tensor = dict()
        self._mut_tensor_data = dict()

        if kwargs:
            unexpected_keys = ', '.join(list(kwargs.keys()))
            raise TypeError(f'Local session got unexpected arguments: {unexpected_keys}')

    @property
    def endpoint(self):
        return self._endpoint

    @endpoint.setter
    def endpoint(self, endpoint):
        if endpoint is not None:
            raise ValueError('Local session cannot set endpoint')
        self._endpoint = endpoint

    @property
    def session_id(self):
        return self._session_id

    @property
    def executor(self):
        return self._executor

    @property
    def context(self):
        return self._context

    @context.setter
    def context(self, new_context):
        self._context = new_context
        self._executor.storage = new_context

    def get_cpu_count(self):
        return self._context.get_ncores()

    @property
    def executed_tileables(self):
        return self._executor.stored_tileables.keys()

    def run(self, *tileables, **kw):
        with self.context:
            if self._executor is None:
                raise RuntimeError('Session has closed')
            dest_gpu = all(tileable.op.gpu for tileable in tileables)
            if dest_gpu:
                self._executor._engine = 'cupy'
            else:
                self._executor._engine = None
            if 'n_parallel' not in kw:
                if dest_gpu:
                    # GPU
                    cnt = cuda_count() or 0
                    if cnt == 0:
                        raise RuntimeError('No GPU found for execution. '
                                           'Make sure NVML library is in your library path.')
                    kw['n_parallel'] = cnt
                else:
                    # CPU
                    kw['n_parallel'] = cpu_count()
            # set number of running cores
            self.context.set_ncores(kw['n_parallel'])
            res = self._executor.execute_tileables(tileables, **kw)
            return res

    def _update_tileable_shape(self, tileable):
        from .optimizes.tileable_graph import tileable_optimized

        new_nsplits = self._executor.get_tileable_nsplits(tileable)
        tiled = get_tiled(tileable, mapping=tileable_optimized)
        for t in (tileable, tiled):
            t._update_shape(tuple(sum(nsplit) for nsplit in new_nsplits))
        tiled.nsplits = new_nsplits

    def fetch(self, *tileables, n_parallel: int = None, **kw):
        if self._executor is None:
            raise RuntimeError('Session has closed')
        if n_parallel is None:
            kw['n_parallel'] = cpu_count()
        return self._executor.fetch_tileables(tileables, **kw)

    def fetch_log(self, tileables, offsets=None, sizes=None):  # pragma: no cover
        raise NotImplementedError('`fetch_log` is not implemented for local executor')

    def create_mutable_tensor(self, name, shape, dtype, fill_value=None, *args, **kwargs):
        from .tensor.core import MutableTensor, MutableTensorData
        if name in self._mut_tensor:
            raise ValueError(f"The mutable tensor named '{name}' already exists.")
        mut_tensor = MutableTensor(data=MutableTensorData(name=name, op=None,
                                                          shape=shape, dtype=np.dtype(dtype)))
        self._mut_tensor[name] = mut_tensor
        if fill_value is None:
            self._mut_tensor_data[name] = np.zeros(shape, dtype=dtype)
        else:
            self._mut_tensor_data[name] = np.full(shape, fill_value, dtype=dtype)
        return mut_tensor

    def get_mutable_tensor(self, name):
        if name not in self._mut_tensor:
            raise ValueError(f"The mutable tensor named '{name}' doesn't exist, or has already been sealed.")
        return self._mut_tensor.get(name)

    def write_mutable_tensor(self, tensor, index, value):
        if tensor.name not in self._mut_tensor:
            raise ValueError(f"The mutable tensor named '{tensor.name}' doesn't exist, or has already been sealed.")
        tensor_data = self._mut_tensor_data.get(tensor.name)
        tensor_data[index] = value

    def seal(self, tensor):
        from .tensor.datasource.array import ArrayDataSource
        if tensor.name not in self._mut_tensor:
            raise ValueError(f"The mutable tensor named '{tensor.name}' doesn't exist, or has already been sealed.")
        mut_tensor = self._mut_tensor.pop(tensor.name)
        tensor_data = self._mut_tensor_data.pop(tensor.name)
        sealed_tensor = ArrayDataSource(tensor_data, dtype=mut_tensor.dtype)(shape=mut_tensor.shape)
        self._executor.execute_tileables([sealed_tensor])
        return sealed_tensor

    def get_named_tileable_infos(self, name):
        return self._context.get_named_tileable_infos(name)

    def decref(self, *keys):
        self._executor.decref(*keys)

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self._executor = None


class ClusterSession(object):
    def __init__(self, endpoint, session_id=None, **kwargs):
        from .api import MarsAPI
        from .context import DistributedContext

        self._endpoint = endpoint
        # dict structure: {tileable_key -> graph_key, tileable_ids}
        # dict value is a tuple object which records graph key and tilable id
        self._executed_tileables = dict()
        self._api = MarsAPI(self._endpoint)

        if session_id is None:
            # create session on the cluster side
            self._session_id = uuid.uuid4()
            self._api.create_session(self._session_id)
        else:
            # Get the session actor ref using given session_id
            self._session_id = session_id
            if not self._api.has_session(self._session_id):
                raise ValueError(f'The session with id = {self._session_id} doesn\'t exist')

        self._context = DistributedContext(endpoint, self._session_id)

        if kwargs:
            unexpected_keys = ', '.join(list(kwargs.keys()))
            raise TypeError(f'Local cluster session got unexpected arguments: {unexpected_keys}')

    @property
    def session_id(self):
        return self._session_id

    @property
    def endpoint(self):
        return self._endpoint

    @endpoint.setter
    def endpoint(self, endpoint):
        from .api import MarsAPI

        self._endpoint = endpoint
        self._api = MarsAPI(self._endpoint)

    @property
    def context(self):
        return self._context

    def get_cpu_count(self):
        return self._context.get_total_ncores()

    def _get_tileable_graph_key(self, tileable_key):
        return self._executed_tileables[tileable_key][0]

    def _set_tileable_graph_key(self, tileable, graph_key):
        tileable_key = tileable.key
        tileable_id = tileable.id
        if tileable_key in self._executed_tileables:
            self._executed_tileables[tileable_key][1].add(tileable_id)
        else:
            self._executed_tileables[tileable_key] = graph_key, {tileable_id}

    def _update_tileable_shape(self, tileable):
        graph_key = self._get_tileable_graph_key(tileable.key)
        new_nsplits = self._api.get_tileable_nsplits(self._session_id, graph_key, tileable.key)
        tileable._update_shape(tuple(sum(nsplit) for nsplit in new_nsplits))
        tileable.nsplits = new_nsplits

    def create_mutable_tensor(self, name, shape, dtype, *args, **kwargs):
        from .tensor.utils import create_mutable_tensor
        shape, dtype, chunk_size, chunk_keys, chunk_eps = \
            self._api.create_mutable_tensor(self._session_id, name, shape,
                                            dtype, *args, **kwargs)
        return create_mutable_tensor(name, chunk_size, shape, dtype, chunk_keys, chunk_eps)

    def get_mutable_tensor(self, name):
        from .tensor.utils import create_mutable_tensor
        shape, dtype, chunk_size, chunk_keys, chunk_eps = \
            self._api.get_mutable_tensor(self._session_id, name)
        return create_mutable_tensor(name, chunk_size, shape, dtype, chunk_keys, chunk_eps)

    def write_mutable_tensor(self, tensor, index, value):
        chunk_records_to_send = tensor._do_write(index, value)
        self._api.send_chunk_records(self._session_id, tensor.name, chunk_records_to_send)

    def get_workers_meta(self):
        return self._api.get_workers_meta()

    def seal(self, tensor):
        from .tensor.utils import create_fetch_tensor
        chunk_records_to_send = tensor._do_flush()
        self._api.send_chunk_records(self._session_id, tensor.name, chunk_records_to_send)

        graph_key_hex, tensor_key, tensor_id, tensor_meta = self._api.seal(self._session_id, tensor.name)
        self._executed_tileables[tensor_key] = uuid.UUID(graph_key_hex), {tensor_id}

        # Construct Tensor on the fly.
        shape, dtype, chunk_size, chunk_keys, _ = tensor_meta
        return create_fetch_tensor(chunk_size, shape, dtype, tensor_key=tensor_key, chunk_keys=chunk_keys)

    def run(self, *tileables, **kw):
        from .utils import build_tileable_graph
        from .scheduler.graph import GraphState
        from .errors import ExecutionFailed

        timeout = kw.pop('timeout', -1)
        fetch = kw.pop('fetch', True)
        compose = kw.pop('compose', True)
        name = kw.pop('name', None)
        if not isinstance(name, (tuple, list)):
            name = [name]
        if kw:
            raise TypeError(f'run got unexpected key arguments {kw!r}')

        # those executed tileables should fetch data directly, submit the others
        run_tileables = [t for t in tileables if t.key not in self._executed_tileables]

        graph = build_tileable_graph(run_tileables, set(self._executed_tileables.keys()))
        targets = [t.key for t in run_tileables]
        graph_key = uuid.uuid4()

        if len(graph) > 0:
            # submit graph to local cluster
            self._api.submit_graph(self._session_id, serialize_graph(graph),
                                   graph_key, targets, compose=compose, names=name)

            ctx = get_context()
            yield_info = None
            if ctx is not None:
                yield_info = ctx.yield_execution_pool()

            try:
                exec_start_time = time.time()
                time_elapsed = 0
                check_interval = options.check_interval
                while timeout <= 0 or time_elapsed < timeout:
                    timeout_val = min(check_interval, timeout - time_elapsed) if timeout > 0 else check_interval
                    self._api.wait_graph_finish(self._session_id, graph_key, timeout=timeout_val)
                    graph_state = self._api.get_graph_state(self._session_id, graph_key)
                    if graph_state == GraphState.SUCCEEDED:
                        break
                    if graph_state == GraphState.FAILED:
                        exc_info = self._api.get_graph_exc_info(self._session_id, graph_key)
                        if exc_info is not None:
                            exc = exc_info[1].with_traceback(exc_info[2])
                            raise ExecutionFailed('Graph execution failed.') from exc
                        else:
                            raise ExecutionFailed('Graph execution failed with unknown reason')
                    time_elapsed = time.time() - exec_start_time

                if 0 < timeout < time.time() - exec_start_time:
                    raise TimeoutError
            finally:
                if ctx is not None:
                    ctx.acquire_execution_pool(yield_info)

            for t in tileables:
                self._set_tileable_graph_key(t, graph_key)

        if not fetch:
            return
        else:
            return self.fetch(*tileables)

    def _is_executed(self, tileable):
        # if tileble.key in executed tileables
        # or it's a fetch already
        return tileable.key in self._executed_tileables or \
               isinstance(tileable.op, Fetch)

    def fetch(self, *tileables):
        from .utils import sort_dataframe_result
        from .serialize import dataserializer
        from .tensor.indexing import TensorIndex
        from .dataframe.indexing.iloc import DataFrameIlocGetItem, SeriesIlocGetItem

        tileable_results = []
        for tileable in tileables:
            # TODO: support DataFrame getitem
            if not self._is_executed(tileable) and \
                    isinstance(tileable.op, (TensorIndex, DataFrameIlocGetItem, SeriesIlocGetItem)):
                to_fetch_tileable = tileable.inputs[0]

                indexes = tileable.op.indexes
                if not all(isinstance(ind, (slice, Integral)) for ind in indexes):
                    raise ValueError('Only support fetch data slices')
            else:
                to_fetch_tileable = tileable
                indexes = None
            if not self._is_executed(to_fetch_tileable):
                raise ValueError('Cannot fetch the unexecuted tileable')

            key = to_fetch_tileable.key
            compressions = dataserializer.get_supported_compressions()
            if getattr(to_fetch_tileable, 'chunks', None) is not None:
                # fetch which owns fetch chunks inside remote
                chunk_indexes = [c.index for c in to_fetch_tileable.chunks]
                chunk_keys = [c.key for c in to_fetch_tileable.chunks]
                nsplits = to_fetch_tileable.nsplits
                if any(np.isnan(np.sum(ns)) for ns in nsplits):
                    # unknown chunk shape exists
                    # try to calculate nsplits
                    chunk_metas = self._api.get_chunk_metas(self._session_id, chunk_keys)
                    chunk_idx_to_shape = {idx: cm.chunk_shape for idx, cm in
                                          zip(chunk_indexes, chunk_metas)}
                    nsplits = calc_nsplits(chunk_idx_to_shape)
                result = self._api.fetch_chunks_data(
                    self._session_id, chunk_indexes, chunk_keys, nsplits,
                    serial=False, index_obj=indexes, compressions=compressions)
            else:
                result = self._api.fetch_data(self._session_id, key, serial=False,
                                              index_obj=indexes, compressions=compressions)
            tileable_results.append(sort_dataframe_result(tileable, result))
        return tileable_results

    def fetch_tileable_op_logs(self, tileable_op_key, offsets=None, sizes=None):
        return self._context.fetch_tileable_op_logs(
            tileable_op_key, chunk_op_key_to_offsets=offsets,
            chunk_op_key_to_sizes=sizes)

    def fetch_log(self, tileables, offsets=None, sizes=None):
        from .custom_log import fetch

        return fetch(tileables, self, offsets=offsets, sizes=sizes)

    def decref(self, *keys):
        for tileable_key, tileable_id in keys:
            if tileable_key not in self._executed_tileables:
                continue
            _graph_key, ids = self._executed_tileables[tileable_key]
            if tileable_id in ids:
                ids.remove(tileable_id)
                # for those same key tileables, do decref only when all those tileables are garbage collected
                if len(ids) != 0:
                    continue
                self.delete_data(tileable_key)

    def delete_data(self, tileable_key, wait=False):
        if tileable_key not in self._executed_tileables:
            return
        graph_key, _ids = self._executed_tileables[tileable_key]
        self._api.delete_data(self._session_id, graph_key, tileable_key, wait=wait)
        del self._executed_tileables[tileable_key]

    def get_named_tileable_infos(self, name):
        return self._context.get_named_tileable_infos(name)

    def get_tileable_chunk_metas(self, tileable_key):
        return self._api.get_tileable_chunk_metas(self._session_id, tileable_key)

    def __enter__(self):
        return self

    def __exit__(self, *_):
        for key in list(self._executed_tileables.keys()):
            self.delete_data(key, wait=True)
        self._api.delete_session(self._session_id)


class Session(object):
    _default_session_local = threading.local()

    def __init__(self, endpoint=None, backend=None, **kwargs):
        self._endpoint = endpoint
        self._backend = backend
        self._kws = kwargs
        self._init()

    def _init(self):
        endpoint, kwargs = self._endpoint, self._kws
        if self._backend is None:
            if endpoint is not None:
                if 'http' in endpoint:
                    # connect to web
                    self._init_web_session(endpoint, **kwargs)
                else:
                    # connect to local cluster
                    self._init_cluster_session(endpoint, **kwargs)
            else:
                try:
                    endpoint = os.environ['MARS_SCHEDULER_ADDRESS']
                    session_id = os.environ.get('MARS_SESSION_ID', None)
                    kwargs['session_id'] = session_id
                    self._init_cluster_session(endpoint, **kwargs)
                except KeyError:
                    self._init_local_session(**kwargs)
        elif self._backend == 'ray':
            self._init_ray_session(**kwargs)
        else:  # pragma: no cover
            raise ValueError('Either endpoint or backend should '
                             'be provided to create a session')

    def _init_local_session(self, **kwargs):
        self._sess = LocalSession(**kwargs)

    def _init_web_session(self, endpoint, **kwargs):
        from .web.session import Session as WebSession

        self._sess = WebSession(endpoint, **kwargs)

    def _init_cluster_session(self, endpoint, **kwargs):
        self._sess = ClusterSession(endpoint, **kwargs)

    def _init_ray_session(self, **kwargs):
        from .ray.core import RaySession

        self._sess = RaySession(**kwargs)

    def __getstate__(self):
        return self._endpoint, self._kws, self.session_id, self._backend

    def __setstate__(self, state):
        self._endpoint, self._kws, session_id, self._backend = state
        self._init()
        self._sess._session_id = session_id

    def run(self, *tileables, **kw):
        from . import tensor as mt

        fetch = kw.get('fetch', True)
        ret_list = False
        if len(tileables) == 1 and isinstance(tileables[0], (tuple, list)):
            ret_list = True
            tileables = tileables[0]
        elif len(tileables) > 1:
            ret_list = True

        tileables = tuple(mt.tensor(t) if not isinstance(t, ENTITY_TYPE) else t
                          for t in tileables)
        result = self._sess.run(*tileables, **kw)

        for t in tileables:
            t._attach_session(self)

        for t in tileables:
            if getattr(t, 'shape', None) is not None and \
                    any(np.isnan(s) for s in t.shape):
                self._sess._update_tileable_shape(t)

        if fetch:
            ret = []
            for r, t in zip(result, tileables):
                if hasattr(t, 'isscalar') and t.isscalar() and getattr(r, 'size', None) == 1:
                    ret.append(r.item())
                else:
                    ret.append(r)
            if ret_list:
                return ret
            return ret[0]

    def fetch(self, *tileables, **kw):
        ret_list = False
        if len(tileables) == 1 and isinstance(tileables[0], (tuple, list)):
            ret_list = True
            tileables = tileables[0]
        elif len(tileables) > 1:
            ret_list = True

        result = self._sess.fetch(*tileables, **kw)

        ret = []
        for r, t in zip(result, tileables):
            if hasattr(t, 'isscalar') and t.isscalar() and getattr(r, 'size', None) == 1:
                ret.append(r.item())
            else:
                ret.append(r)
        if ret_list:
            return ret
        return ret[0]

    @property
    def endpoint(self):
        return self._sess.endpoint

    @endpoint.setter
    def endpoint(self, endpoint):
        self._sess.endpoint = endpoint

    @property
    def session_id(self):
        return self._sess.session_id

    def get_cpu_count(self):
        return self._sess.get_cpu_count()

    def decref(self, *keys):
        if hasattr(self._sess, 'decref'):
            self._sess.decref(*keys)

    def __getattr__(self, attr):
        obj = self._sess.__getattribute__(attr)
        return obj

    def __enter__(self):
        self._sess.__enter__()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._sess.__exit__(exc_type, exc_val, exc_tb)

    def close(self):
        self.__exit__(None, None, None)

    def as_default(self):
        Session._default_session_local.default_session = self
        return self

    @classmethod
    def _set_default_session(cls, session):
        cls._default_session_local.default_session = session
        return session

    @classproperty
    def default(self):
        return getattr(Session._default_session_local, 'default_session', None)

    @classmethod
    def default_or_local(cls):
        default_session = getattr(Session._default_session_local, 'default_session', None)
        if default_session is not None:
            return default_session

        return cls._set_default_session(Session())

    def create_mutable_tensor(self, name, shape, dtype, fill_value=None, *args, **kwargs):
        return self._sess.create_mutable_tensor(name, shape, dtype, fill_value=fill_value, *args, **kwargs)

    def get_mutable_tensor(self, name):
        return self._sess.get_mutable_tensor(name)

    def write_mutable_tensor(self, tensor, index, value):
        self._sess.write_mutable_tensor(tensor, index, value)

    def seal(self, tensor):
        return self._sess.seal(tensor)


def new_session(scheduler=None, backend=None, **kwargs):
    return Session(scheduler, backend=backend, **kwargs)
