# Copyright (c) ModelScope Contributors. All rights reserved.
import base64
import io
import json
import numpy as np
import os
import time
import uuid
from copy import deepcopy
from dataclasses import asdict, dataclass, field, fields
from PIL import Image
from pydantic import AfterValidator, BaseModel, Field, PlainSerializer, field_validator
from typing import Annotated, Any, Dict, List, Literal, Optional, Tuple, Union

from swift.template import Messages, Tool
from swift.utils import remove_response


def serialize_ndarray(value):
    if value is None:
        return None
    if isinstance(value, np.ndarray):
        return {
            'data': base64.b64encode(value.tobytes()).decode('ascii'),
            'shape': value.shape,
            'dtype': str(value.dtype),
            '__ndarray__': True
        }
    return value


def deserialize_ndarray(value):
    if value is None:
        return None
    if isinstance(value, dict) and value.get('__ndarray__'):
        data = base64.b64decode(value['data'])
        return np.frombuffer(data, dtype=value['dtype']).reshape(value['shape'])
    return value


NumpyArray = Annotated[Any, PlainSerializer(serialize_ndarray, return_type=Dict), AfterValidator(deserialize_ndarray)]


@dataclass
class InferRequest:
    """
    Data structure for inference requests.

    Attributes:
        messages (Messages):
            The input conversation in messages format. Each message is a dict containing at least
            a "role" field (e.g., "user", "assistant", "system") and a "content" field.
            Example:
                [{
                    "role": "user",
                    "content": [
                        {
                            "type": "image",  # can also be audio/video
                            "image": "<url/path/base64/PIL.Image>",
                        },
                        {"type": "text", "text": "Please describe the picture."},
                    ],
                }]
            The above is equivalent to:
                [{"role": "user", "content": "<image>Please describe the picture."}]
            with an additional argument:
                images = ["<url/path/base64/PIL.Image>"]

        images (List[Union[str, Image.Image]]):
            Optional, a list of images associated with the request.
            Each image can be a URL, local path, base64 string, or PIL.Image object.

        audios (List[str]):
            Optional, a list of audio resources associated with the request.

        videos (List[str]):
            Optional, a list of video resources associated with the request.

        tools (Optional[List[Tool]]):
            An optional list of tools. These should be organized in the agent_template format for
            tools requested by the system, for example 'react_en'.

        objects (Dict[str, Any]):
            Container for additional multimodal objects, grouped by type (key).
    """
    messages: Messages

    images: List[Union[str, Image.Image]] = field(default_factory=list)
    audios: List[str] = field(default_factory=list)
    videos: List[str] = field(default_factory=list)

    tools: Optional[List[Tool]] = None
    objects: Dict[str, Any] = field(default_factory=dict)
    chat_template_kwargs: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        for key in ['images', 'audios', 'videos']:
            val = getattr(self, key)
            if isinstance(val, str):
                setattr(self, key, [val])
        assert isinstance(self.messages, list), f'messages: {self.messages}'

    @staticmethod
    def remove_response(messages) -> Optional[str]:
        return remove_response(messages)

    @staticmethod
    def _to_printable(obj, key: Optional[str] = None):
        if isinstance(obj, str) and key not in {'content', 'text'} and len(obj) >= 1000:
            return f'<<<base64:{obj[:50]}..>>>'
        elif isinstance(obj, list):
            res = []
            for item in obj:
                res.append(InferRequest._to_printable(item))
            return res
        elif isinstance(obj, dict):
            res = {}
            for k, v in obj.items():
                res[k] = InferRequest._to_printable(v, key=k)
            return res
        return obj

    def to_printable(self):
        return InferRequest._to_printable(asdict(self))


@dataclass
class RolloutInferRequest(InferRequest):
    """
    An inference request class for rollout scenarios.

    This class extends `InferRequest` and specifically overrides the `images` attribute
    to be a list of strings for compatibility with POST requests. Each string may
    represent an image URL or a Base64-encoded image.

    Inherits all fields from `InferRequest`:
        messages (Messages):
            Input conversation messages, supporting multimodal content.
        audios (List[str]):
            List of audio resources associated with the request.
        videos (List[str]):
            List of video resources associated with the request.
        tools (Optional[List[Tool]]):
            List of tools, organized by the agent template (e.g. 'react_en').
        objects (Dict[str, Any]):
            Optional container for additional multimodal objects.

    Additional / Overridden fields:
        images (List[str]):
            List of image resources, each as a string (URL or base64).
        data_dict (Dict):
            Optional dictionary for extra request data.
        uuid (Optional[str]):
            Optional unique identifier for this request instance.
    """
    images: List[str] = field(default_factory=list)
    data_dict: Dict = field(default_factory=dict)
    uuid: Optional[str] = None


