"""
Copyright (c) 2023 by FlashInfer 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.
"""

import contextlib
import importlib.util

from .version import __version__ as __version__
from .version import __git_commit__ as __git_commit__
from .version import __git_version__ as __git_version__  # backward compat


from . import jit as jit
from .activation import gelu_and_mul as gelu_and_mul
from .activation import gelu_tanh_and_mul as gelu_tanh_and_mul
from .activation import silu_and_mul as silu_and_mul
from .activation import (
    silu_and_mul_scaled_nvfp4_experts_quantize as silu_and_mul_scaled_nvfp4_experts_quantize,
)
from .attention import BatchAttention as BatchAttention
from .attention import (
    BatchAttentionWithAttentionSinkWrapper as BatchAttentionWithAttentionSinkWrapper,
)
from .autotuner import autotune as autotune
from .cascade import (
    BatchDecodeWithSharedPrefixPagedKVCacheWrapper as BatchDecodeWithSharedPrefixPagedKVCacheWrapper,
)
from .cascade import (
    BatchPrefillWithSharedPrefixPagedKVCacheWrapper as BatchPrefillWithSharedPrefixPagedKVCacheWrapper,
)
from .cascade import (
    MultiLevelCascadeAttentionWrapper as MultiLevelCascadeAttentionWrapper,
)
from .cascade import merge_state as merge_state
from .cascade import merge_state_in_place as merge_state_in_place
from .cascade import merge_states as merge_states
from .decode import (
    BatchDecodeMlaWithPagedKVCacheWrapper as BatchDecodeMlaWithPagedKVCacheWrapper,
)
from .decode import (
    BatchDecodeWithPagedKVCacheWrapper as BatchDecodeWithPagedKVCacheWrapper,
)
from .decode import (
    CUDAGraphBatchDecodeWithPagedKVCacheWrapper as CUDAGraphBatchDecodeWithPagedKVCacheWrapper,
)
from .decode import (
    fast_decode_plan as fast_decode_plan,
)
from .decode import cudnn_batch_decode_with_kv_cache as cudnn_batch_decode_with_kv_cache
from .decode import single_decode_with_kv_cache as single_decode_with_kv_cache
from .quantization.fp4_quantization import (
    block_scale_interleave,
    nvfp4_block_scale_interleave,
    e2m1_and_ufp8sf_scale_to_float,
    fp4_quantize,
    mxfp4_dequantize_host,
    mxfp4_dequantize,
    mxfp4_quantize,
    nvfp4_quantize,
    nvfp4_quantize_paged_kv_cache,
    nvfp4_batched_quantize,
    shuffle_matrix_a,
    shuffle_matrix_sf_a,
    scaled_fp4_grouped_quantize,
    silu_and_mul_nvfp4_quantize,
    get_fp4_quantization_module,
    nvfp4_kv_dequantize,
    nvfp4_kv_dequantize_paged,
    nvfp4_kv_quantize,
)
from .quantization.fp8_quantization import (
    mxfp8_dequantize_host,
    mxfp8_grouped_quantize,
    mxfp8_quantize,
)
from .fused_moe import (
    cutlass_fused_moe,
    reorder_rows_for_gated_act_gemm,
    trtllm_bf16_moe,
    trtllm_bf16_routed_moe,
    trtllm_fp4_block_scale_moe,
    trtllm_fp4_block_scale_routed_moe,
    trtllm_fp8_block_scale_moe,
    trtllm_fp8_block_scale_routed_moe,
    trtllm_fp8_per_tensor_scale_moe,
    trtllm_fp8_per_tensor_scale_routed_moe,
)

# CuteDSL high-level APIs (conditionally if cute_dsl available)
with contextlib.suppress(ImportError):
    from .fused_moe import (
        cute_dsl_fused_moe_nvfp4 as cute_dsl_fused_moe_nvfp4,
        CuteDslMoEWrapper as CuteDslMoEWrapper,
        b12x_fused_moe as b12x_fused_moe,
        B12xMoEWrapper as B12xMoEWrapper,
    )
    from .gdn_prefill import chunk_gated_delta_rule as chunk_gated_delta_rule
