"""
`NetlabClient` is the main entrypoint to the NETLAB+ SDK.

Only `NetlabClient` should be created directly. See the mixin classes for a full list
of methods:

- :py:class:`netlab.api.ClassApiMixin`
- :py:class:`netlab.api.PodApiMixin`
- :py:class:`netlab.api.ReservationApiMixin`
- :py:class:`netlab.api.SystemApiMixin`
- :py:class:`netlab.api.LabApiMixin`
- :py:class:`netlab.api.UserApiMixin`
- :py:class:`netlab.api.VmApiMixin`

**Usage**

::

    async with NetlabClient() as client:
        client: NetlabConnection

        status = await client.system_status_get()
"""

import asyncio
import json
import logging
import ssl
import uuid
from typing import Any, Callable, Dict, NoReturn, Optional, Union, cast
from weakref import WeakValueDictionary

from typing_extensions import TypeAlias

from .api import ClassApiMixin, LabApiMixin, PodApiMixin, ReservationApiMixin, SystemApiMixin, UserApiMixin, VmApiMixin
from .auth import SystemConfig, get_system_config
from .config import NetlabServerConfig
from .errors import ResponseFormatError, error_decode, hdr
from .errors.common import NetlabConnectionClosedError
from .serializer import deserialize, serialize

log = logging.getLogger()


UGLY_TASKS = {
    "vm.inventory.import.task",
    "vm.datacenter.test.task",
    "vm.datacenter.discover.vms.task",
    "vm.datacenter.discover.hosts.task",
}

HDR_TASKS = {
    "pod.clone.task",
    "pod.remove.task",
}

EventQueue: TypeAlias = asyncio.Queue[asyncio.Future[Dict[str, Any]]]
ProgressCallback: TypeAlias = Callable[[Dict[str, Any]], None]
AnyFuture: TypeAlias = asyncio.Future[Any]


__all__ = ["NetlabClient", "NetlabConnection", "EventStream"]


class EventStream(object):
    """
    Used to collect events from a netlab subscription.

    :meta private:
    """

    def __init__(self, connection, queue):
        self._connection = connection
        self._queue = queue

    async def get(self) -> Any:
        """
        Get the next event from a subscription.
        """
        msg = await self._connection._wait_next_event(self._queue)
        return msg


class EventStreamFactory(object):
    """
    Used to create an event stream.

    :meta private:

    **Usage**

    ::

        async with NetlabClient() as client:
            async with client.subscribe('RESERVATION.ADDED') as events:
                event = await events.get()
    """

    _connection: "NetlabConnection"
    _queue: EventQueue
    _handle: uuid.UUID
    _event: str
    _criteria: Dict[str, Any]

    def __init__(self, connection, event, criteria):
        self._connection = connection
        self._queue = asyncio.Queue()
        self._handle = uuid.uuid4()
        self._event = event
        self._criteria = criteria

    async def __aenter__(self) -> EventStream:
        self._connection._event_mapping[self._handle] = self._queue

        await self._connection.call("event.subscribe", event=self._event, criteria=self._criteria, handle=self._handle)
        log.debug("subscription created %s %s", self._handle, self._event)
        es = EventStream(self._connection, self._queue)

        return es

    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
        if self._connection.alive():
            await self._connection.call("event.unsubscribe", handle=self._handle, event=self._event)
            del self._connection._event_mapping[self._handle]
            log.debug("subscription deleted %s %s", self._handle, self._event)


def _id_from_config(config: SystemConfig) -> tuple[str, int]:
    return (config.host, config.port)


