#
#    Copyright (c) 2023 Project CHIP Authors
#
#    Licensed under the Apache License, Version 2.0 (the 'License');
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an 'AS IS' BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

import subprocess
import time
from dataclasses import dataclass

import websockets

from .hooks import WebSocketRunnerHooks
from .runner import TestRunner

_KEEP_ALIVE_TIMEOUT_IN_SECONDS = 40


@dataclass
class WebSocketRunnerConfig:
    server_address: str = 'localhost'
    server_port: int = '9002'
    server_path: str = None
    server_arguments: str = None
    hooks: WebSocketRunnerHooks = WebSocketRunnerHooks()


class WebSocketRunner(TestRunner):
    def __init__(self, config: WebSocketRunnerConfig):
        self._client = None
        self._server = None
        self._hooks = config.hooks

        self._server_connection_url = self._make_server_connection_url(
            config.server_address, config.server_port)
        self._server_startup_command = self._make_server_startup_command(
            config.server_path, config.server_arguments, config.server_port)

    async def start(self):
        self._server = await self._start_server(self._server_startup_command)
        self._client = await self._start_client(self._server_connection_url)

    async def stop(self):
        await self._stop_client(self._client)
        await self._stop_server(self._server)
        self._client = None
        self._server = None

    async def execute(self, request):
        instance = self._client
        if instance:
            await instance.send(request)
            return await instance.recv()
        return None

    async def _start_client(self, url, max_retries=4, interval_between_retries=1):
        if max_retries:
            start = time.time()
            try:
                self._hooks.connecting(url)
                connection = await websockets.connect(url, ping_timeout=_KEEP_ALIVE_TIMEOUT_IN_SECONDS)
                duration = round((time.time() - start) * 1000, 0)
                self._hooks.success(duration)
                return connection
            except Exception:
                duration = round((time.time() - start) * 1000, 0)
                self._hooks.failure(duration)
                self._hooks.retry(interval_between_retries)
                time.sleep(interval_between_retries)
                return await self._start_client(url, max_retries - 1, interval_between_retries + 1)

        self._hooks.abort(url)
        raise Exception(f'Connecting to {url} failed.')

    async def _stop_client(self, instance):
        if instance:
            await instance.close()

    async def _start_server(self, command):
        instance = None
        if command:
            instance = subprocess.Popen(command, stdout=subprocess.DEVNULL)
        return instance

    async def _stop_server(self, instance):
        if instance:
            instance.kill()

    def _make_server_connection_url(self, address: str, port: int):
        return 'ws://' + address + ':' + str(port)

    def _make_server_startup_command(self, path: str, arguments: str, port: int):
        if path is None:
            return None
        elif arguments is None:
            return [path] + ['--port', str(port)]
        else:
            return [path] + [arg.strip() for arg in arguments.split(' ')] + ['--port', str(port)]