from .gemm import SegmentGEMMWrapper as SegmentGEMMWrapper
from .gemm import bmm_bf16 as bmm_bf16
from .gemm import bmm_fp8 as bmm_fp8
from .gemm import bmm_mxfp8 as bmm_mxfp8
from .gemm import mm_bf16 as mm_bf16
from .gemm import mm_fp4 as mm_fp4
from .gemm import mm_nvfp4_svdquant as mm_nvfp4_svdquant
from .gemm import nvfp4_quantize_smooth as nvfp4_quantize_smooth
from .gemm import svdquant_linear as svdquant_linear
from .gemm import mm_bf16_fp4 as mm_bf16_fp4
from .gemm import prepare_bf16_fp4_weights as prepare_bf16_fp4_weights
from .gemm import mm_fp8 as mm_fp8
from .gemm import mm_mxfp8 as mm_mxfp8
from .gemm import tgv_gemm_sm100 as tgv_gemm_sm100
from .grouped_mm import grouped_mm_bf16 as grouped_mm_bf16
from .grouped_mm import grouped_mm_fp8 as grouped_mm_fp8
from .grouped_mm import grouped_mm_mxfp8 as grouped_mm_mxfp8
from .grouped_mm import grouped_mm_fp4 as grouped_mm_fp4
from .kda_prefill import (
    RecurrentKDAPrefillWorkspace as RecurrentKDAPrefillWorkspace,
)
from .kda import recurrent_kda as recurrent_kda
from .kda_decode import fused_kda_decode as fused_kda_decode
from .mla import BatchMLAPagedAttentionWrapper as BatchMLAPagedAttentionWrapper
from . import mhc as mhc
from . import msa_ops as msa_ops
from .norm import fused_add_rmsnorm as fused_add_rmsnorm
from .norm import fused_add_rmsnorm_quant as fused_add_rmsnorm_quant
from .norm import layernorm as layernorm
from .norm import layernorm_quant as layernorm_quant
from .norm import gemma_fused_add_rmsnorm as gemma_fused_add_rmsnorm
from .norm import gemma_rmsnorm as gemma_rmsnorm
from .norm import rmsnorm as rmsnorm
from .norm import rmsnorm_quant as rmsnorm_quant
from .norm import fused_rmsnorm_silu as fused_rmsnorm_silu
from .norm import fused_qk_rmsnorm_rope as fused_qk_rmsnorm_rope
from . import nvfp4_attention_sm120 as nvfp4_attention_sm120
from .nvfp4_attention_sm120 import (
    nvfp4_attention_sm120_fwd as nvfp4_attention_sm120_fwd,
)
from .nvfp4_attention_sm120 import (
    nvfp4_attention_sm120_quantize_qkv as nvfp4_attention_sm120_quantize_qkv,
)
from .norm import (
    fused_dit_residual_layernorm_scale_shift as fused_dit_residual_layernorm_scale_shift,
)
from .norm import (
    fused_dit_gate_residual_layernorm_scale_shift as fused_dit_gate_residual_layernorm_scale_shift,
)
from .norm import (
    fused_dit_gate_residual_layernorm_gamma_beta as fused_dit_gate_residual_layernorm_gamma_beta,
)

try:
    from .norm import rmsnorm_fp4quant as rmsnorm_fp4quant
    from .norm import add_rmsnorm_fp4quant as add_rmsnorm_fp4quant
except (ImportError, AttributeError):
    pass  # nvidia-cutlass-dsl not installed
