# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.

import asyncio
import json
import uuid
from http import HTTPStatus
from typing import Tuple, Union

import aiohttp

from dashscope.api_entities.base_request import AioBaseRequest
from dashscope.api_entities.dashscope_response import DashScopeAPIResponse
from dashscope.common.constants import (
    DEFAULT_REQUEST_TIMEOUT_SECONDS,
    SERVICE_503_MESSAGE,
    WEBSOCKET_ERROR_CODE,
)
from dashscope.common.error import (
    RequestFailure,
    UnexpectedMessageReceived,
    UnknownMessageReceived,
)
from dashscope.common.logging import logger
from dashscope.common.utils import async_to_sync
from dashscope.protocol.websocket import (
    ACTION_KEY,
    ERROR_MESSAGE,
    ERROR_NAME,
    EVENT_KEY,
    HEADER,
    TASK_ID,
    ActionType,
    EventType,
    WebsocketStreamingMode,
)


class WebSocketRequest(AioBaseRequest):
    def __init__(
        self,
        url: str,
        api_key: str,
        stream: bool = True,
        ws_stream_mode: str = WebsocketStreamingMode.OUT,
        is_binary_input: bool = False,
        timeout: int = DEFAULT_REQUEST_TIMEOUT_SECONDS,
        flattened_output: bool = False,
        pre_task_id=None,
        user_agent: str = "",
    ) -> None:
        super().__init__(user_agent=user_agent)
        # pylint: disable=pointless-string-statement
        """HttpRequest.

        Args:
            url (str): The request url.
            api_key (str): The api key.
            method (str): The http method(GET|POST).
            stream (bool, optional): Is stream request. Defaults to False.
            timeout (int, optional): Total request timeout.
                Defaults to DEFAULT_REQUEST_TIMEOUT_SECONDS.
        """
        self.url = url
        self.stream = stream
        self.flattened_output = flattened_output
        if timeout is None:
            self.timeout = DEFAULT_REQUEST_TIMEOUT_SECONDS
        else:
            self.timeout = timeout  # type: ignore[has-type]
        self.ws_stream_mode = ws_stream_mode
        self.is_binary_input = is_binary_input

        self.headers = {
            "Authorization": f"bearer {api_key}",
            **self.headers,
        }

        self.task_headers = {
            "streaming": self.ws_stream_mode,
        }
        self.pre_task_id = pre_task_id

    def add_headers(self, headers):
        self.headers = {**self.headers, **headers}

    def call(self):
        response = async_to_sync(self.connection_handler())
        if self.stream:
            return (item for item in response)
        else:
            output = next(response)
            try:
                next(response)
            except StopIteration:
                pass
            return output

    async def close(self):
        if self.ws is not None and not self.ws.closed:
            await self.ws.close()

    async def aio_call(self):
        response = self.connection_handler()
        if self.stream:
            return (item async for item in response)
        else:
            result = await response.__anext__()
            try:
                await response.__anext__()
            except StopAsyncIteration:
                pass
            return result

    async def connection_handler(self):  # pylint: disable=too-many-branches
        try:
            task_id = None
            async with aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(
                    total=self.timeout,
                ),
            ) as session:
                async with session.ws_connect(
                    self.url,
                    headers=self.headers,
                    heartbeat=6000,
                ) as ws:
                    await self._start_task(ws)  # send start task action.
                    task_id = self.task_headers["task_id"]
                    await self._wait_for_task_started(
                        ws,
                    )  # wait for task started event. # noqa E501
                    if self.ws_stream_mode == WebsocketStreamingMode.NONE:
                        if self.is_binary_input:  # send the binary package
                            data = self.data.get_batch_binary_data()
                            await ws.send_bytes(list(data.values())[0])
                        (
                            is_binary,
                            result,
                        ) = await self._receive_batch_data_task(  # noqa E501
                            ws,
                        )
                        # do not need send finished task message.
                        yield self._to_DashScopeAPIResponse(
                            task_id,
                            is_binary,
                            result,
                        )
                    elif self.ws_stream_mode == WebsocketStreamingMode.IN:
                        # server is in, we send streaming out.
                        await self._send_continue_task_data(ws)
                        (
                            is_binary,
                            result,
                        ) = await self._receive_batch_data_task(  # noqa E501
                            ws,
                        )
                        # do not need send finished task message.
                        yield self._to_DashScopeAPIResponse(
                            task_id,
                            is_binary,
                            result,
                        )
                    elif self.ws_stream_mode == WebsocketStreamingMode.OUT:
                        # we send batch data, server streaming output data.
                        if self.is_binary_input:  # send only binary package.
                            data = self.data.get_batch_binary_data()
                            await ws.send_bytes(list(data.values())[0])
                        async for is_binary, message in self._receive_streaming_data_task(  # noqa E501  # pylint: disable=line-too-long
                            ws,
                        ):
                            yield self._to_DashScopeAPIResponse(
                                task_id,
                                is_binary,
                                message,
                            )
                    else:  # duplex mode
                        asyncio.create_task(self._send_continue_task_data(ws))
                        async for is_binary, message in self._receive_streaming_data_task(  # noqa E501  # pylint: disable=line-too-long
                            ws,
                        ):
                            yield self._to_DashScopeAPIResponse(
                                task_id,
                                is_binary,
                                message,
                            )
        except RequestFailure as e:
            yield DashScopeAPIResponse(
                request_id=e.request_id,
                status_code=e.http_code,
                output=None,
                code=e.name,
                message=e.message,
            )
        except aiohttp.ClientConnectorError as e:
            logger.exception(e)
            yield DashScopeAPIResponse(
                request_id="",
                status_code=-1,
                code="ClientConnectorError",
                message=str(e),
            )
        except aiohttp.WSServerHandshakeError as e:
            code = e.status
            msg = e.message
            if e.status in [HTTPStatus.FORBIDDEN, HTTPStatus.UNAUTHORIZED]:
                msg = "Unauthorized, your api-key is invalid!"
            elif e.status == HTTPStatus.SERVICE_UNAVAILABLE:
                msg = SERVICE_503_MESSAGE
            else:
                pass
            yield DashScopeAPIResponse(
                request_id=task_id,
                status_code=code,
                code=code,
                message=msg,
            )
        except BaseException as e:
            logger.exception(e)
            yield DashScopeAPIResponse(
                request_id="",
                status_code=-1,
                code="Unknown",
                message=f"Error type: {type(e)}, message: {e}",
            )

    def _to_DashScopeAPIResponse(self, task_id, is_binary, result):
        if is_binary:
            return DashScopeAPIResponse(
                request_id=task_id,
                status_code=HTTPStatus.OK,
                output=result,
            )
        else:
            # get output and usage.
            output = {}
            usage = {}
            if "output" in result:
                output = result["output"]
            if "usage" in result:
                usage = result["usage"]
            return DashScopeAPIResponse(
                request_id=task_id,
                status_code=HTTPStatus.OK,
                output=output,
                usage=usage,
            )

    async def _receive_streaming_data_task(self, ws):
        # check if request stream data, re return an iterator,
        # otherwise we collect data and return user.
        # no matter what, the response is streaming
        is_binary_output = False
        while True:  # pylint: disable=R1702
            msg = await ws.receive()
            await self._check_websocket_unexpected_message(msg)
            if msg.type == aiohttp.WSMsgType.TEXT:
                msg_json = msg.json()
                logger.debug("Receive %s event", msg_json[HEADER][EVENT_KEY])
                if msg_json[HEADER][EVENT_KEY] == EventType.GENERATED:
                    payload = msg_json["payload"]
                    yield False, payload
                elif msg_json[HEADER][EVENT_KEY] == EventType.FINISHED:
                    payload = None
                    if "payload" in msg_json:
                        payload = msg_json["payload"]
                    logger.debug(payload)
                    if payload:
                        yield False, payload
                    else:
                        if not self.stream:
                            if is_binary_output:
                                yield True, payload
                            else:
                                yield False, payload
                    break
                elif msg_json[HEADER][EVENT_KEY] == EventType.FAILED:
                    self._on_failed(msg_json)
                else:
                    error = f"Receive unknown message: {msg_json}"
                    logger.error(error)
                    raise UnknownMessageReceived(error)
            elif msg.type == aiohttp.WSMsgType.BINARY:
                is_binary_output = True
                yield True, msg.data

    def _on_failed(self, details):
        error = RequestFailure(
            request_id=details[HEADER][TASK_ID],
            http_code=WEBSOCKET_ERROR_CODE,
            name=details[HEADER][ERROR_NAME],
            message=details[HEADER][ERROR_MESSAGE],
        )
        logger.error(error)
        raise error

    async def _start_task(self, ws):
        if self.pre_task_id is not None:
            self.task_headers["task_id"] = self.pre_task_id
        else:
            self.task_headers["task_id"] = uuid.uuid4().hex  # create task id.
        task_header = {**self.task_headers, ACTION_KEY: ActionType.START}
        # for binary data, the start action has no input, only parameters.
        start_data = self.data.get_websocket_start_data()
        message = self._build_up_message(task_header, start_data)
        logger.debug("Send start task: %s", message)
        await ws.send_str(message)

    async def _send_finished_task(self, ws):
        task_header = {**self.task_headers, ACTION_KEY: ActionType.FINISHED}
        payload = {"input": {}}
        message = self._build_up_message(task_header, payload)
        logger.debug("Send finish task: %s", message)
        await ws.send_str(message)

    async def _send_continue_task_data(self, ws):
        headers = {
            "task_id": self.task_headers["task_id"],
            "action": "continue-task",
        }
        for input_item in self.data.get_websocket_continue_data():
            if self.is_binary_input:
                if len(input_item) > 0:
                    if isinstance(input_item, bytes):
                        await ws.send_bytes(input_item)
                        logger.debug(
                            "Send continue task with bytes: %s",
                            len(input_item),
                        )
                    else:
                        await ws.send_bytes(list(input_item.values())[0])
                        logger.debug(
                            "Send continue task with list[byte]: %s",
                            len(input_item),
                        )
            else:
                if len(input_item) > 0:
                    message = self._build_up_message(
                        headers=headers,
                        payload=input_item,
                    )
                    logger.debug("Send continue task: %s", message)
                    await ws.send_str(message)
            await asyncio.sleep(0.000001)

        # data send completed, and send task completed.
        await self._send_finished_task(ws)

    async def _receive_batch_data_task(
        self,
        ws,
    ) -> Tuple[bool, Union[str, bytes]]:
        """_summary_

        Args:
            ws (connection): The ws connection.

        Raises:
            UnknownMessageReceived: The message is unexpected.

        Returns:
            Tuple[bool, str]: is output is binary, output
        """
        while True:
            msg = await ws.receive()
            await self._check_websocket_unexpected_message(msg)
            if msg.type == aiohttp.WSMsgType.TEXT:
                msg_json = msg.json()
                logger.debug("Receive %s event", msg_json[HEADER][EVENT_KEY])
                if msg_json[HEADER][EVENT_KEY] == EventType.GENERATED:
                    payload = msg_json["payload"]
                    return False, payload
                elif msg_json[HEADER][EVENT_KEY] == EventType.FINISHED:
                    payload = msg_json["payload"]
                    return False, payload
                elif msg_json[HEADER][EVENT_KEY] == EventType.FAILED:
                    self._on_failed(msg_json)
                else:
                    error = f"Receive unknown message: {msg_json}"
                    logger.error(error)
                    raise UnknownMessageReceived(error)
            elif msg.type == aiohttp.WSMsgType.BINARY:
                return True, msg.data  # get binary result data.

    async def _wait_for_task_started(self, ws):
        while True:
            msg = await ws.receive()
            await self._check_websocket_unexpected_message(msg)
            if msg.type == aiohttp.WSMsgType.TEXT:
                msg_json = msg.json()
                logger.debug("Receive %s event", msg_json[HEADER][EVENT_KEY])
                if msg_json[HEADER][EVENT_KEY] == EventType.STARTED:
                    return
                elif msg_json[HEADER][EVENT_KEY] == EventType.FAILED:
                    self._on_failed(msg_json)
                else:
                    raise UnexpectedMessageReceived(
                        "Receive unexpected message, expect task-started, "
                        f"real: {msg_json[HEADER][EVENT_KEY]}.",
                    )
            elif msg.type == aiohttp.WSMsgType.BINARY:
                raise UnexpectedMessageReceived(
                    "Receive unexpected binary message when wait for task-started",  # noqa E501
                )

    async def _check_websocket_unexpected_message(self, msg):
        if msg.type == aiohttp.WSMsgType.CLOSED:
            details = f"WSMsgType.CLOSE, data: {msg.data}, extra: {msg.extra}"
            logger.error("Connection unexpected closed!")
            raise UnexpectedMessageReceived(
                f"Receive unexpected websocket close message, "
                f"details: {details}",
            )
        if msg.type == aiohttp.WSMsgType.ERROR:
            details = f"WSMsgType.ERROR, data: {msg.data}, extra: {msg.extra}"
            logger.error("Connection error: %s", details)
            raise UnexpectedMessageReceived(
                f"Receive unexpected websocket error message "
                f"details: {details}.",
            )

    def _build_up_message(self, headers, payload):
        message = {"header": headers, "payload": payload}
        return json.dumps(message)