class NetlabConnection(
    ClassApiMixin, PodApiMixin, ReservationApiMixin, SystemApiMixin, LabApiMixin, UserApiMixin, VmApiMixin
):
    """
    An active connection to a NETLAB+ used to communicate with it. Should only be created by a
    :py:class:`NetlabClient`.
    """

    _system_version: str
    _event_handler_task: asyncio.Future[NoReturn]
    _keep_alive_task: asyncio.Future[NoReturn]

    def __init__(self, reader, writer, config):
        self._reader: asyncio.StreamReader = reader
        self._writer: asyncio.StreamWriter = writer
        self._config = config
        self._event_mapping: WeakValueDictionary[uuid.UUID, EventQueue] = WeakValueDictionary()

    async def _connect(self) -> None:
        self._event_handler_task = asyncio.ensure_future(self._event_handler())
        self._keep_alive_task = asyncio.ensure_future(self._keep_alive())
        config = self._config
        await self.call("user.authenticate", user=config.user, token=config.token)
        result = await self.system_status_get()
        self._server_version = result["sys_sdn_version"]

    async def _disconnect(self) -> None:
        self._writer.close()
        # await self._writer.wait_closed()
        self._keep_alive_task.cancel()
        self._event_handler_task.cancel()
        await asyncio.wait([self._event_handler_task, self._keep_alive_task])

    async def _keep_alive(self) -> NoReturn:
        missed = 0
        while True:
            await asyncio.sleep(15)
            try:
                await asyncio.wait_for(self.call("internal.mbusd.ping"), 2.0)
            except asyncio.TimeoutError:
                missed += 1
                log.debug("keep alive packet lost (%d/3)", missed)
                if missed >= 3:
                    raise NetlabConnectionClosedError("Keep alive packet lost.") from None
                continue
            missed = 0
            log.debug("keep alive packet sent")

    async def _event_handler(self) -> NoReturn:
        while True:
            bmsg = await self._reader.readuntil(b"\n")

            if log.isEnabledFor(logging.DEBUG):
                log.debug("data recv", extra={"data": json.loads(bmsg.decode("utf-8"))})
            msg = deserialize(bmsg.decode("utf-8"))

            if "id" in msg:
                ident = uuid.UUID(msg["id"])
                payload = msg
            elif "handle" in msg:
                ident = uuid.UUID(msg["handle"])
                payload = msg["params"]
            else:
                raise ResponseFormatError('message did not contain "handle" or "id"')

            queue = self._event_mapping.get(ident)

            if queue is None:
                continue

            responder: AnyFuture = asyncio.Future()

            if "result" in payload or "event" in msg:
                responder.set_result(msg)
                queue.put_nowait(responder)
            elif "error" in payload:
                responder.set_exception(error_decode(msg["error"]))
                queue.put_nowait(responder)
            else:
                raise ResponseFormatError('message did not contain "result" or "error"')

    async def _call_method(self, ident: uuid.UUID, data: Dict[str, Any]) -> Any:
        msg = serialize(data)
        if log.isEnabledFor(logging.DEBUG):
            log.debug("data sent", extra={"data": json.loads(msg)})
        bmsg = msg.encode("utf-8") + b"\n"
        self._writer.write(bmsg)

        try:
            await self._writer.drain()
        except ConnectionResetError as e:
            raise NetlabConnectionClosedError() from e

        queue: EventQueue = asyncio.Queue()
        self._event_mapping[ident] = queue

        msg = await self._wait_next_event(queue)

        if "result" in msg:
            return msg["result"]
        elif msg.get("error") is not None:
            raise error_decode(msg["error"])
        else:
            raise ResponseFormatError('message did not contain "result" or "error"')

    async def _call_hdr(
        self,
        ident: uuid.UUID,
        data: Dict[str, Any],
        progress: Optional[ProgressCallback] = None,
    ) -> Any:
        response = await self._call_task(ident, data, progress=progress)

        method = data["method"]
        status_code = response["status_t"]

        if status_code in ("PARTIAL", "FAILED", "UNKNOWN"):
            hdr_error = hdr.HDRError(f"HDR task {method} failed with {status_code}")
            hdr_error.status_code = status_code
            hdr_error.method = method
            hdr_error.ident = ident
            hdr_error.response = response
            raise hdr_error

        return response

    async def _call_task(
        self,
        ident: uuid.UUID,
        data: Dict[str, Any],
        progress: Optional[ProgressCallback] = None,
    ) -> Any:
        handle = uuid.uuid4()
        queue: EventQueue = asyncio.Queue()
        self._event_mapping[handle] = queue

        data["params"]["notify_complete"] = True
        data["params"]["notify_handle"] = handle
        if progress is not None:
            data["params"]["notify_progress"] = True

        await self._call_method(ident, data)

        while True:
            event = await self._wait_next_event(queue)

            if event.get("event") == "TASK.PROGRESS":
                if progress is not None:
                    params = event.get("params", {})
                    progress(params)
                continue

            if "params" not in event:
                raise ResponseFormatError('message did not contain "params"')
            msg = event["params"]

            if "result" in msg:
                return msg["result"]
            elif msg.get("error") is not None:
                raise error_decode(msg["error"])
            else:
                raise ResponseFormatError('message did not contain "result" or "error"')

    async def _call_task_ugly(self, ident: uuid.UUID, data: Dict[str, Any]) -> Any:
        task_id = await self._call_method(ident, data)

        while True:
            status = await self.call("task.check", task_id=task_id)
            if status["is_complete"]:
                return status["result"]
            await asyncio.sleep(1)

    async def _wait_next_event(self, queue: EventQueue) -> Any:
        queue_task: AnyFuture = asyncio.ensure_future(queue.get())

        try:
            completed_tasks, _ = await asyncio.wait(
                [queue_task, self._keep_alive_task, self._event_handler_task],
                return_when=asyncio.FIRST_COMPLETED,
            )

            completed_task = completed_tasks.pop()

            if completed_task is queue_task:
                future = completed_task.result()
                result = await future
                return result
            else:
                raise NetlabConnectionClosedError() from completed_task.exception()
        finally:
            queue_task.cancel()
            try:
                await queue_task
            except asyncio.CancelledError:
                pass

    def alive(self) -> bool:
        return not bool(self._keep_alive_task.done() or self._event_handler_task.done())

    async def call(self, method: str, **kwargs: Any) -> Any:
        """

        :param method: The NETLAB+ method name.
        :param kwargs: The arguments to pass to netlab. The reserved kwarg
            ``_progress`` (a ``Callable[[int], None]``) is consumed by the
            client to receive 0-100 progress updates from polled task methods
            (currently HDR tasks); it is never sent to the server.

        .. warning::

            In most cases, it is better to use the higher level methods, rather than this
            low level method.
        """
        progress = kwargs.pop("_progress", None)
        ident = uuid.uuid4()
        data = {"id": ident, "jsonrpc": "2.0", "method": method, "params": kwargs}

        if method in UGLY_TASKS:
            return await self._call_task_ugly(ident, data)
        if method in HDR_TASKS:
            return await self._call_hdr(ident, data, progress=progress)
        elif method.endswith(".task"):
            return await self._call_task(ident, data, progress=progress)
        else:
            return await self._call_method(ident, data)

    def subscribe(self, event: str, **criteria: Any) -> EventStreamFactory:
        """
        Subscribe to a NETLAB+ event.

        :meta private:

        :param event: The NETLAB+ event name.
        :param critera: Any filtering criteria.

        :return: Used to accept new events.
        """

        es = EventStreamFactory(self, event, criteria)
        return es


