# region youmi
import copy
from http import HTTPStatus
from pprint import pformat
from typing import Dict, Iterator, List, Literal, Optional, Union

import dashscope
from qwen_agent.llm.base import (
    ModelServiceError,
    _truncate_input_messages_roughly,
    register_llm,
)
from qwen_agent.llm.schema import (
    ASSISTANT,
    DEFAULT_SYSTEM_MESSAGE,
    SYSTEM,
    USER,
    Message,
)
from qwen_agent.llm.text_base import BaseTextChatModel
from qwen_agent.log import logger
from qwen_agent.settings import DEFAULT_MAX_INPUT_TOKENS
from qwen_agent.utils.utils import has_chinese_messages, merge_generate_cfgs

from rtp_llm.config.py_config_modules import StaticConfig

# end region


@register_llm("qwen_dashscope")
class QwenChatAtDS(BaseTextChatModel):

    def __init__(self, cfg: Optional[Dict] = None):
        super().__init__(cfg)
        self.model = self.model or "qwen-max"
        initialize_dashscope(cfg)

    def _chat_stream(
        self,
        messages: List[Message],
        delta_stream: bool,
        generate_cfg: dict,
    ) -> Iterator[List[Message]]:
        messages = [msg.model_dump() for msg in messages]
        logger.debug(f"LLM Input:\n{pformat(messages, indent=2)}")
        response = dashscope.Generation.call(
            self.model,
            messages=messages,  # noqa
            result_format="message",
            stream=True,
            **generate_cfg,
        )
        if delta_stream:
            return self._delta_stream_output(response)
        else:
            return self._full_stream_output(response)

    def _chat_no_stream(
        self,
        messages: List[Message],
        generate_cfg: dict,
    ) -> List[Message]:
        messages = [msg.model_dump() for msg in messages]
        logger.debug(f"LLM Input:\n{pformat(messages, indent=2)}")
        response = dashscope.Generation.call(
            self.model,
            messages=messages,  # noqa
            result_format="message",
            stream=False,
            **generate_cfg,
        )
        if response.status_code == HTTPStatus.OK:
            return [Message(ASSISTANT, response.output.choices[0].message.content)]
        else:
            raise ModelServiceError(code=response.code, message=response.message)

    def _continue_assistant_response(
        self,
        messages: List[Message],
        generate_cfg: dict,
        stream: bool,
    ) -> Iterator[List[Message]]:
        prompt = self._build_text_completion_prompt(messages)
        logger.debug(f"LLM Input:\n{pformat(prompt, indent=2)}")
        response = dashscope.Generation.call(
            self.model,
            prompt=prompt,  # noqa
            result_format="message",
            stream=True,
            use_raw_prompt=True,
            **generate_cfg,
        )
        it = self._full_stream_output(response)
        if stream:
            return it  # streaming the response
        else:
            *_, final_response = it  # return the final response without streaming
            return final_response

    # region youmi generate completion prompt
    def generate_completion_prompt(
        self,
        messages: List[Union[Message, Dict]],
        functions: Optional[List[Dict]] = None,
        extra_generate_cfg: Optional[Dict] = None,
    ) -> str:
        """copy from LLM chat interface.

        Args:
            messages: Inputted messages.
            functions: Inputted functions for function calling. OpenAI format supported.
        Returns:
            the generated prompt
        """
        generate_cfg = merge_generate_cfgs(
            base_generate_cfg=self.generate_cfg, new_generate_cfg=extra_generate_cfg
        )
        if "lang" in generate_cfg:
            lang: Literal["en", "zh"] = generate_cfg.pop("lang")
        else:
            lang: Literal["en", "zh"] = "zh" if has_chinese_messages(messages) else "en"

        messages = copy.deepcopy(messages)

        _return_message_type = "dict"
        new_messages = []
        for msg in messages:
            if isinstance(msg, dict):
                new_messages.append(Message(**msg))
            else:
                new_messages.append(msg)
                _return_message_type = "message"
        messages = new_messages

        if messages[0].role != SYSTEM:
            messages = [Message(role=SYSTEM, content=DEFAULT_SYSTEM_MESSAGE)] + messages

        # Not precise. It's hard to estimate tokens related with function calling and multimodal items.
        messages = _truncate_input_messages_roughly(
            messages=messages,
            max_tokens=generate_cfg.pop("max_input_tokens", DEFAULT_MAX_INPUT_TOKENS),
        )

        messages = self._preprocess_messages(messages, lang=lang)

        if "function_choice" in generate_cfg:
            fn_choice = generate_cfg["function_choice"]
            valid_fn_choices = [
                f.get("name", f.get("name_for_model", None)) for f in (functions or [])
            ]
            valid_fn_choices = ["auto", "none"] + [f for f in valid_fn_choices if f]
            if fn_choice not in valid_fn_choices:
                raise ValueError(
                    f"The value of function_choice must be one of the following: {valid_fn_choices}. "
                    f'But function_choice="{fn_choice}" is received.'
                )
            if fn_choice == "auto":
                del generate_cfg["function_choice"]
            if fn_choice == "none":
                raise NotImplementedError(
                    'Not implemented function_choice="none" yet.'
                )  # TODO:

        if functions:
            fncall_mode = True
        else:
            fncall_mode = False
            for k in ["parallel_function_calls", "function_choice"]:
                if k in generate_cfg:
                    del generate_cfg[k]

        if fncall_mode:
            messages = self._prepend_fncall_system(messages, functions, lang=lang)
            prompt = self._build_text_completion_prompt(messages)
            return prompt
        else:
            prompt = self._build_text_completion_prompt(messages)
            return prompt

    # end region

    @staticmethod
    def _build_text_completion_prompt(messages: List[Message]) -> str:
        im_start = "<|im_start|>"
        im_end = "<|im_end|>"
        if messages[0].role == SYSTEM:
            sys = messages[0].content
            assert isinstance(sys, str)
            prompt = f"{im_start}{SYSTEM}\n{sys}{im_end}"
        else:
            prompt = f"{im_start}{SYSTEM}\n{DEFAULT_SYSTEM_MESSAGE}{im_end}"
        if messages[-1].role != ASSISTANT:
            messages.append(Message(ASSISTANT, ""))
        for msg in messages:
            assert isinstance(msg.content, str)
            if msg.role == USER:
                query = msg.content.lstrip("\n").rstrip()
                prompt += f"\n{im_start}{USER}\n{query}{im_end}"
            elif msg.role == ASSISTANT:
                response = msg.content.lstrip("\n").rstrip()
                prompt += f"\n{im_start}{ASSISTANT}\n{response}{im_end}"
        assert prompt.endswith(im_end)
        prompt = prompt[: -len(im_end)]
        return prompt

    @staticmethod
    def _delta_stream_output(response) -> Iterator[List[Message]]:
        last_len = 0
        delay_len = 5
        in_delay = False
        text = ""
        for chunk in response:
            if chunk.status_code == HTTPStatus.OK:
                text = chunk.output.choices[0].message.content
                if (len(text) - last_len) <= delay_len:
                    in_delay = True
                    continue
                else:
                    in_delay = False
                    real_text = text[:-delay_len]
                    now_rsp = real_text[last_len:]
                    yield [Message(ASSISTANT, now_rsp)]
                    last_len = len(real_text)
            else:
                raise ModelServiceError(code=chunk.code, message=chunk.message)
        if text and (in_delay or (last_len != len(text))):
            yield [Message(ASSISTANT, text[last_len:])]

    @staticmethod
    def _full_stream_output(response) -> Iterator[List[Message]]:
        for chunk in response:
            if chunk.status_code == HTTPStatus.OK:
                yield [Message(ASSISTANT, chunk.output.choices[0].message.content)]
            else:
                raise ModelServiceError(code=chunk.code, message=chunk.message)


def initialize_dashscope(cfg: Optional[Dict] = None) -> None:
    cfg = cfg or {}

    api_key = cfg.get("api_key", "")
    base_http_api_url = cfg.get("base_http_api_url", None)
    base_websocket_api_url = cfg.get("base_websocket_api_url", None)

    if not api_key:
        api_key = StaticConfig.model_config.dashscope_api_key
    if not base_http_api_url:
        base_http_api_url = StaticConfig.model_config.dashscope_http_url
    if not base_websocket_api_url:
        base_websocket_api_url = StaticConfig.model_config.dashscope_websocket_url

    api_key = api_key.strip()
    dashscope.api_key = api_key
    if base_http_api_url is not None:
        dashscope.base_http_api_url = base_http_api_url.strip()
    if base_websocket_api_url is not None:
        dashscope.base_websocket_api_url = base_websocket_api_url.strip()
