from __future__ import annotations

import asyncio
import inspect
import logging
import os
import sys
from typing import Literal

from . import utils
from .controller import Controller
from .http import HttpHeader
from .protocol import CoreServer
from .state import State
from .ui import VirtualNodeManager
from .utils import share
from .utils.argument_parser import ArgumentParser
from .utils.namespace import Translator

logger = logging.getLogger(__name__)

ClientType = Literal["vue2", "vue3"]
BackendType = Literal["aiohttp", "generic", "tornado", "jupyter"]
ExecModeType = Literal["main", "desktop", "task", "coroutine"]

DEFAULT_CLIENT_TYPE: ClientType = "vue3"


def set_default_client_type(value: ClientType) -> None:
    global DEFAULT_CLIENT_TYPE  # noqa: PLW0603
    DEFAULT_CLIENT_TYPE = value


class Server:
    """
    Server implementation for trame.
    This is the core object that manage client/server communication but also
    holds a state and controller instance.
    With trame a server instance should be retrieved by using **trame.app.get_server()**

    Known options:
      - log_network: False (path to log file)
      - ws_max_msg_size: 10000000 (bytes)
      - ws_heart_beat: 30
      - desktop_debug: False

    :param name: A name identifier for a given server
    :type name: str, optional (default: trame)

    :param **options: Gather any keyword arguments into options
    :type options: Dict
    """

    def __init__(
        self,
        name="trame",
        vn_constructor=None,
        translator=None,
        parent_server=None,
        **options,
    ) -> None:
        # Core internal variables
        self._parent_server = parent_server
        self._translator = translator or Translator()
        self._name = share(parent_server, "_name", name)
        self._options = share(parent_server, "_options", options)
        self._client_type = share(parent_server, "_client_type", None)
        self._http_header = share(parent_server, "_http_header", HttpHeader())

        # use parent_server instead of local version
        self._server = None
        self._running_stage = 0  # 0: off / 1: pending / 2: running
        self._running_port = 0
        self._running_future = None
        self._www = None
        self.serve = {}  # HTTP static endpoints
        self._loaded_modules = set()
        self._loaded_module_dicts = []
        self._cli_parser = None
        self._root_protocol = None
        self._protocols_to_configure = []

        # ENV variable mapping settings
        self.hot_reload = "--hot-reload" in sys.argv or bool(
            os.getenv("TRAME_HOT_RELOAD", None)
        )
        if parent_server is None:
            self._options["log_network"] = self._options.get(
                "log_network", os.environ.get("TRAME_LOG_NETWORK")
            )
            self._options["ws_max_msg_size"] = self._options.get(
                "ws_max_msg_size", os.environ.get("TRAME_WS_MAX_MSG_SIZE") or 10000000
            )
            self._options["ws_heart_beat"] = self._options.get(
                "ws_heart_beat", os.environ.get("TRAME_WS_HEART_BEAT") or 30
            )
            self._options["desktop_debug"] = self._options.get(
                "desktop_debug", os.environ.get("TRAME_DESKTOP_DEBUG")
            )
            # reset default wslink startup message
            os.environ["WSLINK_READY_MSG"] = ""

        # Shared state + reserve internal keys
        if parent_server is None:
            self._state = State(
                self.translator, commit_fn=self._push_state, hot_reload=self.hot_reload
            )
            for key in ["scripts", "module_scripts", "styles", "vue_use", "mousetrap"]:
                self._state[f"trame__{key}"] = []
            self._state.trame__client_only = ["trame__busy"]
            self._state.trame__busy = 1
            self._state.trame__favicon = None
            self._state.trame__title = "Trame"
        else:
            self._state = State(
                self.translator,
                internal=parent_server._state,
                commit_fn=self._push_state,
                hot_reload=self.hot_reload,
            )

        # Controller
        if parent_server is None:
            self._controller = Controller(self.translator, hot_reload=self.hot_reload)
        else:
            self._controller = Controller(
                self.translator,
                internal=parent_server._controller,
                hot_reload=self.hot_reload,
            )

        # Server only context
        if parent_server is None:
            self._context = State(self.translator, hot_reload=self.hot_reload)
        else:
            self._context = State(
                self.translator,
                internal=parent_server._context,
                hot_reload=self.hot_reload,
            )

        # UI (FIXME): use for translator
        self._ui = share(parent_server, "_ui", VirtualNodeManager(self, vn_constructor))

    def create_child_server(self, translator=None, prefix=None) -> Server:
        translator = translator or Translator(prefix=prefix)
        return Server(translator=translator, parent_server=self)

    # -------------------------------------------------------------------------
    # State management helpers
    # -------------------------------------------------------------------------

    def _push_state(self, state):
        if self.protocol:
            self.protocol.push_state_change(state)

    # -------------------------------------------------------------------------
    # Initialization helper
    # -------------------------------------------------------------------------

    @property
    def http_headers(self):
        """Return http header helper so they can be applied before the server start."""
        return self._http_header

    def enable_module(self, module, **kwargs):
        """
        Expend server using a module definition which can be used to serve custom
        client code or assets, load/initialize resources (js, css, vue),
        register custom protocols and even execute custom code.

        Any previously seem module will be automatically skipped.

        The attributes that are getting processed in a module are the following:
          - setup(server, **kwargs): Function called first
          - scripts = []           : List all JavaScript URL that should be loaded
          - module_scripts = []    : List all JavaScript URL as type=module to load
          - styles  = []           : List all CSS URL that should be loaded
          - vue_use = ['libName', ('libName2', { **options })]: List Vue plugin to load
          - state = {}             : Set of variable to add to state
          - serve = { data: '/path/on/fs' }: Set of endpoints to serve static content
          - www = '/path/on/fs'    : Path served as main web content

        :param module: A module to enable or a dict()
        :param kwargs: Any optional parameters needed for your module setup() function.
        """
        if self.root_server != self:
            return self.root_server.enable_module(module, **kwargs)

        # Make sure definitions is a dict while skipping already loaded module
        definitions = module
        if isinstance(definitions, dict):
            if definitions in self._loaded_module_dicts:
                return False

            self._loaded_module_dicts.append(definitions)
        elif definitions in self._loaded_modules:
            return False
        else:
            self._loaded_modules.add(definitions)
            definitions = definitions.__dict__

        if "setup" in definitions:
            definitions["setup"](self, **kwargs)

        for key in ["scripts", "module_scripts", "styles", "vue_use"]:
            if key in definitions:
                self.state[f"trame__{key}"] += definitions[key]

        if "state" in definitions:
            self.state.update(definitions["state"])

        if "serve" in definitions:
            self.serve.update(definitions["serve"])

        if "www" in definitions:
            self._www = definitions["www"]

        # Reduce vue_use to merge options
        utils.reduce_vue_use(self.state)

        return True

    # -------------------------------------------------------------------------
    # Call methods
    # -------------------------------------------------------------------------

    def js_call(self, ref: str | None = None, method: str | None = None, *args):
        """
        Python call method on JS element.

        :param ref: ref name of the widget element
        :type ref: str

        :param method: name of the method that should be called
        :type method: str

        :param *args: set of parameters needed for the function
        """
        if self.protocol:
            self.protocol.push_actions(
                [
                    {
                        "type": "method",
                        "ref": ref,
                        "method": method,
                        "args": list(args),
                    },
                ],
            )

    # -------------------------------------------------------------------------
    # Annotations
    # -------------------------------------------------------------------------

    @property
    def change(self):
        """
        Use as decorator `@server.change(key1, key2, ...)` so the decorated function
        will be called like so `_fn(**state)` when any of the listed key name
        is getting modified from either client or server.

        :param *_args: A list of variable name to monitor
        :type *_args: str
        """
        return self._state.change

    # -------------------------------------------------------------------------

    @property
    def trigger(self):
        """
        Use as decorator `@server.trigger(name)` so the decorated function
        will be able to be called from the client by doing `click="trigger(name)"`.

        :param name: A name to use for that trigger
        :type name: str
        """
        return self._controller.trigger

    # -------------------------------------------------------------------------
    # From a function get its trigger name and register it if need be
    # -------------------------------------------------------------------------

    @property
    def trigger_name(self):
        """
        Given a function this method will register a trigger and returned its name.
        If manually registered, the given name at the time will be returned.

        :return: The trigger name for that function
        :rtype: str
        """
        return self._controller.trigger_name

    # -------------------------------------------------------------------------
    # App properties
    # -------------------------------------------------------------------------

    @property
    def name(self) -> str:
        """Name of server"""
        return self._name

    @property
    def root_server(self) -> Server:
        """Root server to start"""
        if self._parent_server:
            return self._parent_server.root_server
        return self

    @property
    def translator(self):
        """Translator of the server"""
        return self._translator

    @property
    def options(self):
        """Server options provided at instantiation time"""
        return self._options

    @property
    def client_type(self) -> ClientType:
        """Specify the client type. Either 'vue2' or 'vue3' for now."""
        if self._client_type is None:
            return DEFAULT_CLIENT_TYPE  # default
        return self._client_type

    @client_type.setter
    def client_type(self, value: ClientType) -> None:
        """Should only be called once before any widget initialization"""
        if self._client_type is None:
            self._client_type = value

        if self.client_type != value:
            msg = (
                f"Trying to switch client_type from {self._client_type} to {value}."
                "The client_type can only be set once."
            )
            raise TypeError(msg)

    @property
    def cli(self):
        """argparse parser"""
        if self.root_server != self:
            return self.root_server.cli

        if self._cli_parser:
            return self._cli_parser

        self._cli_parser = ArgumentParser(
            description="Kitware trame",
            allow_abbrev=False,
        )

        # Trame specific args
        self._cli_parser.add_argument(
            "--server",
            help="Prevent your browser from opening at startup",
            action="store_true",
        )
        self._cli_parser.add_argument(
            "--banner",
            help="Print trame banner",
            action="store_true",
        )
        self._cli_parser.add_argument(
            "--app",
            help="Use OS built-in browser",
            action="store_true",
        )
        self._cli_parser.add_argument(
            "--no-http",
            help="Do not serve anything over http",
            dest="no_http",
            action="store_true",
        )
        self._cli_parser.add_argument(
            "--authKeyFile",
            help="""Path to a File that contains the Authentication key for clients
                    to connect to the WebSocket.
                    This takes precedence over '-a, --authKey' from wslink.""",
        )
        self._cli_parser.add_argument(
            "--hot-reload",
            help="""Automatically reload state/controller callback functions for every
                    function call. This allows live editing of the functions. Functions
                    located in the site-packages directories are skipped.""",
            action="store_true",
        )
        self._cli_parser.add_argument(
            "--trame-args",
            help="""If specified, trame will ignore all other arguments, and only the contents
                    of the `--trame-args` will be used. For example:
                    `--trame-args="-p 8081 --server"`. Alternatively, the environment variable
                    `TRAME_ARGS` may be set instead.""",
        )
        self._cli_parser.add_argument(
            "--follow-symlinks",
            dest="static_follow_symlinks",
            help="""flag for allowing to follow symlinks that lead outside
                    the static root directory, by default it's not allowed
                    and HTTP/404 will be returned on access.
                    Enabling follow_symlinks can be a security risk,
                    and may lead to a directory transversal attack.
                    You do NOT need this option to follow symlinks which point
                    to somewhere else within the static directory, this option
                    is only used to break out of the security sandbox.
                    Enabling this option is highly discouraged, and only
                    expected to be used for edge cases in a local development
                    setting where remote users do not have access to the server.""",
            action="store_true",
        )

        CoreServer.add_arguments(self._cli_parser)

        return self._cli_parser

    @property
    def state(self) -> State:
        """
        :return: The server shared state
        :rtype: trame_server.state.State
        """
        return self._state

    @property
    def context(self) -> State:
        """
        The server-only context (not shared with the client).

        :return: The server context state
        :rtype: trame_server.state.State
        """
        return self._context

    @property
    def controller(self) -> Controller:
        """
        :return: The server controller
        :rtype: trame_server.controller.Controller
        """
        return self._controller

    @property
    def ui(self) -> VirtualNodeManager:
        """
        :return: The server VirtualNode manager
        :rtype: trame_server.ui.VirtualNodeManager
        """
        return self._ui

    @property
    def running(self) -> bool:
        """Return True if the server is currently starting or running."""
        if self.root_server != self:
            return self.root_server.running

        return self._running_stage > 1

    @property
    def network_completion(self):
        """Return a future to await if you want to ensure that any pending network call
        have been issued before locking the server"""
        return asyncio.ensure_future(self.context.network_monitor.completion())

    @property
    def ready(self):
        """Return a future that will resolve once the server is ready"""
        if self.root_server != self:
            return self.root_server.ready

        if self._running_future is None:
            self._running_future = asyncio.get_running_loop().create_future()

        return self._running_future

    def _resolve_ready_future(self, result=None, exception=None) -> None:
        if self.ready.done():
            return

        if exception:
            self.ready.set_exception(exception)
        else:
            self.ready.set_result(result)

    # -------------------------------------------------------------------------
    # API for network handling
    # -------------------------------------------------------------------------

    def get_server_state(self):
        """Return the current server state"""
        return {
            "name": self._name,
            "state": self.state.initial,
        }

    def clear_state_client_cache(self, *state_names):
        protocol = self.protocol
        if protocol:
            protocol.clear_state_client_cache(*state_names)

    # -------------------------------------------------------------------------

    def add_protocol_to_configure(self, configure_protocol_fn):
        """
        Register function that will be called with a wslink.ServerProtocol
        when the server start and is ready for registering new wslink.Protocol.

        :param configure_protocol_fn: A function to be called later with a
                                      wslink.ServerProtocol as argument.
        """
        if self.root_server != self:
            self.root_server.add_protocol_to_configure(configure_protocol_fn)
            return

        self._protocols_to_configure.append(configure_protocol_fn)

    @property
    def protocol(self):
        """Return the server root protocol"""
        if self.root_server != self:
            return self.root_server.protocol

        return self._root_protocol

    # -------------------------------------------------------------------------

    def protocol_call(self, method, *args, **kwargs):
        """
        Call a registered protocol method

        :param method: Method registration name
        :type method: str
        :param *args: Set of args to use for that method call
        :param **kwargs: Set of keyword arguments to use for that method call
        :return: transparently return what the called function returns
        """
        if self.protocol:
            pair = self.protocol.getRPCMethod(method)
            if pair:
                obj, func = pair
                return func(obj, *args, **kwargs)
            return None

        error = "Protocol does not exist yet"
        raise ValueError(error)

    def force_state_push(self, *key_names):
        """
        Should only be needed when client corrupted its data and need the server need to send it again.

        :param *args: Set of key names to be send again to the client.
        """
        self.protocol_call(
            "trame.force.push", *[self._translator.translate_key(k) for k in key_names]
        )

    # -------------------------------------------------------------------------
    # Server handling (start/stop/port)
    # -------------------------------------------------------------------------

    def start(
        self,
        port: int | None = None,
        thread: bool = False,
        open_browser: bool | None = None,
        show_connection_info: bool = True,
        disable_logging: bool = False,
        backend: BackendType | None = None,
        follow_symlinks: bool | None = None,
        exec_mode: ExecModeType = "main",
        timeout: int | None = None,
        host: str | None = None,
        **kwargs,
    ):
        """
        Start the server by listening to the provided port or using the
        `--port, -p` command line argument.
        If the server is already starting or started, any further call will be skipped.

        When the exec_mode="main" or "desktop", the method will be blocking.
        If exec_mode="task", the method will return a scheduled task.
        If exec_mode="coroutine", the method will return a coroutine which
        will need to be scheduled by the user.

        :param port: A port number to listen to. When 0 is provided
                     the system will use a random open port.
        :param thread: If the server run in a thread which means
                       we should disable interruption listeners
        :param open_browser: Should we open the system browser with app url.
                             Using the `--server` command line argument is
                             similar to setting it to False.
        :param show_connection_info: Should we print connection URL at startup?
        :param disable_logging: Ask wslink to disable logging
        :param backend: aiohttp by default but could be generic or tornado.
                        This can also be set with the environment variable ``TRAME_BACKEND``.
                        Defaults to ``'aiohttp'``.
        :param exec_mode: main/desktop/task/coroutine
                          specify how the start function should work
        :param timeout: How much second should we wait before automatically
                        stopping the server when no client is connected.
                        Setting it to 0 will disable such auto-shutdown.
        :param host: The hostname used to bind the server. This can also be
                     set with the environment variable ``TRAME_DEFAULT_HOST``.
                     Defaults to ``'localhost'``.
        :param **kwargs: Keyword arguments for capturing optional parameters
                         for wslink server and/or desktop browser
        """
        if self.root_server != self:
            self.root_server.start(
                port=port,
                thread=thread,
                open_browser=open_browser,
                show_connection_info=show_connection_info,
                disable_logging=disable_logging,
                backend=backend,
                exec_mode=exec_mode,
                timeout=timeout,
                host=host,
                **kwargs,
            )
            return None

        if self._running_stage:
            return None

        # Try to bind client if none were added
        if self._www is None:
            from trame_client import module  # noqa: PLC0415

            self.enable_module(module)

        # Apply any header change needed
        self._http_header.apply()

        # Trigger on_server_start life cycle callback
        if self.controller.on_server_start.exists():
            self.controller.on_server_start(self)

        CoreServer.bind_server(self)
        options = self.cli.parse_known_args()[0]

        if backend is None:
            backend = os.environ.get("TRAME_BACKEND", "aiohttp")

        if follow_symlinks is not None:
            options.static_follow_symlinks = follow_symlinks

        if open_browser is None:
            open_browser = not os.environ.get("TRAME_SERVER")

        if options.host == "localhost":
            if host is None:
                host = os.environ.get("TRAME_DEFAULT_HOST", "localhost")
            options.host = host

        if timeout is not None:
            options.timeout = timeout

        if port is not None:
            options.port = port

        if not options.content:
            options.content = self._www

        if thread:
            options.nosignalhandlers = True

        if options.banner:
            from .utils.banner import print_banner  # noqa: PLC0415

            self.controller.on_server_ready.add(print_banner)

        if options.app:
            exec_mode = "desktop"

        if exec_mode == "desktop":
            from .utils.desktop import start_browser  # noqa: PLC0415

            options.port = 0
            exec_mode, show_connection_info, open_browser = "main", False, False
            self.controller.on_server_ready.add(
                lambda **_: start_browser(self, **kwargs)
            )

        # Allow for older wslink versions where this was not an attribute
        reverse_url = getattr(options, "reverse_url", None)

        if not reverse_url and show_connection_info and exec_mode != "task":
            from .utils.server import print_informations  # noqa: PLC0415

            self.controller.on_server_ready.add(lambda **_: print_informations(self))

        if (
            not reverse_url
            and open_browser
            and exec_mode != "task"
            and not options.server
        ):
            from .utils.browser import open_browser  # noqa: PLC0415

            self.controller.on_server_ready.add(lambda **_: open_browser(self))

        if len(self.serve):
            endpoints = []
            for key in self.serve:
                value = self.serve[key]
                if isinstance(value, (list, tuple)):
                    # tuple are use to describe sync loading (affect client)
                    endpoints.append(f"{key}={value[0]}")
                else:
                    endpoints.append(f"{key}={value}")
            options.fsEndpoints = "|".join(endpoints)

        # Reset http delivery
        if options.no_http:
            options.content = ""
            options.fsEndpoints = ""

        self._server_options = options
        CoreServer.configure(options)

        self._running_stage = 1
        task = CoreServer.server_start(
            options,
            **{  # Do a proper merging/override
                **kwargs,
                "disableLogging": disable_logging,
                "backend": backend,
                "exec_mode": exec_mode,
            },
        )

        # Manage exit life cycle unless coroutine
        if exec_mode == "main":
            self._running_stage = 0
            if self.controller.on_server_exited.exists():
                loop = asyncio.get_event_loop()
                for exit_task in self.controller.on_server_exited(
                    **self.state.to_dict()
                ):
                    if inspect.isawaitable(exit_task):
                        loop.run_until_complete(exit_task)
                    elif callable(exit_task):
                        result = exit_task()
                        if inspect.isawaitable(result):
                            loop.run_until_complete(result)
        elif hasattr(task, "add_done_callback"):

            def on_done(task: asyncio.Task) -> None:
                try:
                    task.result()
                    self._running_stage = 0
                    if self.controller.on_server_exited.exists():
                        self.controller.on_server_exited(**self.state.to_dict())
                except asyncio.CancelledError:
                    self._resolve_ready_future(result=False)
                except Exception as e:  # pylint: disable=broad-except
                    logging.exception("Exception raised by task = %r", task)
                    self._resolve_ready_future(exception=e)

            task.add_done_callback(on_done)

        return task

    async def stop(self) -> None:
        """Coroutine for stopping the server"""
        if self.root_server != self:
            await self.root_server.stop()
        elif self._running_stage:
            await self._server.stop()
            self._running_future = None
        self._running_stage = 0

    @property
    def port(self) -> int:
        """Once started, you can retrieve the port used"""
        if self.root_server != self:
            return self.root_server.port

        return self._running_port

    @property
    def server_options(self):
        """Once started, you can retrieve the server options used"""
        return self._server_options