def random_uuid() -> str:
    return str(uuid.uuid4().hex)


@dataclass
class Model:
    id: str  # model_type

    object: str = 'model'
    created: int = field(default_factory=lambda: int(time.time()))
    owned_by: str = 'ms-swift'


@dataclass
class ModelList:
    data: List[Model]
    object: str = 'list'


@dataclass
class RequestConfig:
    """NOTE: The following behavior is inconsistent with the OpenAI API.
    Default values for OpenAI:
        temperature = 1.
        top_k = -1
        top_p = 1.
        repetition_penalty = 1.
    """
    max_tokens: Optional[int] = None  # None: max_model_len - num_tokens
    # None: use deploy_args
    temperature: Optional[float] = None
    top_k: Optional[int] = None
    top_p: Optional[float] = None
    repetition_penalty: Optional[float] = None
    num_beams: int = 1
    stop: Optional[List[str]] = field(default_factory=list)

    seed: Optional[int] = None
    stream: bool = False
    logprobs: bool = False
    top_logprobs: Optional[int] = None
    prompt_logprobs: Optional[int] = None

    n: int = 1
    best_of: Optional[int] = None
    presence_penalty: float = 0.
    frequency_penalty: float = 0.
    length_penalty: float = 1.
    # Return token_ids additionally (non-stream)
    return_details: bool = False
    # vLLM structured outputs (guided decoding)
    structured_outputs_regex: Optional[str] = None

    def __post_init__(self):
        if self.stop is None:
            self.stop = []


@dataclass
class CompletionRequestMixin:
    model: str
    prompt: str


@dataclass
class EmbeddingRequestMixin:
    input: str
    model: str
    encoding_format: Literal['float', 'base64'] = 'float'


@dataclass
class ChatCompletionRequestMixin:
    model: str
    messages: Messages
    tools: Optional[List[Tool]] = None
    tool_choice: Optional[Union[str, Dict]] = None
    chat_template_kwargs: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        if self.tool_choice is None:
            self.tool_choice = 'none' if self.tools is None else 'auto'

        if self.tools:
            if self.tool_choice == 'none':
                self.tools = None
            elif isinstance(self.tool_choice, dict):
                name = self.tool_choice['function']['name']
                tool = next(tool for tool in self.tools if tool['function']['name'] == name)
                if tool is None:
                    raise ValueError(f"Tool choice '{name}' not found in tools.")
                self.tools = [tool]


@dataclass
class MultiModalRequestMixin:
    images: List[str] = field(default_factory=list)
    audios: List[str] = field(default_factory=list)
    videos: List[str] = field(default_factory=list)
    objects: Dict[str, Any] = field(default_factory=dict)

    @staticmethod
    def to_base64(mm_data: Union[str, Image.Image, bytes]) -> str:
        if isinstance(mm_data, dict) and 'bytes' in mm_data:
            mm_data = mm_data['bytes'] or mm_data['path']
        if isinstance(mm_data, str) and not os.path.isfile(mm_data):
            # base64 or url
            return mm_data
        if isinstance(mm_data, str):
            # local_path
            with open(mm_data, 'rb') as f:
                bytes_ = f.read()
        elif isinstance(mm_data, Image.Image):
            bytes_io = io.BytesIO()
            mm_data.save(bytes_io, format='png')
            bytes_ = bytes_io.getvalue()
        else:
            bytes_ = mm_data
        img_base64: str = base64.b64encode(bytes_).decode('utf-8')
        return img_base64

    def __post_init__(self):
        for key in ['images', 'audios', 'videos']:
            values = getattr(self, key)
            if isinstance(values, str):
                values = [values]
                setattr(self, key, values)
            for i, val in enumerate(values):
                values[i] = self.to_base64(val)


@dataclass
class CompletionRequest(RequestConfig, MultiModalRequestMixin, CompletionRequestMixin):

    def __post_init__(self):
        RequestConfig.__post_init__(self)
        MultiModalRequestMixin.__post_init__(self)


@dataclass
class EmbeddingRequest(RequestConfig, MultiModalRequestMixin, EmbeddingRequestMixin):

    def __post_init__(self):
        RequestConfig.__post_init__(self)
        MultiModalRequestMixin.__post_init__(self)

    def parse(self) -> Tuple['InferRequest', 'RequestConfig']:
        data = asdict(self)
        res = []
        for cls_type in [InferRequest, RequestConfig]:
            parameters = set(f.name for f in fields(cls_type))
            _data = {k: v for k, v in data.items() if k in parameters}
            res.append(cls_type(**_data))
        return tuple(res)


