# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.
"""Tornado handlers for invoking notebook cells as web APIs."""

import json
import os
from functools import partial
from typing import Optional

import tornado.web
from tornado.concurrent import Future
from tornado.log import access_log

from ..mixins import CORSMixin, JSONErrorsMixin, TokenAuthorizationMixin
from .errors import CodeExecutionError, UnsupportedMethodError
from .request_utils import (
    format_request,
    headers_to_dict,
    parameterize_path,  # noqa: F401
    parse_args,
    parse_body,
)


class NotebookAPIHandler(
    TokenAuthorizationMixin, CORSMixin, JSONErrorsMixin, tornado.web.RequestHandler
):
    """Executes code from a notebook cell in response to HTTP requests at the
    route registered in association with this class.

    Supports the GET, POST, PUT, and DELETE HTTP methods.

    Parameters
    ----------
    sources : dict
        Maps HTTP verb strings to annotated cells extracted from a notebook
    response_sources
        Maps HTTP verb strings to ResponseInfo annotated cells extracted from a
        notebook
    kernel_pool
        Instance of services.kernels.ManagedKernelPool
    kernel_name
        Kernel spec name used to launch the kernel pool
    kernel_language
        Kernel spec language used to make language specific operations

    Attributes
    ----------
    See parameters: they are stored as passed as instance variables.

    See Also
    --------
    services.cell.parser.APICellParser for detail about how the source cells
    are identified, parsed, and associated with HTTP verbs and paths.
    """

    def initialize(self, sources, response_sources, kernel_pool, kernel_name, kernel_language=""):
        self.kernel_pool = kernel_pool
        self.sources = sources
        self.kernel_name = kernel_name
        self.response_sources = response_sources
        self.kernel_language = kernel_language

    def finish_future(self, future, result_accumulator):
        """Resolves the promise to respond to a HTTP request handled by a
        kernel in the pool.

        Defines precedence for the kind of response:

        1. If any error occurred, resolve with a CodeExecutionError exception
        2. If any stream message was collected from the kernel, resolve with
            the joined string of stream messages
        3. If an execute_result was collected, resolve with that result object
            JSON encoded
        4. Resolve with an empty string

        Parameters
        ----------
        future : tornado.concurrent.Future
            Promise of a future response to an API request
        result_accumulator : dict
            Dictionary of results from a kernel with at least the keys error,
            stream, and result
        """
        if result_accumulator["error"]:
            future.set_exception(CodeExecutionError(result_accumulator["error"]))
        elif len(result_accumulator["stream"]) > 0:
            future.set_result("".join(result_accumulator["stream"]))
        elif result_accumulator["result"]:
            future.set_result(json.dumps(result_accumulator["result"]))
        else:
            # If nothing was set, return an empty value
            future.set_result("")

    def on_recv(self, result_accumulator, future, parent_header, msg):
        """Collects ipoub messages associated with code execution request
        identified by `parent_header`.

        Continues collecting until an execution state of idle is reached.
        The first three parameters are typically applied in a partial.

        Parameters
        ----------
        result_accumulator : dict
            Accumulates data from `execute_result`, `stream` and `error`
            messages under keys `result`, `stream`, and `error` respectively
            across multiple invocations of this callback.
        future : tornado.concurrent.Future
            Promise to resolve when the kernel goes idle
        parent_header : dict
            Parent header from an `execute` request, used to identify messages
            that relate to its execution vs other executions
        msg : dict
            Kernel message received from the iopub channel
        """
        if msg["parent_header"]["msg_id"] == parent_header:
            # On idle status, exit our loop
            if (
                msg["header"]["msg_type"] == "status"
                and msg["content"]["execution_state"] == "idle"
            ):
                self.finish_future(future, result_accumulator)
            # Store the execute result
            elif msg["header"]["msg_type"] == "execute_result":
                result_accumulator["result"] = msg["content"]["data"]
            # Accumulate the stream messages
            elif msg["header"]["msg_type"] == "stream":
                # Only take stream output if it is on stdout or if the kernel
                # is non-confirming and does not name the stream
                if "name" not in msg["content"] or msg["content"]["name"] == "stdout":
                    result_accumulator["stream"].append(msg["content"]["text"])
            # Store the error message
            elif msg["header"]["msg_type"] == "error":
                error_name = msg["content"]["ename"]
                error_value = msg["content"]["evalue"]
                result_accumulator["error"] = f"Error {error_name}: {error_value} \n"

    def execute_code(self, kernel_client, kernel_id, source_code):
        """Executes `source_code` on the kernel specified.

        Registers a callback for iopub messages. Promises to return the output
        of the execution in the future after the kernel returns to its idle
        state.

        Parameters
        ----------
        kernel_client : object
            Client to use to execute the code
        kernel_id : str
            ID of the kernel from the pool that will execute the request
        source_code : str
            Source code to execute

        Returns
        -------
        tornado.concurrent.Future
            Promise of execution result

        Raises
        ------
        CodeExecutionError
            If the kernel returns any error
        """
        future = Future()
        result_accumulator = {"stream": [], "error": None, "result": None}
        parent_header = kernel_client.execute(source_code)
        on_recv_func = partial(self.on_recv, result_accumulator, future, parent_header)
        self.kernel_pool.on_recv(kernel_id, on_recv_func)
        return future

    async def _handle_request(self):
        """Turns an HTTP request into annotated notebook code to execute on a
        kernel.

        Sets the HTTP response code, headers, and response body based on the
        result of the kernel execution. Then finishes the Tornado response.
        """
        self.response_future = Future()
        kernel_client, kernel_id = await self.kernel_pool.acquire()
        try:
            # Method not supported
            if self.request.method not in self.sources:
                raise UnsupportedMethodError(self.request.method)

            # Set the Content-Type and status to default values
            self.set_header("Content-Type", "text/plain")
            self.set_status(200)

            # Get the source to execute in response to this request
            source_code = self.sources[self.request.method]
            # Build the request dictionary
            request = json.dumps(
                {
                    "body": parse_body(self.request),
                    "args": parse_args(self.request.query_arguments),
                    "path": self.path_kwargs,
                    "headers": headers_to_dict(self.request.headers),
                }
            )
            # Turn the request string into a valid code string
            request_code = format_request(request, self.kernel_language)

            # Run the request and source code and yield until there's a result
            access_log.debug(f"Request code for notebook cell is: {request_code}")
            await self.execute_code(kernel_client, kernel_id, request_code)
            source_result = await self.execute_code(kernel_client, kernel_id, source_code)

            # If a response code cell exists, execute it
            if self.request.method in self.response_sources:
                response_code = self.response_sources[self.request.method]
                response_future = self.execute_code(kernel_client, kernel_id, response_code)

                # Wait for the response and parse the json value
                response_result = await response_future
                response = json.loads(response_result)

                # Copy all the header values into the tornado response
                if "headers" in response:
                    for header in response["headers"]:
                        self.set_header(header, response["headers"][header])

                # Set the status code if it exists
                if "status" in response:
                    self.set_status(response["status"])

            # Write the result of the source code execution
            if source_result:
                self.write(source_result)
        # If there was a problem executing an code, return a 500
        except CodeExecutionError as err:
            self.write(str(err))
            self.set_status(500)
        # An unsupported method was called on this handler
        except UnsupportedMethodError:
            self.set_status(405)
        finally:
            # Always make sure we release the kernel and finish the request
            self.response_future.set_result(None)
            self.kernel_pool.release(kernel_id)
            await self.finish()

    async def get(self, **kwargs):
        await self._handle_request()
        await self.response_future

    async def post(self, **kwargs):
        await self._handle_request()
        await self.response_future

    async def put(self, **kwargs):
        await self._handle_request()
        await self.response_future

    async def delete(self, **kwargs):
        await self._handle_request()
        await self.response_future

    def options(self, **kwargs):
        self.finish()


class NotebookDownloadHandler(
    TokenAuthorizationMixin, CORSMixin, JSONErrorsMixin, tornado.web.StaticFileHandler
):
    """Handles requests to download the annotated notebook behind the web API."""

    def initialize(self, path: str, default_filename: Optional[str] = None):
        self.dirname, self.filename = os.path.split(path)
        super().initialize(self.dirname)

    async def get(self, include_body: bool = True):
        res = await super().get(path=self.filename, include_body=include_body)
        return res