class NetlabClient(object):
    """
    Async Context Manager for creating connections to NETLAB+.
    """

    _connection: Optional[NetlabConnection]

    def __init__(
        self,
        system: Optional[str] = "default",
        config: Union[None, NetlabServerConfig, Dict[str, Any], str] = None,  # TODO Mypy cannot infer TypedDict
        config_path: Optional[str] = None,
    ):
        """
        :param system: The name to identify a NETLAB+ system. Defaults to **'default'**.
        :param config: Configuration options. See below for all possible options for the :attr:`config`.
        :param config_path: File path location of a JSON config file to be used instead of the default location.
        """
        self._config = get_system_config(system, config, config_path)
        self._connection = None

    async def __aenter__(self) -> NetlabConnection:
        assert self._connection is None
        config = self._config
        try:
            log.debug("connecting to %s", config.host)

            ssl_context: ssl.SSLContext
            if isinstance(config.ssl, ssl.SSLContext):
                ssl_context = config.ssl
            else:
                ssl_context = ssl.create_default_context()
                if self._config.ssl_ciphers:
                    ssl_context.set_ciphers(self._config.ssl_ciphers)
                if config.ssl in ("self_signed", "noverify"):
                    ssl_context.check_hostname = False
                    ssl_context.verify_mode = ssl.CERT_NONE
                if cast(Any, config.ssl) not in ("default", "self_signed", "noverify"):
                    log.warning("`ssl` should be an `ssl.SSLContext`, `'default'`, or `'noverify'`.")

            reader, writer = await asyncio.open_connection(
                host=config.host,
                port=config.port,
                ssl=ssl_context,
                server_hostname=config.server_hostname,
                limit=config.message_byte_limit,
            )
            log.debug("tcp connection to %s", config.host)
        except ConnectionAbortedError as e:
            raise NetlabConnectionClosedError() from e
        client = NetlabConnection(reader, writer, config)
        self._connection = client
        try:
            await client._connect()
            log.debug("completed connection to %s", config.host)
        except Exception:
            await client._disconnect()
            raise
        return client

    async def __aexit__(self, exc_type, exc_value, traceback):
        assert self._connection, "__aexit__ called without __aenter__"
        await self._connection._disconnect()
        self._connection = None
        log.debug("successfully disconnected from %s", self._config.host)
