# Copyright (c) ModelScope Contributors. All rights reserved.
import asyncio
import hashlib
import inspect
import json
import os
import pickle
import time
import torch
import torch.nn.functional as F
import transformers
from copy import deepcopy
from packaging import version
from PIL import Image
from queue import Queue
from threading import Thread
from torch import nn
from tqdm import tqdm
from transformers import GenerationConfig, LogitsProcessorList
from transformers.utils import is_torch_npu_available
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union

from swift.metrics import Metric
from swift.model import get_model_processor
from swift.template import Template
from swift.tuners import Swift
from swift.utils import get_last_valid_indices, patch_kernels, safe_snapshot_download, to_device
from .infer_engine import InferEngine
from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
                       ChatCompletionStreamResponse, ChatMessage, DeltaMessage, EmbeddingResponse,
                       EmbeddingResponseData, InferRequest, RequestConfig, random_uuid)
from .utils import AdapterRequest, InferStreamer, LogitsStreamer, TokensIteratorStreamer, prepare_generation_config

_TRANSFORMERS_GE_5_2 = version.parse(transformers.__version__) >= version.parse('5.2.0')
_kernels_patched = False


class _GenerationConfig(GenerationConfig):

    def __repr__(self) -> str:
        parameters = inspect.signature(self.to_json_string).parameters
        kwargs = {}
        if 'ignore_metadata' in parameters:
            kwargs['ignore_metadata'] = True
        gen_kwargs = json.loads(self.to_json_string(**kwargs))
        gen_kwargs.pop('transformers_version', None)
        return f'GenerationConfig({gen_kwargs})'