from .page import append_paged_kv_cache as append_paged_kv_cache
from .page import append_paged_mla_kv_cache as append_paged_mla_kv_cache
from .page import get_batch_indices_positions as get_batch_indices_positions
from .page import get_seq_lens as get_seq_lens
from .page import (
    nvfp4_quantize_append_paged_kv_cache as nvfp4_quantize_append_paged_kv_cache,
)
from .page import (
    nvfp4_quantize_append_paged_kv_cache_with_slot_mapping as nvfp4_quantize_append_paged_kv_cache_with_slot_mapping,
)
from .pod import PODWithPagedKVCacheWrapper as PODWithPagedKVCacheWrapper
from .pod import BatchPODWithPagedKVCacheWrapper as BatchPODWithPagedKVCacheWrapper
from .prefill import (
    BatchPrefillWithPagedKVCacheWrapper as BatchPrefillWithPagedKVCacheWrapper,
)
from .prefill import (
    BatchPrefillWithRaggedKVCacheWrapper as BatchPrefillWithRaggedKVCacheWrapper,
)
from .prefill import single_prefill_with_kv_cache as single_prefill_with_kv_cache
from .prefill import (
    single_prefill_with_kv_cache_return_lse as single_prefill_with_kv_cache_return_lse,
)
from .prefill import trtllm_fmha_v2_prefill as trtllm_fmha_v2_prefill
from .prefill import (
    trtllm_sage_attention_quantize as trtllm_sage_attention_quantize,
)
from .quantization import packbits as packbits
from .quantization import segment_packbits as segment_packbits
from .rope import apply_llama31_rope as apply_llama31_rope
from .rope import apply_llama31_rope_inplace as apply_llama31_rope_inplace
from .rope import apply_llama31_rope_pos_ids as apply_llama31_rope_pos_ids
from .rope import (
    apply_llama31_rope_pos_ids_inplace as apply_llama31_rope_pos_ids_inplace,
)
from .rope import apply_rope as apply_rope
from .rope import apply_rope_inplace as apply_rope_inplace
from .rope import apply_rope_pos_ids as apply_rope_pos_ids
from .rope import apply_rope_pos_ids_inplace as apply_rope_pos_ids_inplace
from .rope import apply_rope_with_cos_sin_cache as apply_rope_with_cos_sin_cache
from .rope import (
    apply_rope_with_cos_sin_cache_inplace as apply_rope_with_cos_sin_cache_inplace,
)
from .sampling import chain_speculative_sampling as chain_speculative_sampling
from .sampling import min_p_sampling_from_probs as min_p_sampling_from_probs
from .sampling import sampling_from_logits as sampling_from_logits
from .sampling import sampling_from_probs as sampling_from_probs
from .sampling import softmax as softmax
from .sampling import top_k_mask_logits as top_k_mask_logits
from .sampling import top_k_renorm_probs as top_k_renorm_probs
from .sampling import top_k_sampling_from_probs as top_k_sampling_from_probs
from .sampling import (
    top_k_top_p_sampling_from_logits as top_k_top_p_sampling_from_logits,
)
from .sampling import top_k_top_p_sampling_from_probs as top_k_top_p_sampling_from_probs
from .sampling import top_p_renorm_probs as top_p_renorm_probs
from .sampling import top_p_sampling_from_probs as top_p_sampling_from_probs
from .tllm_enums import (
    SfLayout,
    ActivationType,
    RoutingMethodType,
    is_gated_activation as is_gated_activation,
)
from . import topk as topk
from .topk import top_k as top_k
from .topk import top_k_page_table_transform as top_k_page_table_transform
from .topk import top_k_ragged_transform as top_k_ragged_transform
from .topk import TopKTieBreak as TopKTieBreak
from .topk_varlen.topk_varlen import top_k_varlen as top_k_varlen
from .sparse import BlockSparseAttentionWrapper as BlockSparseAttentionWrapper
from .sparse import (
    VariableBlockSparseAttentionWrapper as VariableBlockSparseAttentionWrapper,
)
from .trtllm_low_latency_gemm import (
    prepare_low_latency_gemm_weights as prepare_low_latency_gemm_weights,
)
from .utils import next_positive_power_of_2 as next_positive_power_of_2
from .xqa import xqa as xqa
from .xqa import xqa_mla as xqa_mla
from . import mamba as mamba
from .fi_trace import fi_trace as fi_trace

# ---------------------------------------------------------------------------
# Trace Apply (opt-in): zero-code kernel substitution from the FlashInfer Trace.
# Activated only when FLASHINFER_TRACE_APPLY=1. This is the single edit to
# existing flashinfer code required by the Trace Apply design — the package
# itself lives in flashinfer/trace_apply/. Failures here must never break a
# normal import, so the install is best-effort.
# ---------------------------------------------------------------------------
import os as _os

if _os.environ.get("FLASHINFER_TRACE_APPLY", "0") not in ("0", "", "false", "False"):
    try:
        from . import trace_apply as trace_apply

        trace_apply._enable_apply_from_env()
    except Exception as _trace_apply_err:  # noqa: BLE001
        import logging as _logging

        _logging.getLogger("flashinfer.trace_apply").warning(
            "FLASHINFER_TRACE_APPLY is set but enabling Trace Apply failed: %s "
            "(continuing without Trace Apply).",
            _trace_apply_err,
        )


# ---------------------------------------------------------------------------
# Import-time version log: emit one line when FLASHINFER_LOGLEVEL >= 1 so
# that crash logs contain the exact commit without any manual archaeology.
# Respects FLASHINFER_LOGDEST (stdout / stderr / filepath) the same way
# api_logging.py does; defaults to stdout.
# ---------------------------------------------------------------------------
def _log_import_version() -> None:
    # Wrapped in a private function so no temp variables leak into the
    # flashinfer module namespace.  Two-level protection:
    #   inner try  – safely resolve the commit hash; falls back to "unknown"
    #                if __git_commit__ is missing or malformed so the log
    #                line is still emitted rather than silently suppressed.
    #   outer try  – absorbs every other failure (non-integer LOGLEVEL env
    #                var, closed/None stdout or stderr, unwritable log file)
    #                so a logging misconfiguration can never block the import.
    try:
        if int(_os.environ.get("FLASHINFER_LOGLEVEL", "0")) < 1:
            return
        try:
            _short = __git_commit__[:8] if __git_commit__ != "unknown" else "unknown"
        except Exception:
            _short = "unknown"
        _line = f"FlashInfer {__version__} (commit {_short})\n"
        _dest = _os.environ.get("FLASHINFER_LOGDEST", "stdout").replace(
            "%i", str(_os.getpid())
        )
        if _dest == "stderr":
            import sys as _sys

            _sys.stderr.write(_line)
            _sys.stderr.flush()
        elif _dest not in ("stdout", ""):
            with open(_dest, "a") as _f:
                _f.write(_line)
        else:
            print(_line, end="", flush=True)
    except Exception:
        pass  # never let import-time logging crash the import


_log_import_version()
del _log_import_version
