import time
import asyncio
import argparse
from typing import Optional

from ....enums import RemoteWorkspaceState
from ..asyncio.base import AsyncZMQRuntime
from ....helper import cached_property, colored
from ....excepts import DaemonConnectivityError
from .client import PeaDaemonClient, WorkspaceDaemonClient


class JinadRuntime(AsyncZMQRuntime):
    """Runtime procedure for Jinad."""

    def __init__(self, args: 'argparse.Namespace', ctrl_addr: str):
        super().__init__(args, ctrl_addr)
        # Need the `proper` control address to send `activate` and `deactivate` signals, from the pea in the `main`
        # process.
        self.host = args.host
        self.port_expose = args.port_expose
        # NOTE: args.timeout_ready is always set to -1 for JinadRuntime so that wait_for_success doesn't fail in Pea,
        # so it can't be used for Client timeout.
        self.workspace_api = WorkspaceDaemonClient(
            host=self.host,
            port=self.port_expose,
            logger=self.logger,
        )
        self.pea_api = PeaDaemonClient(
            host=self.host,
            port=self.port_expose,
            logger=self.logger,
        )
        # Uploads Pea context to remote & Creates remote Pea using :class:`PeaDaemonClient`
        if self._remote_id:
            self.logger.success(
                f'created a remote {self.pea_api.kind}: {colored(self._remote_id, "cyan")}'
            )

    async def async_run_forever(self):
        """
        Streams log messages using websocket from remote server
        """
        self._logging_task = asyncio.create_task(
            self._sleep_forever()
            if self.args.quiet_remote_logs
            else self.pea_api.logstream(self._remote_id)
        )

    async def async_cancel(self):
        """
        Cancels the logging task
        """
        self._logging_task.cancel()

    def teardown(self):
        """
        Terminates the remote Workspace/Pod/Pea using :class:`JinadAPI`
        """
        self.pea_api.delete(id=self._remote_id)
        # TODO: don't fail if workspace deletion fails. all peas would make this call. can be optimized
        self.workspace_api.delete(id=self.args.workspace_id)
        super().teardown()

    def create_workspace(self):
        """Create a workspace on remote (includes file upload & docker build)

        :raises DaemonConnectivityError: if remote daemon is not reachable
        :raises RuntimeError: if workspace creation fails
        """
        if not self.workspace_api.alive:
            self.logger.error(
                f'fail to connect to the daemon at {self.host}:{self.port_expose}, please check:\n'
                f'- is there a typo in {self.host}?\n'
                f'- on {self.host}, are you running `docker run --network host jinaai/jina:latest-daemon`?\n'
                f'- on {self.host}, have you set the security policy to public for all required ports?\n'
                f'- on local, are you behind VPN or proxy?'
            )
            raise DaemonConnectivityError

        sleep = 2
        retries = 100
        for retry in range(retries):
            workspace_status = self.workspace_api.get(id=self.args.workspace_id)
            if not workspace_status:
                raise DaemonConnectivityError
            state = workspace_status.get('state', None)
            if not state:
                self.logger.info(
                    f'creating workspace {colored(self.args.workspace_id, "cyan")} on remote. This might take some time.'
                )
                self.workspace_api.post(
                    dependencies=self.args.upload_files,
                    workspace_id=self.args.workspace_id,
                )
            elif state in [
                RemoteWorkspaceState.PENDING,
                RemoteWorkspaceState.CREATING,
                RemoteWorkspaceState.UPDATING,
            ]:
                if retry % 10 == 0:
                    self.logger.info(
                        f'workspace {self.args.workspace_id} is getting created on remote. waiting..'
                    )
                time.sleep(sleep)
            elif state == RemoteWorkspaceState.ACTIVE:
                self.logger.success(
                    f'successfully created a remote workspace: {colored(self.args.workspace_id, "cyan")}'
                )
                break
            else:
                raise RuntimeError(f'remote workspace creation failed')

    @cached_property
    def _remote_id(self) -> Optional[str]:
        """Creates a workspace & a pea on remote

        :return: id of rempte pea
        """
        self.create_workspace()
        pea_id = self.pea_api.post(self.args)
        if not pea_id:
            raise RuntimeError('remote pea creation failed')
        return pea_id

    async def _sleep_forever(self):
        """Sleep forever, no prince will come."""
        await asyncio.sleep(1e10)