@dataclass
class ChatCompletionRequest(RequestConfig, MultiModalRequestMixin, ChatCompletionRequestMixin):

    def __post_init__(self):
        RequestConfig.__post_init__(self)
        MultiModalRequestMixin.__post_init__(self)
        ChatCompletionRequestMixin.__post_init__(self)
        self.convert_to_base64()

    def convert_to_base64(self):
        for message in self.messages:
            content = message['content']
            if isinstance(content, str):
                continue
            for item in content:
                key: str = item['type']
                if key == 'text':
                    continue

                key_origin = key
                value = item[key]
                if key.endswith('_url'):
                    key = key[:-len('_url')]
                is_dict = False
                if isinstance(value, dict):
                    is_dict = True
                    value = value['url']
                if isinstance(value, str) and (value.startswith('data:') or value.startswith('http')
                                               or len(value) > 200):
                    continue

                # local_path / PIL.Image
                if isinstance(value, str) and os.path.isfile(value):
                    suffix = os.path.splitext(value)[1][1:].lower()
                elif isinstance(value, Image.Image):
                    suffix = 'jpeg'
                else:
                    raise ValueError(f'value: {value}')
                mm_data_base64 = self.to_base64(value)
                new_value = f'data:{key}/{suffix};base64,{mm_data_base64}'
                if is_dict:
                    new_value = {'url': new_value}
                item[key_origin] = new_value

    def parse(self) -> Tuple['InferRequest', 'RequestConfig']:
        data = asdict(self)
        res = []
        for cls_type in [InferRequest, RequestConfig]:
            parameters = set(f.name for f in fields(cls_type))
            _data = {k: v for k, v in data.items() if k in parameters}
            res.append(cls_type(**_data))
        return tuple(res)

    @classmethod
    def from_cmpl_request(cls, cmpl_request: Union[CompletionRequest, EmbeddingRequest]) -> 'ChatCompletionRequest':
        cmpl_request = asdict(cmpl_request)
        if 'prompt' in cmpl_request:
            prompt = cmpl_request.pop('prompt')
        else:
            prompt = cmpl_request.pop('input')
        cmpl_request['messages'] = [{'role': 'user', 'content': prompt}]
        if 'encoding_format' in cmpl_request:
            cmpl_request.pop('encoding_format')
        return cls(**cmpl_request)


@dataclass
class UsageInfo:
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int


@dataclass
class Function:
    name: str
    arguments: Optional[Union[str, Any]]

    def __post_init__(self):
        if not isinstance(self.arguments, str):
            self.arguments = json.dumps(self.arguments, ensure_ascii=False)
        self.name = self.name.strip()
        self.arguments = self.arguments.strip()


@dataclass
class ChatCompletionMessageToolCall:
    function: Function
    type: str = 'function'
    id: str = field(default_factory=lambda: f'toolcall-{random_uuid()}')


@dataclass
class ChatMessage:
    role: Literal['system', 'user', 'assistant']
    content: Union[str, List[Dict[str, Any]], int, float, List[float]]
    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
    reasoning_content: Optional[str] = None


@dataclass
class ChatCompletionResponseChoice:
    index: int
    message: ChatMessage
    finish_reason: Literal['stop', 'length', None]
    logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None
    token_ids: Optional[List[int]] = None
    routed_experts: Optional[NumpyArray] = None

    def to_cmpl_choice(self) -> 'CompletionResponseChoice':
        self = deepcopy(self)
        assert not self.message.tool_calls, f'message: {self.message}'
        return CompletionResponseChoice(self.index, self.message.content, self.finish_reason, self.logprobs)


@dataclass
class EmbeddingResponseData:
    object: str = 'embedding'
    index: int = 0
    embedding: List[str] = field(default_factory=lambda: [])


@dataclass
class EmbeddingResponse:
    model: str
    data: List[EmbeddingResponseData]
    usage: UsageInfo
    id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}')
    object: str = 'list'
    created: int = field(default_factory=lambda: int(time.time()))


@dataclass
class CompletionResponseChoice:
    index: int
    text: str
    finish_reason: Literal['stop', 'length', None]
    logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None


@dataclass
class ChatCompletionResponse:
    model: str
    choices: List[ChatCompletionResponseChoice]
    usage: UsageInfo
    id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}')
    object: str = 'chat.completion'
    created: int = field(default_factory=lambda: int(time.time()))
    prompt_token_ids: Optional[List[int]] = None
    prompt_logprobs: Optional[List] = None
    images_size: Optional[List[Tuple[int, int]]] = None

    def to_cmpl_response(self) -> 'CompletionResponse':
        self = deepcopy(self)
        choices = [choice.to_cmpl_choice() for choice in self.choices]
        id_ = f'cmpl{self.id[len("chatcmpl"):]}'
        return CompletionResponse(
            self.model,
            choices,
            self.usage,
            id_,
            created=self.created,
            prompt_token_ids=self.prompt_token_ids,
            prompt_logprobs=self.prompt_logprobs,
        )


