# Copyright 2026 The HuggingFace Inc. team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Resolves a `ContinuousBatchingConfig` into a fully-specified config ready for cache and runner creation. Each
helper mutates the config in place; `resolve_continuous_batching_config` orchestrates them in the required order."""

from copy import deepcopy
from math import ceil

import torch

from ...configuration_utils import PretrainedConfig
from ...generation.configuration_utils import CompileConfig, ContinuousBatchingConfig
from ...modeling_flash_attention_utils import lazy_import_paged_flash_attention
from ...utils import is_torch_xpu_available
from ...utils.generic import is_flash_attention_requested
from .requests import logger
from .utils import WorkloadHints


FALLBACK_DEFAULTS = {
    "max_requests_per_batch": 1024,
    "max_blocks_per_request": 32,
    "q_padding_interval_size": 64,
    "kv_padding_interval_size": 64 * 256,  # 64 blocks of 256 tokens ie. 16384 tokens
    "max_cached_graphs": 32,
}


def resolve_continuous_batching_config(
    config: PretrainedConfig,
    cb_config: ContinuousBatchingConfig,
    workload_hints: WorkloadHints | None,
    has_logit_processors: bool,
) -> ContinuousBatchingConfig:
    """Returns a deep-copied and fully-resolved `ContinuousBatchingConfig`. The original `cb_config` is not mutated."""
    cb_config = deepcopy(cb_config)

    # Look at whether the user explicitly asked for the decode fast path before we assign a default value
    user_requested_decode_path = cb_config.max_blocks_per_request is not None
    # Same for cuda graphs, if the user signals they want CUDA graphs via any padding/cached-graph parameter
    cuda_graph_requested = any(
        [cb_config.q_padding_interval_size, cb_config.kv_padding_interval_size, cb_config.max_cached_graphs]
    )

    # Resolve missing attributes for which we have hints. Must happen before no-hints resolve.
    resolve_using_hints(cb_config, workload_hints)

    # Resolve remaining missing attributes. Must happen before decode fast path is checked.
    resolve_without_hints(cb_config)

    # Check if the decode fast path is available. Must happen before the compile config.
    ensure_decode_fast_path_is_available(config, cb_config, user_requested_decode_path)

    # Decide if compile should be used. Must happen before CUDA graphs are decided.
    resolve_compile_configs(
        cb_config=cb_config,
        fallback_compile_config=getattr(config, "compile_config", None),
        is_flash_attn=is_flash_attention_requested(config),
        decode_fast_path_available=cb_config.max_blocks_per_request > 0,
    )

    # Decide if CUDA graphs should be used. Should happen after compile configs are decided.
    is_attn_mask_needed = not is_flash_attention_requested(config)
    decide_use_cuda_graphs(
        cb_config=cb_config, is_attn_mask_needed=is_attn_mask_needed, cuda_graph_requested=cuda_graph_requested
    )

    # Decide if asynchronous batching should be used. Should happen after CUDA graphs are decided.
    decide_use_async_batching(cb_config=cb_config, is_attn_mask_needed=is_attn_mask_needed)

    # Resolve the max memory percent. This can happen anytime before cache creation.
    resolve_max_memory_percent(cb_config=cb_config, has_logit_processors=has_logit_processors)
    return cb_config


def resolve_using_hints(cb_config: ContinuousBatchingConfig, workload_hints: WorkloadHints | None) -> None:
    """Fills some attributes from the workload hints, when the user did not set it explicitly: `max_blocks_per_request`
    and `max_requests_per_batch`."""
    # The max number of blocks per request is an even number large enough to hold the max request length
    if cb_config.max_blocks_per_request is None and workload_hints is not None:
        max_sequence_length = workload_hints.max_prompt_length + workload_hints.max_generated_length
        if max_sequence_length > 0:
            blocks_per_request = int(ceil(max_sequence_length / cb_config.block_size)) + 1
            cb_config.max_blocks_per_request = blocks_per_request + (blocks_per_request % 2)
    # The maximum number of requests per batch is the minimum of the workload hints and the fallback default
    if cb_config.max_requests_per_batch is None and workload_hints is not None:
        if workload_hints.num_requests > 0:  # guard against bad hints
            max_requests_per_batch = min(workload_hints.num_requests, FALLBACK_DEFAULTS["max_requests_per_batch"])
        else:
            max_requests_per_batch = FALLBACK_DEFAULTS["max_requests_per_batch"]
        cb_config.max_requests_per_batch = max_requests_per_batch


def resolve_without_hints(cb_config: ContinuousBatchingConfig) -> None:
    """Fills any remaining unset/sentinel attribute with a fallback default."""
    if cb_config.max_requests_per_batch is None:
        cb_config.max_requests_per_batch = FALLBACK_DEFAULTS["max_requests_per_batch"]
    if cb_config.max_blocks_per_request is None:
        cb_config.max_blocks_per_request = FALLBACK_DEFAULTS["max_blocks_per_request"]
    if cb_config.q_padding_interval_size == 0:
        cb_config.q_padding_interval_size = FALLBACK_DEFAULTS["q_padding_interval_size"]
    if cb_config.kv_padding_interval_size == 0:
        cb_config.kv_padding_interval_size = FALLBACK_DEFAULTS["kv_padding_interval_size"]
    if cb_config.max_cached_graphs == 0:
        cb_config.max_cached_graphs = FALLBACK_DEFAULTS["max_cached_graphs"]


def ensure_decode_fast_path_is_available(
    config: PretrainedConfig, cb_config: ContinuousBatchingConfig, user_requested: bool
) -> None:
    """Ensures the decode fast path is available. If it is not, set the max blocks per request to 0. If it is
    available, and no user-provided max blocks per request, set it to the fallback default."""
    # Then, if the decode fast path is not turned off, check if it is available
    if cb_config.max_blocks_per_request != 0:
        # NOTE: For CUDA, block table should be available with FA2 and FA3, but there seems to be an issue with FA2 atm
        cuda_available = torch.cuda.is_available()
        fa_cuda = is_flash_attention_requested(config, version=3) and cuda_available
        # XPU support is given through its kernel variation `kernels-community/flash-attn2`
        xpu_available = is_torch_xpu_available()
        fa_xpu = is_flash_attention_requested(config, version=2) and xpu_available
        if fa_cuda or fa_xpu:  # Block table is only supported on these
            flash_attn_with_kvcache = lazy_import_paged_flash_attention(config._attn_implementation)[1]
            # Throw a warning only if the decode fast path was requested by the user
            if flash_attn_with_kvcache is None:
                if user_requested:
                    logger.warning(
                        f"Although {cb_config.max_blocks_per_request = }, the decode fast path is not available "
                        f"because `flash_attn_with_kvcache` is not available for {config._attn_implementation = }."
                    )
                cb_config.max_blocks_per_request = 0
        # Specific warning for unsupported attention implementation/device combinations
        else:
            if user_requested:
                logger.warning(
                    f"Although {cb_config.max_blocks_per_request = }, the decode fast path is not available "
                    "because the attention implementation and device combination is not supported. Supported "
                    "combinations are Flash Attention 3 on CUDA, or Flash Attention 2 on XPU through "
                    "`kernels-community/flash-attn2`. "
                    f"Got {config._attn_implementation = }, {cuda_available = }, {xpu_available = }."
                )
            cb_config.max_blocks_per_request = 0


def resolve_compile_configs(
    cb_config: ContinuousBatchingConfig,
    fallback_compile_config: CompileConfig | None,
    is_flash_attn: bool,
    decode_fast_path_available: bool,
) -> None:
    """Resolve if the compile configs for varlen and decode paths, modifying these attributes in place if needed.
    Default config use full compile over regional compile, because the throughput is significantly higher (~15%)"""
    default_mode = "max-autotune-no-cudagraphs" if cb_config.default_compile_level >= 2 else "default"
    default_dynamic = cb_config.default_compile_level <= 2
    # For each config, priority is: explicit config, default config, fallback config, None
    if cb_config.varlen_compile_config is None:
        if cb_config.default_compile_level > 0:
            # TODO: now that max_seqlen_k is bucketted, is that still True?
            # We don't use compile with flash varlen, because max_seqlen_k is volatile and introduces recompilations
            if is_flash_attn:
                varlen_config = None
            else:
                varlen_config = CompileConfig(mode=default_mode, fullgraph=True, dynamic=default_dynamic)
        elif fallback_compile_config is not None:
            varlen_config = fallback_compile_config
        else:
            varlen_config = None
    else:
        varlen_config = cb_config.varlen_compile_config

    if cb_config.decode_compile_config is None:
        if cb_config.default_compile_level > 0:
            # Paged attention is wrapped in @torch.compiler.disable so we can't use fullgraph
            decode_config = CompileConfig(mode=default_mode, fullgraph=False, dynamic=default_dynamic)
        elif fallback_compile_config is not None:
            decode_config = fallback_compile_config
        else:
            decode_config = None
    else:
        decode_config = cb_config.decode_compile_config

    # For decode, we throw a warning if the fast decode path is not available and a compile config was found
    if not decode_fast_path_available and cb_config.decode_compile_config is not None:
        decode_config = None
        logger.warning("A decode_compile_config was set but fast decode path is not available. Ignoring it.")

    # Log what will be compiled
    if varlen_config is not None:
        logger.info(f"Varlen path will be compiled with {varlen_config.to_dict()}")
    if decode_config is not None:
        logger.info(f"Decode path will be compiled with {decode_config.to_dict()}")
    # Modify in place
    cb_config.varlen_compile_config = varlen_config
    cb_config.decode_compile_config = decode_config


def decide_use_cuda_graphs(
    cb_config: ContinuousBatchingConfig, is_attn_mask_needed: bool, cuda_graph_requested: bool
) -> None:
    """Decides whether or not to use cuda graphs for continuous batching. If the user specified this in the config
    or if they specified a parameter related to cuda graphs, they are turned on. Otherwise, we use a heuristic
    based on the attention implementation: we turn on cuda graphs if and only if no attention mask is needed.

    This function modifies the `use_cuda_graph` attribute of the config in place, to a tuple of booleans.
    """
    # If cuda is not available, we cannot use cuda graphs
    if not torch.cuda.is_available():
        intended_use_cuda_graph = any(cb_config.cuda_graph_booleans)
        if intended_use_cuda_graph:  # throw a warning only if the user intended to use cuda graphs
            logger.warning(
                f"{cb_config.use_cuda_graph = } but {torch.cuda.is_available() = }: turning off cuda graphs"
            )
        cb_config.use_cuda_graph = (False, False)

    # Else if use_cuda_graph is specified, we follow the user's choice and make sure it is a tuple of booleans
    elif cb_config.use_cuda_graph is not None:
        if isinstance(cb_config.use_cuda_graph, bool):
            cb_config.use_cuda_graph = (cb_config.use_cuda_graph, cb_config.use_cuda_graph)

    # Else if the user specified a parameter related to cuda graphs, we activate cuda graphs
    elif cuda_graph_requested:
        cb_config.use_cuda_graph = (True, True)

    # Otherwise we have a default heuristic based on the attention implementation:
    # attention implementations where an attention mask is needed suffer a lot more from the padding associated
    # with cuda graphs, so default is to turn cuda graphs off for those implementations
    else:
        use_cuda_graph = []
        for compile_config in [cb_config.varlen_compile_config, cb_config.decode_compile_config]:
            # No compile config means we decide on attention
            if compile_config is None:
                use_cuda_graph.append(not is_attn_mask_needed)
                continue
            # Otherwise we disable cuda graphs if the compile config uses them
            options = torch._inductor.list_mode_options().get(compile_config.mode, compile_config.options)
            compile_uses_cudagraphs = options.get("triton.cudagraphs", False)
            if compile_uses_cudagraphs:
                logger.warning(
                    f"Compile config {compile_config.mode = } uses cudagraphs, which usually does not work well with "
                    "continuous batching. We recommend using mode 'default' or 'max-autotune-no-cudagraphs' instead."
                )
            use_cuda_graph.append(not compile_uses_cudagraphs and not is_attn_mask_needed)
        cb_config.use_cuda_graph = tuple(use_cuda_graph)

    logger.info(f"Using cuda graphs for (varlen, decode) paths: {cb_config.use_cuda_graph}")


def decide_use_async_batching(cb_config: ContinuousBatchingConfig, is_attn_mask_needed: bool) -> None:
    """Returns whether or not to use asynchronous batching for continuous batching. If the user specified this in
    the config, we follow their choice. Otherwise, we turn on asynchronous batching if and only if CUDA graphs are
    turned on and no attention mask is needed.

    This function modifies the `use_async_batching` attribute of the config in place.
    """
    # If the user specifies to use async or not, no need to decide ourselves
    if cb_config.use_async_batching is None:
        use_cuda_graphs = any(cb_config.cuda_graph_booleans)
        cb_config.use_async_batching = use_cuda_graphs and not is_attn_mask_needed
        logger.info(
            f"No behavior specified for use_async_batching, choosing {cb_config.use_async_batching = } because "
            f"{use_cuda_graphs = } and {is_attn_mask_needed = }. If you want to save memory, you can "
            "disable asynchronous batching but it will degrade performance."
        )


def resolve_max_memory_percent(cb_config: ContinuousBatchingConfig, has_logit_processors: bool) -> None:
    if cb_config.max_memory_percent is None:
        cb_config.max_memory_percent = 0.8 if has_logit_processors else 0.9


def update_cb_config_after_cache_creation(
    cb_config: ContinuousBatchingConfig,
    num_blocks: int,
    max_batch_tokens: int,
    use_prefix_sharing: bool,
) -> None:
    """Updates the continuous batching config with the concrete values inferred during the creation of the cache."""
    # Memoize concrete values
    cb_config.num_blocks = num_blocks
    cb_config.max_batch_tokens = max_batch_tokens
    # Cap the number of max requests per batch to the max tokens per batch
    cb_config.max_requests_per_batch = min(cb_config.max_requests_per_batch, max_batch_tokens)
    # And if there is no prefix sharing, we can cap the number of request per batch (1 request = 1 block at least)
    if not use_prefix_sharing:
        cb_config.max_requests_per_batch = min(cb_config.max_requests_per_batch, num_blocks)
    # TODO: should we align the max number of request per batch to a multiple of 32 ?
