# Copyright (c) Metakernel Development Team.
# Distributed under the terms of the Modified BSD License.
from __future__ import annotations

import time
from typing import Any

from metakernel import Magic, MetaKernel, option


class Slice:
    """Utility class for making slice ranges."""

    def __getitem__(self, item: Any) -> Any:
        return item


slicer = Slice()  ## instance to parse slices


class ParallelMagic(Magic):
    client: Any = None
    view: Any = None
    view_load_balanced: Any = None
    module_name: Any = None
    class_name: Any = None
    kernel_name: Any = None
    ids: Any = None
    retval: Any = None
    retry = False

    @option(
        "-k",
        "--kernel_name",
        action="store",
        default="default",
        help="arbitrary name given to reference kernel",
    )
    @option(
        "-i",
        "--ids",
        action="store",
        default=None,
        help="the machine ids to use from the cluster",
    )
    def line_parallel(
        self,
        module_name: str,
        class_name: str,
        kernel_name: str = "default",
        ids: str | None = None,
    ) -> None:
        """
        %parallel MODULE CLASS [-k NAME] [-i [...]] - construct an interface to the cluster.

        Example:

            %parallel bash_kernel BashKernel
            %parallel bash_kernel BashKernel -k bash
            %parallel bash_kernel BashKernel -i [0,2:5,9,...]

        cluster_size and cluster_rank variables are set upon
        initialization of the remote node (if the kernel
        supports %set).

        Use %px or %%px to send code to the cluster.
        """
        from ipyparallel import Client  # type: ignore[import-untyped]

        count = 1
        while count <= 5:
            try:
                self.client = Client()
                break
            except Exception:
                print("Waiting on cluster to start...")
                time.sleep(2)
            count += 1
        if count == 6:
            raise Exception("Cluster was not started.")
        if ids is None:
            count = 1
            while count <= 5:
                try:
                    self.view = self.client[:]
                    break
                except Exception:
                    print("Waiting for engines...")
                    time.sleep(2)
                count += 1
            if count == 6:
                raise Exception("Engines were not started.")
        else:
            # ids[:] = slice(None, None, None)
            # ids[1:3] = slice(1, 3, None)
            # ids[1:3:1] = slice(1, 3, 1)
            # ids[1, 2, ...] = [1, 2, Ellipsis]
            # ids[1, 2:4, ...] = [1, slice(2, 4, None), Ellipsis]
            try:
                ids_slice = eval(f"slicer{ids}")  # slicer[0,...,7]
            except Exception:
                ids_slice = slicer[:]
            if isinstance(ids_slice, (slice, int)):
                count = 1
                while count <= 5:
                    try:
                        self.view = self.client[ids_slice]
                        break
                    except Exception:
                        print("Waiting for engines...")
                        time.sleep(2)
                    count += 1
                if count == 6:
                    raise Exception("Engines were not started.")
            else:  # tuple of indexes/slices
                # FIXME: if so, handle Ellipsis
                view = None
                for item in ids_slice:
                    count = 1
                    while count <= 5:
                        try:
                            client = self.client[item]
                            if view:
                                ## FIXME: can't do this:
                                view.append(client)
                            else:
                                view = client
                            break
                        except Exception:
                            print("Waiting on cluster to start...")
                            time.sleep(2)
                        count += 1
                    if count == 6:
                        raise Exception("Cluster was not started.")
                self.view = view
        self.view_load_balanced = self.client.load_balanced_view()
        self.module_name = module_name
        self.class_name = class_name
        self.kernel_name = kernel_name
        self.view.execute(
            f"""
import os
for key, value in {self.kernel.env!s}.items():
    os.environ[key] = value
try:
    kernels
except Exception:
    kernels = {{}}
from {module_name} import {class_name}
kernels['{kernel_name}'] = {class_name}()
## FIXME: kernels['{kernel_name}'].kernel = kernel
""",
            block=True,
        )

        self.view[
            f"kernels['{kernel_name}'].set_variable(\"cluster_size\", {len(self.client)})"
        ]
        self.client[:].scatter("cluster_rank", self.client.ids, flatten=True)
        self.view[
            f"kernels['{kernel_name}'].set_variable(\"cluster_rank\", cluster_rank)"
        ]
        ## So that these are available in the host kernel:
        self.kernel.set_variable("cluster_size", len(self.client))
        self.kernel.set_variable("cluster_rank", -1)
        self.retval = None

    @option(
        "-k",
        "--kernel_name",
        action="store",
        default=None,
        help="kernel name given to use for execution",
    )
    @option(
        "-e",
        "--evaluate",
        action="store_true",
        default=False,
        help=(
            "evaluate code in the current kernel, too. The current "
            + "kernel should be of the same language as the cluster."
        ),
    )
    @option(
        "-s",
        "--set_variable",
        action="store",
        default=None,
        help="set the variable with the parallel results rather than returning them",
    )
    def line_px(
        self,
        expression: str,
        kernel_name: str | None = None,
        evaluate: bool = False,
        set_variable: str | None = None,
    ) -> None:
        """
        %px EXPRESSION - send EXPRESSION to the cluster.

        Example:

            %px sys.version
            %px -k scheme (define x 42)
            %px x
            %px cluster_rank

        cluster_size and cluster_rank variables are set upon
        initialization of the remote node (if the kernel
        supports %set).

        Use %parallel to initialize the cluster.
        """
        results = None
        expression = str(expression)
        if kernel_name is None:
            kernel_name = self.kernel_name
        if self.retry:
            count = 1
            while count <= 5:
                try:
                    results = self.view[
                        f"kernels['{kernel_name}'].do_execute_direct(\"{self._clean_code(expression)}\")"
                    ]
                    break
                except Exception:
                    print("Waiting on cluster clients to start...")
                    time.sleep(2)
                count += 1
            if count == 6:
                raise Exception("Cluster clients have not started.")
            self.retry = False
        else:
            try:
                results = self.view[
                    f"kernels['{kernel_name}'].do_execute_direct(\"{self._clean_code(expression)}\")"
                ]
            except Exception as e:
                results = str(e)
        if set_variable is None:
            self.retval = results
        else:
            self.kernel.set_variable(set_variable, results)
            self.retval = None
        if evaluate:
            self.code = expression

    def _clean_code(self, expr: str) -> str:
        return expr.strip().replace('"', '\\"').replace("\n", "\\n")

    ## px --kernel NAME
    @option(
        "-k",
        "--kernel_name",
        action="store",
        default=None,
        help="kernel name given to use for execution",
    )
    @option(
        "-e",
        "--evaluate",
        action="store_true",
        default=False,
        help=(
            "evaluate code in the current kernel, too. The current "
            + "kernel should be of the same language as the cluster."
        ),
    )
    @option(
        "-s",
        "--set_variable",
        action="store",
        default=None,
        help="set the variable with the parallel results rather than returning them",
    )
    def cell_px(
        self,
        kernel_name: str | None = None,
        evaluate: bool = False,
        set_variable: str | None = None,
    ) -> None:
        """
        %%px - send cell to the cluster.

        Example:

            %%px
            (define x 42)

        Use %parallel to initialize the cluster.
        """
        if kernel_name is None:
            kernel_name = self.kernel_name
        results = self.view[
            f"kernels['{kernel_name}'].do_execute_direct(\"{self._clean_code(self.code)}\")"
        ]
        if set_variable is None:
            self.retval = results
        else:
            self.kernel.set_variable(set_variable, results)
            self.retval = None
        self.evaluate = evaluate

    @option(
        "-s",
        "--set_variable",
        action="store",
        default=None,
        help="set the variable with the parallel results rather than returning them",
    )
    def line_pmap(
        self,
        function_name: str,
        args: str,
        kernel_name: str | None = None,
        set_variable: str | None = None,
    ) -> None:
        """
        %pmap FUNCTION [ARGS1,ARGS2,...] - ("parallel map") call a FUNCTION on args

        This line magic will apply a function name to all of the
        arguments given one at a time using a dynamic load balancing scheduler.

        Currently, the args are provided as a Python expression (with no spaces).

        You must first setup a cluster using the %parallel magic.

        Examples:

            %pmap function-name-in-language range(10)
            %pmap function-name-in-language [1,2,3,4]
            %pmap run_experiment range(1,100,5)
            %pmap run_experiment ["test1","test2","test3"]
            %pmap f [(1,4,7),(2,3,5),(7,2,2)]

        The function name must be a function that is available on all
        nodes in the cluster. For example, you could:

            %%px
            (define myfunc
               (lambda (n)
                 (+ n 1)))

        to define myfunc on all machines (use %%px -e to also define
        it in the running notebook or console). Then you can apply it
        to a list of arguments:

            %%pmap myfunc range(100)

        The load balancer will run myfunc on the next available node
        in the cluster.

        Note: not all languages may support running a function via this magic.
        """
        if kernel_name is None:
            kernel_name = self.kernel_name

        # To make sure we can find `kernels`:
        try:
            from ipyparallel.util import interactive  # type: ignore[import-untyped]
        except ImportError:
            from IPython.parallel.util import (  # type: ignore[import-not-found]
                interactive,
            )
        f = interactive(
            lambda arg, kname=kernel_name, fname=function_name: kernels[  # type:ignore[name-defined]  # noqa: F821
                kname
            ].do_function_direct(fname, arg)
        )
        results = self.view_load_balanced.map_async(f, eval(args))
        if set_variable is None:
            self.retval = results
        else:
            self.kernel.set_variable(set_variable, results)
            self.retval = None

    def post_process(self, retval: Any) -> Any:
        try:
            ## any will crash on numpy arrays
            if isinstance(self.retval, list) and not any(self.retval):
                return None
        except Exception:
            pass
        return self.retval


def register_magics(kernel: MetaKernel) -> None:
    kernel.register_magics(ParallelMagic)