class RolloutOutput(BaseModel):
    """
    Output structure for rollout.

    Attributes:
        response (ChatCompletionResponse):
            The model's response

        messages (Optional[Messages]):
            (Optional) Conversation history for the final rollout; required for multi-turn scenarios.
            NOTE:
                - If provided, this messages sequence will overwrite the original messages.
                - If not provided, 'response' will be appended as the latest turn in the original messages.
                - For multi-turn training, you need to manually return the updated messages, including the full history.
                - The messages should include the latest assistant response as the final message.

        response_token_ids (Optional[List[List[int]]]):
            (Optional) Token IDs generated at each rollout turn.
            If provided, the training process will skip tokenizing the response.

        response_loss_mask (Optional[List[List[int]]]):
            (Optional) Loss masks corresponding to each rollout turn.
            If provided, the training process will skip computing loss masks for the response (as controlled by the `loss_scale` parameter). # noqa

        rollout_infos (Dict[str, Any]):
            (Optional) Additional rollout information. This must be JSON-serializable.
    """
    response: ChatCompletionResponse
    # multi turn
    messages: Optional[Messages] = None
    response_token_ids: List[List[int]] = Field(default_factory=list)
    response_loss_mask: List[List[int]] = Field(default_factory=list)
    rollout_infos: Dict[str, Any] = Field(default_factory=dict)
    # rollout logprobs for each turn (used for rollout importance sampling correction in multi-turn scenarios)
    rollout_logprobs: List[List[float]] = Field(default_factory=list)

    prompt_logprobs: Optional[List] = None

    @field_validator('response_token_ids', 'response_loss_mask', 'rollout_logprobs', mode='before')
    @classmethod
    def _wrap_flat_list(cls, v):
        if isinstance(v, list) and v and isinstance(v[0], (int, float)):
            return [v]
        return v

    def model_post_init(self, __context):
        # Ensure multimodal data in rollout_infos is serializable (e.g., images to base64)
        super().model_post_init(__context)
        self.mminfo_to_serializable()

    def mminfo_to_serializable(self):
        mm_keys = ['images', 'audios', 'videos']

        for key, values in self.rollout_infos.items():
            if key in mm_keys:
                if not isinstance(values, list):
                    values = [values]
                for i, value in enumerate(values):
                    values[i] = MultiModalRequestMixin.to_base64(value)
                self.rollout_infos[key] = values


@dataclass
class CompletionResponse:
    model: str
    choices: List[CompletionResponseChoice]
    usage: UsageInfo
    id: str = field(default_factory=lambda: f'cmpl-{random_uuid()}')
    object: str = 'text_completion'
    created: int = field(default_factory=lambda: int(time.time()))
    prompt_token_ids: Optional[List[int]] = None
    prompt_logprobs: Optional[List] = None


@dataclass
class DeltaMessage:
    role: Literal['system', 'user', 'assistant', None] = None
    content: Optional[str] = None
    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
    reasoning_content: Optional[str] = None


@dataclass
class ChatCompletionResponseStreamChoice:
    index: int
    delta: DeltaMessage
    finish_reason: Literal['stop', 'length', None]
    logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None

    def to_cmpl_choice(self) -> 'CompletionResponseStreamChoice':
        self = deepcopy(self)
        assert not self.delta.tool_calls
        return CompletionResponseStreamChoice(self.index, self.delta.content, self.finish_reason, self.logprobs)


@dataclass
class CompletionResponseStreamChoice:
    index: int
    text: str
    finish_reason: Literal['stop', 'length', None]
    logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None


@dataclass
class ChatCompletionStreamResponse:
    model: str
    choices: List[ChatCompletionResponseStreamChoice]
    usage: Optional[UsageInfo] = None
    id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}')
    object: str = 'chat.completion.chunk'
    created: int = field(default_factory=lambda: int(time.time()))

    def to_cmpl_response(self) -> 'CompletionStreamResponse':
        self = deepcopy(self)
        choices = [choice.to_cmpl_choice() for choice in self.choices]
        id_ = f'cmpl{self.id[len("chatcmpl"):]}'
        return CompletionStreamResponse(self.model, choices, self.usage, id_, created=self.created)


@dataclass
class CompletionStreamResponse:
    model: str
    choices: List[CompletionResponseStreamChoice]
    usage: Optional[UsageInfo] = None
    id: str = field(default_factory=lambda: f'cmpl-{random_uuid()}')
    object: str = 'text_completion.chunk'
    created: int = field(default_factory=lambda: int(time.time()))


class InitCommunicatorRequest(BaseModel):
    host: str
    port: int
    world_size: int


class UpdateWeightsRequest(BaseModel):
    name: str
    dtype: str
    shape: list[int]
