# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import os
import json
import pickle
from typing import Iterable, Tuple

import torch

import deepspeed.comm as dist

from deepspeed.accelerator import get_accelerator
from deepspeed.comm.comm import init_distributed

from .model_implementations import InferenceV2Policy
from .logging import inference_logger
from .ragged import DSStateManager, RaggedBatchWrapper, PlaceholderSequenceDescriptor
from .scheduling_utils import SchedulingError, SchedulingResult
from .model_implementations.flat_model_helpers import make_param_filename, make_metadata_filename
from .model_implementations.inference_model_base import DSInferenceModelBase

from .config_v2 import RaggedInferenceEngineConfig

INFERENCE_MODEL_TIMER = "model-forward-inference"


class InferenceEngineV2:

    _config: RaggedInferenceEngineConfig
    """
    Configuration of the inference engine.
    """

    _model: DSInferenceModelBase
    """
    Inference model supporting ragged inference.
    """

    _state_manager: DSStateManager
    """
    Persistent state manager for sequences and KV-cache.
    """

    @property
    def free_blocks(self) -> torch.Tensor:
        """
        Number of free KV blocks. This is a tensor of shape [n_kv_cache_groups] where each
        element is the number of free blocks in the corresponding KV cache group.
        """
        return self._state_manager.free_blocks

    @property
    def n_kv_cache_groups(self) -> int:
        """
        Number of KV cache groups.
        """
        return self._state_manager.n_kv_cache_groups

    def model(self) -> DSInferenceModelBase:
        """
        The model implementation.
        """
        return self._model

    def __init__(self, policy: InferenceV2Policy, engine_config: RaggedInferenceEngineConfig) -> None:
        """
        Create the Inference V2 engine.

        Arguments:
            policy (InferenceV2Policy): Policy for the model implementation. This policy object
                will be used to build the model and load the checkpoint associated with it.
            engine_config (RaggedInferenceEngineConfig): Configuration for the inference engine.
        """
        self._config = engine_config
        self._policy = policy
        self._base_mp_group = self._initialize_tp_group()

        # Build model from policy
        inference_logger().info("Building model...")
        self._model = self._policy.build_model(self._config, self._base_mp_group)
        inference_logger().info("Model built.")

        # Create state manager
        self._batch = RaggedBatchWrapper(self._config.state_manager)
        self._state_manager = DSStateManager(self._config.state_manager,
                                             self._model.kv_cache_config(),
                                             base_mp_group=self._base_mp_group)
        self._model.set_state_manager(self._state_manager)

    def _initialize_tp_group(self):
        """
        Implementation of our TP group initialization.
        """
        init_distributed()
        local_rank = int(os.getenv("LOCAL_RANK", 0))
        get_accelerator().set_device(local_rank)

        if local_rank >= self._config.tensor_parallel.tp_size:
            raise RuntimeError("Local rank is greater than TP size, ensure that the TP config is correct.")

        ranks = list(range(self._config.tensor_parallel.tp_size))
        return dist.new_group(ranks=ranks)

    def put(self,
            batch_uids: Iterable[int],
            batch_tokens: Iterable[torch.Tensor],
            do_checks: bool = True) -> torch.Tensor:
        """
        Put a ragged batch onto the inference engine. This will perform one forward and return
        a Tensor of the shape [len(batch_uids), *output_shape]. Logits for the non-final tokens
        are not calculated.

        Arguments:
            batch_uids: Iterable of uids for the batch on the host
            batch_tokens: Iterable of token tensors for the batch on the host
            do_checks: Check schedulability when it is set to True. You can skip this check for better performance when it has already been completed.
        """

        if do_checks:
            token_lens = [len(tokens) for tokens in batch_tokens]
            schedule_check = self.can_schedule(batch_uids, token_lens)
            if schedule_check != SchedulingResult.Success:
                raise SchedulingError(schedule_check)

        self._batch.clear()
        for uid, tokens in zip(batch_uids, batch_tokens):

            host_seq_desc = self._state_manager.get_or_create_sequence(uid)
            self._model.maybe_allocate_kv(host_seq_desc, tokens.numel())
            host_seq_desc.pre_forward(tokens.numel())

            # We can disable checks since we already validated schedulability.
            self._batch.insert_sequence(host_seq_desc, tokens, do_checks=do_checks)

        # Send all metadata to the device
        self._batch.finalize()

        # Prep all data structures for the actual forward (in anticipation of CG in the future)
        # and also to amortize some of the costs in a more straightforward way.
        self._model.prepare_batch(self._batch)

        # Model implementation will pick up in the forward.
        logits = self._model.forward(self._batch)

        # We return one set of logits per sequence in the batch (saves cost on unembedding)
        assert logits.shape[0] == self._batch.current_sequences

        for uid in batch_uids:
            host_seq_desc = self._state_manager.get_sequence(uid)
            host_seq_desc.post_forward()  # Updates sequence metadata.
            self._model.maybe_free_kv(host_seq_desc)

        return logits

    def query(self, uid: int, max_request_tokens: int, max_request_blocks) -> Tuple[int, torch.Tensor]:
        """
        Determine the number of tokens and KV blocks to reserve for a given request. Given a UID
        (this UID may not be recognized by the model yet), this will return the number of tokens
        and blocks to reserve for the request.

        Arguments:
            uid (int): The UID of the sequence (as tracked by the scheduling entity). If
                this is a new sequence (with a UID unknown to the inference engine), then
                an empty placeholder is created to pass to the occupancy logic.
            n_tokens (int): The number of tokens to hypothetically send.

        Returns:
            Tuple[int, Optional[int]]: Tuple of free kv blocks and the number of blocks
                required to schedule the sequence.
        """
        seq_desc = self._state_manager.get_sequence(uid)
        if seq_desc is None:
            if (self._state_manager.n_tracked_sequences == self._config.state_manager.max_tracked_sequences):
                return (0, 0)
            seq_desc = PlaceholderSequenceDescriptor()

        req_tokens, req_blocks = self._model.get_kv_requirements(seq_desc, max_request_tokens, max_request_blocks)

        return (req_tokens, req_blocks)

    def can_schedule(self, uids: Iterable[int], lengths: Iterable[int]) -> SchedulingResult:
        """
        Dry run a batch to determine if it can be scheduled. Placeholder sequences will be
        created for any UIDs that are unknown to the inference engine.

        Arguments:
            uids (Iterable[int]): Iterable of UIDs for the batch
            lengths (Iterable[int]): Iterable of lengths for each sequence of the batch. This lengths
                corresponds to the number of tokens to send in the hypothetical forward; history
                tokens will be determined via UID lookup and future tokens are disregarded.

        Returns:
            bool: True if the batch can be scheduled, False otherwise.
        """

        cur_seqs = self._state_manager.n_tracked_sequences
        free_blocks = self._state_manager.free_blocks
        req_blocks = 0
        batch_len = 0

        if len(uids) > self._config.state_manager.max_ragged_sequence_count:
            # Can only compose a batch from a limited number of sequences
            return SchedulingResult.BatchSequenceLimitExceeded

        for uid, length in zip(uids, lengths):
            seq_desc = self._state_manager.get_sequence(uid)
            if seq_desc is None:
                cur_seqs += 1
                seq_desc = PlaceholderSequenceDescriptor()

            sched_len, sched_blocks = self._model.get_kv_requirements(seq_desc, length, free_blocks)

            if sched_len != length:
                # We ran out of KV cache
                return SchedulingResult.KVCacheLimitExceeded

            batch_len += length
            free_blocks -= sched_blocks

        if cur_seqs > self._config.state_manager.max_tracked_sequences:
            # Would run out of tracking metadata
            return SchedulingResult.EngineSequenceLimitExceeded

        if batch_len > self._config.state_manager.max_ragged_batch_size:
            # Would exceed the maximum batch size
            return SchedulingResult.BatchTokenLimitExceeded

        return SchedulingResult.Success

    def get_remaining_block_capacity(self, uid: int) -> int:
        """
        Get the remaining capacity of the last block already allocated.
        """
        seq_desc = self._state_manager.get_sequence(uid)
        if seq_desc is None:
            return 0
        return self._model.get_remaining_block_capacity(seq_desc)

    def flush(self, uid: int) -> None:
        """
        Remove all state associated with a sequence from the inference engine.

        Arguments:
            uid (int): The UID of the sequence to flush.
        """
        self._state_manager.flush_sequence(uid)

    def serialize(self, save_path: str) -> None:
        """
        Serialize the model to a file.

        Arguments:
            path (str): Path to the file to serialize to.
        """
        param_file_name = make_param_filename(save_path, self._model.tp_rank, self._model.tp_size)
        metadata_file_name = make_metadata_filename(save_path, self._model.tp_rank, self._model.tp_size)

        # Save the flattened parameters

        torch.save(self._model.flattened_params, param_file_name)

        json.dump(self._model.flattened_param_metadata.json(), open(metadata_file_name, "w"))

        if self._model.tp_rank == 0:
            pickle.dump(self._model._config, open(os.path.join(save_path, "ds_model_config.pkl"), "wb"))