"""Token-level middle truncation (ported from your /data1/sora/evalscope/bash/run.py).

Keeps head+tail halves of the token stream -- the industry-standard
middle-truncation for long-context benchmarks (longbench_v2 / mrcr).
The evalside run.py uses the same algorithm, guaranteeing comparable inputs.

Usage in run_eval:  max_input_tokens=131072  (0/off = no truncation)
Requires a tokenizer (transformers) at tokenizer_path or auto from the model.
"""

import os
from functools import lru_cache
from typing import Optional

DEFAULT_TRUNCATION_TOKENS = 32768 * 4  # 131072, mirrors evalside run.py


@lru_cache(maxsize=4)
def _get_tokenizer(tokenizer_path: str):
    if not tokenizer_path or not os.path.exists(tokenizer_path):
        raise FileNotFoundError(
            f'tokenizer not found at {tokenizer_path!r} -- token-level truncation '
            'needs a local tokenizer dir (e.g. /data1/models/DeepSeek-V4-Flash-INT8)')
    from transformers import AutoTokenizer

    return AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)


def truncate_middle_tokens(text: str, max_tokens: int, tokenizer_path: str) -> str:
    """Keep head+tail halves of the token stream; decode back to text."""
    if max_tokens <= 0 or not text:
        return text
    tok = _get_tokenizer(tokenizer_path)
    ids = tok.encode(text, add_special_tokens=False)
    if len(ids) <= max_tokens:
        return text
    keep_head = max_tokens // 2
    keep_tail = max_tokens - keep_head
    return tok.decode(ids[:keep_head] + ids[-keep_tail:], skip_special_tokens=True)


def truncate_messages_middle(messages: list, max_tokens: int, tokenizer_path: str,
                             desired_index: int = 0, window: int = 2) -> list:
    """MRCR-style: when a message list exceeds the budget, keep first/last
    messages plus a window around the desired (needle) message, dropping
    middle spans; middle of each KEPT long message is token-truncated."""
    if max_tokens <= 0 or not messages:
        return messages
    tok = _get_tokenizer(tokenizer_path)
    total = sum(len(tok.encode(m.get('content', '') if isinstance(m, dict) else str(m),
                                add_special_tokens=False)) for m in messages)
    if total <= max_tokens:
        return messages
    n = len(messages)
    keep = set(range(min(2, n))) | set(range(max(0, n - 2), n))
    di = desired_index if isinstance(desired_index, int) and 0 <= desired_index < n else 0
    keep |= set(range(max(0, di - window), min(n, di + window + 1)))
    return [messages[i] for i in sorted(keep)]


def default_tokenizer_path() -> Optional[str]:
    """Candidate local tokenizer for truncation (mirrors evalside default)."""
    env = os.environ.get('EVALHARNESS_TOKENIZER')
    if env and os.path.exists(env):
        return env
    for cand in ('/data1/models/DeepSeek-V4-Flash-INT8',):
        if os.path.exists(cand):
            return cand
    return None