# FILE: nanovllm/utils/context.py
from dataclasses import dataclass
from typing import Optional
import mindtorch


@dataclass
class Context:
    is_prefill: bool = False
    cu_seqlens_q: Optional[mindtorch.Tensor] = None
    cu_seqlens_k: Optional[mindtorch.Tensor] = None
    max_seqlen_q: int = 0
    max_seqlen_k: int = 0
    slot_mapping: Optional[mindtorch.Tensor] = None
    context_lens: Optional[mindtorch.Tensor] = None
    block_tables: Optional[mindtorch.Tensor] = None
    block_size: int = 0

_CONTEXT = Context()

def get_context():
    return _CONTEXT

def set_context(is_prefill, cu_seqlens_q=None, cu_seqlens_k=None, max_seqlen_q=0, max_seqlen_k=0, slot_mapping=None, context_lens=None, block_tables=None, block_size=0):
    global _CONTEXT
    _CONTEXT = Context(is_prefill, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, context_lens, block_tables, block_size)

def reset_context():
    global _CONTEXT
    _CONTEXT = Context()