class TransformersEngine(InferEngine):

    def __init__(
            self,
            model: Union[str, nn.Module],
            *,
            template: Optional[Template] = None,
            adapters: Optional[List[str]] = None,
            max_batch_size: int = 1,  # 0/1: no limit
            reranker_use_activation: bool = True,
            # model kwargs
            torch_dtype: Optional[torch.dtype] = None,
            model_type: Optional[str] = None,
            attn_impl: Optional[str] = None,
            experts_impl: Optional[str] = None,
            device_map: Optional[Union[str, Dict[str, Any]]] = None,
            task_type: Optional[str] = None,
            quantization_config=None,
            model_kwargs: Optional[Dict[str, Any]] = None,
            template_type: Optional[str] = None,
            # hub kwargs
            use_hf: Optional[bool] = None,
            revision: Optional[str] = None,
            hub_token: Optional[str] = None,
            **kwargs):
        if isinstance(adapters, str):
            adapters = [adapters]
        self.adapters = adapters or []
        self.max_batch_size = max_batch_size
        self.reranker_use_activation = reranker_use_activation

        self.torch_dtype = torch_dtype
        self.model_type = model_type
        self.attn_impl = attn_impl
        self.experts_impl = experts_impl
        self.device_map = device_map
        self.task_type = task_type
        self.quantization_config = quantization_config
        self.model_kwargs = model_kwargs

        self.use_hf = use_hf
        self.revision = revision
        self.hub_token = hub_token
        global _kernels_patched
        if _TRANSFORMERS_GE_5_2 and not _kernels_patched:
            if use_hf is not None and 'USE_HF' not in os.environ:
                os.environ['USE_HF'] = str(use_hf)
            patch_kernels()
            _kernels_patched = True
        if isinstance(model, str):
            self.model, processor = self._get_model_processor(model, **kwargs)
            template = self._get_template(processor, template_type=template_type)
        elif isinstance(model, nn.Module):
            self.model = model
            if template is None:
                raise ValueError('`template` is required when `model` is a nn.Module')
        super().__init__(template)
        for adapter in self.adapters:
            self._add_adapter(safe_snapshot_download(adapter, use_hf=self.use_hf, hub_token=self.hub_token))
        self.engine = self.model  # dummy
        self.generation_config = getattr(self.model, 'generation_config', None)
        self._queue = Queue()
        self._task_pool = {}
        self._adapters_pool = {}
        self._task_thread = None

    def _get_model_processor(self, model_id_or_path, **kwargs):
        return get_model_processor(
            model_id_or_path,
            torch_dtype=self.torch_dtype,
            model_type=self.model_type,
            use_hf=self.use_hf,
            hub_token=self.hub_token,
            revision=self.revision,
            device_map=self.device_map,
            quantization_config=self.quantization_config,
            attn_impl=self.attn_impl,
            experts_impl=self.experts_impl,
            task_type=self.task_type,
            model_kwargs=self.model_kwargs,
            **kwargs)

    def _start_infer_worker(self):
        self._task_thread = Thread(target=self._infer_worker, daemon=True)
        self._task_thread.start()

    def _fetch_infer_requests(self):
        while not self._queue.empty():
            infer_request, kwargs, queue = self._queue.get()
            info = hashlib.sha256(pickle.dumps((kwargs['request_config']))).hexdigest()
            if info not in self._task_pool:
                self._task_pool[info] = kwargs, []
            self._task_pool[info][1].append((infer_request, queue))
        if len(self._task_pool) == 0:
            return
        key, (kwargs, data) = next(iter(self._task_pool.items()))
        max_batch_size = self.max_batch_size
        if max_batch_size <= 0:
            max_batch_size = len(data)
        data, remain_data = data[:max_batch_size], data[max_batch_size:]
        if remain_data:
            self._task_pool[key] = kwargs, remain_data
        else:
            self._task_pool.pop(key)
        kwargs = kwargs.copy()
        kwargs['infer_requests'] = [d[0] for d in data]
        queue_list = [d[1] for d in data]
        return kwargs, queue_list

    def _infer_worker(self):
        while True:
            time.sleep(0.01)
            item = self._fetch_infer_requests()
            if item is not None:
                kwargs, queue_list = item
                request_config = kwargs['request_config']
                res_list_or_gen = self._infer(**kwargs)
                if request_config.stream:
                    finished = False
                    while not finished:
                        try:
                            res_list = next(res_list_or_gen)
                        except StopIteration:
                            finished = True
                            res_list = [None] * len(queue_list)
                        for (queue, loop), res in zip(queue_list, res_list):
                            asyncio.run_coroutine_threadsafe(queue.put(res), loop)
                else:
                    for (queue, loop), res in zip(queue_list, res_list_or_gen):
                        asyncio.run_coroutine_threadsafe(queue.put(res), loop)

    def _add_adapter(self, adapter_path: str, adapter_name: Optional[str] = None) -> None:
        self.model = Swift.from_pretrained(self.model, adapter_path, adapter_name)

    def _prepare_generation_config(self, request_config: RequestConfig) -> _GenerationConfig:
        generation_config = prepare_generation_config(self.generation_config, request_config, self.tokenizer)
        generation_config.return_dict_in_generate = True
        if request_config.logprobs:
            generation_config.output_logits = True
        generation_config.num_return_sequences = request_config.n
        return _GenerationConfig(**generation_config.to_dict())

    def _add_stop_words(self, generation_config: _GenerationConfig, request_config: RequestConfig) -> None:
        template_meta = self.template.template_meta
        stop_words = (request_config.stop or []) + template_meta.stop_words
        generation_config.stop_words = self._get_stop_words(stop_words)

    @staticmethod
    def preprocess_logits(batched_logits: Optional[List[torch.Tensor]], batched_generate_ids: torch.Tensor,
                          top_logprobs: Optional[int]):
        top_logprobs = top_logprobs or 1
        batch_size = batched_generate_ids.shape[0]
        if batched_logits is None:
            return None
        batched_logprobs = []
        for i in range(batch_size):
            logprobs_list = []
            generate_ids = batched_generate_ids[i]
            for j, logits in enumerate(batched_logits):
                token = generate_ids[j].item()
                logprobs = torch.log_softmax(logits[i], -1)
                tokens = [token] + logprobs.argsort(descending=True, dim=-1)[:top_logprobs].tolist()
                logprobs_list.append({token: logprobs[token].item() for token in tokens})
            batched_logprobs.append(logprobs_list)
        return batched_logprobs

    @staticmethod
    def _update_batched_logprobs(batched_logprobs: List[torch.Tensor], logits_streamer: Optional[LogitsStreamer],
                                 generate_ids: torch.Tensor, top_logprobs: int) -> None:
        seq_len = generate_ids.shape[1] - len(batched_logprobs[0])
        if logits_streamer is None or seq_len == 0:
            return

        res = []
        for i in range(seq_len):
            res.append(logits_streamer.queue.get())
        new_batched_logprobs = TransformersEngine.preprocess_logits(res, generate_ids[:, -seq_len:], top_logprobs)
        for logprobs, new_logprobs in zip(batched_logprobs, new_batched_logprobs):
            logprobs += new_logprobs

    def _infer_stream(self, inputs: Dict[str, Any], *, generation_config: GenerationConfig,
                      adapter_request: Optional[AdapterRequest], request_config: RequestConfig,
                      template_inputs) -> Iterator[List[Optional[ChatCompletionStreamResponse]]]:
        if generation_config.num_beams != 1:
            error_msg = 'Streaming generation does not support beam search.'
            raise ValueError(error_msg)
        streamer = TokensIteratorStreamer()
        generate_kwargs = {
            'generation_config': generation_config,
            'streamer': streamer,
            **inputs,
        }
        adapter_names = self._get_adapter_names(adapter_request)
        if adapter_names is not None:
            generate_kwargs['adapter_names'] = adapter_names
        num_prompt_tokens = self._get_num_tokens(inputs)

        logits_streamer = None
        if generation_config.output_logits:
            generate_kwargs['logits_processor'] = LogitsProcessorList([LogitsStreamer()])

        def _model_generate(**kwargs):
            if is_torch_npu_available():
                torch.npu.set_device(self.model.device)
            self.template.generate(self.model, **kwargs)

        generate_kwargs = self.template.prepare_generate_kwargs(generate_kwargs, model=self.model)
        thread = Thread(target=_model_generate, kwargs=generate_kwargs)
        thread.start()
        batch_size = inputs['attention_mask'].shape[0]
        all_is_finished = False
        is_finished = [False] * batch_size
        infer_streamers = [InferStreamer(self.template, template_inputs=template_inputs[i]) for i in range(batch_size)]
        request_id_list = [f'chatcmpl-{random_uuid()}' for _ in range(batch_size)]
        token_idxs = [0] * batch_size

        raw_batched_generate_ids = None  # or torch.Tensor: [batch_size, seq_len]
        batched_logprobs = [[] for _ in range(batch_size)]
        while not all_is_finished:
            try:
                batched_tokens = next(streamer)
                if batched_tokens.ndim == 1:
                    batched_tokens = batched_tokens[:, None]

                raw_batched_generate_ids = torch.concat(
                    [batched_tokens]
                    if raw_batched_generate_ids is None else [raw_batched_generate_ids, batched_tokens],
                    dim=1)
            except StopIteration:
                all_is_finished = True

            batched_generate_ids = self.template.get_generate_ids(raw_batched_generate_ids, num_prompt_tokens)
            self._update_batched_logprobs(batched_logprobs, logits_streamer, batched_generate_ids,
                                          request_config.top_logprobs)

            res = []
            for i in range(batched_generate_ids.shape[0]):
                if is_finished[i]:
                    res.append(None)
                    continue
                generate_ids = batched_generate_ids[i]

                # ignore pad_token
                masks = generate_ids != self.tokenizer.pad_token_id
                generate_ids = generate_ids[masks].tolist()
                logprobs_list = None
                if batched_logprobs[i]:
                    logprobs_list = [logprobs for m, logprobs in zip(masks, batched_logprobs[i]) if m.item()]

                is_finished[i] = (
                    all_is_finished or is_finished[i]
                    or len(generate_ids) > 0 and generate_ids[-1] == self.tokenizer.pad_token_id)
                delta_text = infer_streamers[i].get_printable_text(generate_ids, is_finished[i])
                if not delta_text and not is_finished[i]:
                    res.append(None)
                    continue
                logprobs = self._get_logprobs(logprobs_list, generate_ids[token_idxs[i]:], request_config.top_logprobs)
                token_idxs[i] = len(generate_ids)

                usage_info = self._get_usage_info(num_prompt_tokens, len(generate_ids))
                toolcall = None
                if is_finished[i]:
                    toolcall = self._get_toolcall(
                        self.template.decode_generate_ids(generate_ids, template_inputs=template_inputs[i]))
                finish_reason = self._get_finish_reason(generation_config.max_new_tokens, usage_info.completion_tokens,
                                                        is_finished[i])

                choices = [
                    ChatCompletionResponseStreamChoice(
                        index=0,
                        delta=DeltaMessage(role='assistant', content=delta_text, tool_calls=toolcall),
                        finish_reason=finish_reason,
                        logprobs=logprobs)
                ]
                res.append(
                    ChatCompletionStreamResponse(
                        model=self.model_name, choices=choices, usage=usage_info, id=request_id_list[i]))
            if any(res):
                yield res

    def _get_adapter_names(self, adapter_request: Optional[AdapterRequest]) -> Optional[List[str]]:
        if adapter_request is None:
            if self._adapters_pool:
                return ['__base__']
            return
        adapter_name = adapter_request.name
        if adapter_name not in self._adapters_pool:
            self._adapters_pool[adapter_name] = adapter_request
            self._add_adapter(adapter_request.path, adapter_name)
        return [adapter_name]

    def _infer_forward(self, inputs: Dict[str, Any], adapter_request: Optional[AdapterRequest],
                       request_config: RequestConfig, **kwargs):
        call_kwargs = {}
        top_logprobs = request_config.top_logprobs or 20
        adapter_names = self._get_adapter_names(adapter_request)
        if adapter_names is not None:
            call_kwargs['adapter_names'] = adapter_names
        num_prompt_tokens = self._get_num_tokens(inputs)
        inputs.pop('labels', None)
        output = self.model(**inputs, **call_kwargs)
        if hasattr(output, 'logits'):
            logits = output.logits
        elif 'last_hidden_state' in output:
            # embeddings
            logits = output['last_hidden_state']
        else:
            raise NotImplementedError('Only support `logits` or `hidden_state` in output.')
        task_type = self.template.task_type
        if task_type == 'seq_cls':
            preds, logprobs = self.template.decode_seq_cls(logits, top_logprobs)
        elif task_type == 'prm':
            preds = self.template.decode_prm(inputs['input_ids'], logits)
            logprobs = [None] * len(preds)
        elif task_type == 'embedding':
            preds = logits
            logprobs = [None] * len(preds)
        elif task_type in ('reranker', 'generative_reranker'):
            if task_type == 'generative_reranker':
                attention_mask = inputs.get('attention_mask')
                last_valid_indices = -1 if attention_mask is None else get_last_valid_indices(attention_mask)
                batch_indices = torch.arange(logits.shape[0], device=logits.device)
                logits = logits[batch_indices, last_valid_indices]
            preds = logits.float()
            if self.reranker_use_activation:
                preds = F.sigmoid(preds)
            preds = preds.tolist()
            logprobs = [None] * len(preds)
        else:
            raise ValueError(f'Unsupported task_type: {task_type}')

        res = []
        for i, pred in enumerate(preds):
            usage_info = self._get_usage_info(num_prompt_tokens, 1)
            if task_type == 'embedding':
                res.append(
                    EmbeddingResponse(
                        model=self.model_name, usage=usage_info, data=[EmbeddingResponseData(embedding=pred.tolist())]))
            else:
                choices = [
                    ChatCompletionResponseChoice(
                        index=0,
                        message=ChatMessage(role='assistant', content=pred, tool_calls=None),
                        finish_reason='stop',
                        logprobs=logprobs[i])
                ]
                res.append(ChatCompletionResponse(model=self.model_name, choices=choices, usage=usage_info))
        return res

    def _infer_full(self, inputs: Dict[str, Any], *, generation_config: GenerationConfig,
                    adapter_request: Optional[AdapterRequest], request_config: RequestConfig,
                    template_inputs) -> List[ChatCompletionResponse]:
        # bos_token TODO: encoder-decoder
        generate_kwargs = {'generation_config': generation_config, **inputs}
        adapter_names = self._get_adapter_names(adapter_request)
        if adapter_names is not None:
            generate_kwargs['adapter_names'] = adapter_names
        num_prompt_tokens = self._get_num_tokens(inputs)
        generate_kwargs = self.template.prepare_generate_kwargs(generate_kwargs, model=self.model)
        output = dict(self.template.generate(self.model, **generate_kwargs))
        output.pop('past_key_values', None)
        batched_generate_ids = output['sequences']
        batched_generate_ids = self.template.get_generate_ids(batched_generate_ids, num_prompt_tokens)
        self.template.debug_logger({'generate_ids': batched_generate_ids})  # debug
        batched_logprobs = self.preprocess_logits(
            output.get('logits'), batched_generate_ids, request_config.top_logprobs)

        res = []
        num_return_sequences = generation_config.num_return_sequences
        for i in range(inputs['attention_mask'].shape[0]):
            choices = []
            usage_info = self._get_usage_info(num_prompt_tokens, 0)
            for j in range(num_return_sequences):
                batched_index = i * num_return_sequences + j
                generate_ids = batched_generate_ids[batched_index]

                # ignore pad_token
                masks = generate_ids != self.tokenizer.pad_token_id
                generate_ids = generate_ids[masks].tolist()
                logprobs_list = None
                if batched_logprobs is not None:
                    logprobs_list = [
                        logprobs for m, logprobs in zip(masks, batched_logprobs[batched_index]) if m.item()
                    ]

                logprobs = self._get_logprobs(logprobs_list, generate_ids, request_config.top_logprobs)
                usage_info = self._update_usage_info(usage_info, len(generate_ids))
                response = self.template.decode_generate_ids(generate_ids, template_inputs=template_inputs[i])
                finish_reason = self._get_finish_reason(generation_config.max_new_tokens, len(generate_ids), True)
                toolcall = self._get_toolcall(response)
                token_ids = generate_ids if request_config.return_details else None
                choices.append(
                    ChatCompletionResponseChoice(
                        index=j,
                        message=ChatMessage(role='assistant', content=response, tool_calls=toolcall),
                        finish_reason=finish_reason,
                        logprobs=logprobs,
                        token_ids=token_ids))
            prompt_token_ids = None
            images_size = None
            if request_config.return_details:
                if 'input_ids' in inputs:
                    non_pad_indices = (inputs['input_ids'][i] != self.tokenizer.pad_token_id).nonzero()
                    if non_pad_indices.numel() > 0:
                        idx = non_pad_indices.min().item()
                        prompt_token_ids = inputs['input_ids'][i][idx:].tolist()
                if all(isinstance(image, Image.Image) for image in template_inputs[i].images):
                    images_size = [image.size for image in template_inputs[i].images]
            res.append(
                ChatCompletionResponse(
                    model=self.model_name,
                    choices=choices,
                    usage=usage_info,
                    prompt_token_ids=prompt_token_ids,
                    images_size=images_size))
        return res

    async def infer_async(
        self,
        infer_request: InferRequest,
        request_config: Optional[RequestConfig] = None,
        *,
        adapter_request: Optional[AdapterRequest] = None,
        pre_infer_hook=None,
    ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
        if request_config is None:
            request_config = RequestConfig()
        queue = asyncio.Queue()
        self._queue.put((infer_request, {
            'request_config': request_config,
            'adapter_request': adapter_request,
            'pre_infer_hook': pre_infer_hook
        }, (queue, asyncio.get_event_loop())))
        await asyncio.sleep(0)
        if self._task_thread is None:
            self._start_infer_worker()
        if request_config.stream:

            async def _gen_wrapper():
                while True:
                    item = await queue.get()
                    await asyncio.sleep(0)
                    if item is None:
                        break
                    yield item

            return _gen_wrapper()
        else:
            return await queue.get()

    # Ensure `template._post_encode` has no gradient.
    @torch.inference_mode()
    def _infer(
        self,
        infer_requests: List[InferRequest],
        request_config: RequestConfig,
        *,
        adapter_request: Optional[AdapterRequest] = None,
        pre_infer_hook=None,
    ) -> Union[List[ChatCompletionResponse], Iterator[List[Optional[ChatCompletionStreamResponse]]]]:
        self.model.eval()
        request_config = deepcopy(request_config)
        if self.template.use_model:
            self.template.model = self.model

        if self.model_info.task_type == 'causal_lm':
            self.template.set_mode('transformers')

        batched_inputs, error_list = self._batch_encode(infer_requests, strict=getattr(self, 'strict', True))
        if len(batched_inputs) > 0:
            template_inputs = [inputs.pop('template_inputs') for inputs in batched_inputs]
            inputs = to_device(self.template.data_collator(batched_inputs), self.model.device)
            self.template.debug_logger(inputs)  # debug
            if self.model_meta.is_multimodal:
                _, inputs = self.template.pre_forward_hook(self.model, None, inputs)
            if self.model_info.task_type == 'causal_lm':
                self.set_default_max_tokens(request_config, inputs)
                generation_config = self._prepare_generation_config(request_config)
                self._add_stop_words(generation_config, request_config)
            else:
                generation_config = request_config

            kwargs = {
                'inputs': inputs,
                'generation_config': generation_config,
                'adapter_request': adapter_request,
                'request_config': request_config,
                'template_inputs': template_inputs,
            }
            if pre_infer_hook:
                kwargs = pre_infer_hook(kwargs)
        else:
            kwargs = {}
        if request_config.stream:

            def _gen_wrapper():
                if len(kwargs) > 0:
                    for res in self._infer_stream(**kwargs):
                        yield self._add_error_list(res, error_list)
                else:
                    yield self._add_error_list([], error_list)

            return _gen_wrapper()
        else:
            if len(kwargs) > 0:
                infer_func = self._infer_forward if self.template.task_type in {
                    'seq_cls', 'prm', 'embedding', 'reranker', 'generative_reranker'
                } else self._infer_full
                res = infer_func(**kwargs)
            else:
                res = []
            return self._add_error_list(res, error_list)

    def infer(
        self,
        infer_requests: List[InferRequest],
        request_config: Optional[RequestConfig] = None,
        metrics: Optional[List[Metric]] = None,
        *,
        use_tqdm: Optional[bool] = None,
        adapter_request: Optional[AdapterRequest] = None
    ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
        if request_config is None:
            request_config = RequestConfig()
        if request_config.stream:
            return super().infer(
                infer_requests, request_config, metrics, use_tqdm=use_tqdm, adapter_request=adapter_request)
        # Has higher stability than calling super().infer
        if use_tqdm is None:
            use_tqdm = not request_config.stream and len(infer_requests) > 1
        prog_bar = tqdm(total=len(infer_requests), dynamic_ncols=True, disable=not use_tqdm)
        # If self.max_batch_size <= 0, then process all infer_requests at once.
        max_batch_size = self.max_batch_size
        if max_batch_size <= 0:
            max_batch_size = len(infer_requests)
        res = []
        i = 0
        while i < len(infer_requests):
            infer_requests_samples = infer_requests[i:i + max_batch_size]
            res += self._infer(infer_requests_samples, request_config, adapter_request=adapter_request)
            i += max_batch_size
            prog_bar.update(len(infer_requests_samples))
        prog_bar.close()
        self._update_metrics(res, metrics)
        return res
