# #region estimate_token_budget [C:3] [TYPE Module] [SEMANTICS translate, token, budget, estimation, llm]
# @BRIEF Calculate safe batch_size and max_tokens for LLM translation calls based on actual content length and model context window limits.
# @LAYER: Domain
# @RELATION DEPENDS_ON -> [TranslationPreview:Module]
# @RELATION DEPENDS_ON -> [TranslationExecutor:Module]
# @RATIONALE: Prevents LLM truncation (finish_reason=length) by sizing batches within context limits.
#             DeepSeek v4 Flash supports up to 64K context window; output is limited by max_tokens.
# @REJECTED: External tokenizer library — would introduce heavy dependency for estimation only.
#            Fixed batch_size of 50 — causes truncation on long-content rows.

# #region DEFAULT_CONTEXT_WINDOW [TYPE Constant]
# @BRIEF Default context window for DeepSeek v4 Flash: up to 64K tokens.
DEFAULT_CONTEXT_WINDOW = 64000
# #endregion DEFAULT_CONTEXT_WINDOW

# #region DEFAULT_MAX_OUTPUT_TOKENS [TYPE Constant]
# @BRIEF Default max_tokens setting for LLM output (8192 tokens).
DEFAULT_MAX_OUTPUT_TOKENS = 8192
# #endregion DEFAULT_MAX_OUTPUT_TOKENS

# #region REASONING_OVERHEAD [TYPE Constant]
# @BRIEF CoT reasoning overhead tokens for DeepSeek models (~2000 tokens for chain-of-thought).
REASONING_OVERHEAD = 2000
# #endregion REASONING_OVERHEAD

# #region OUTPUT_PER_ROW_PER_LANG [TYPE Constant]
# @BRIEF Estimated output tokens per row per language in JSON response format.
OUTPUT_PER_ROW_PER_LANG = 60
# #endregion OUTPUT_PER_ROW_PER_LANG

# #region PROMPT_BASE_TOKENS [TYPE Constant]
# @BRIEF Base tokens for system prompt + instructions + JSON format specification.
PROMPT_BASE_TOKENS = 300
# #endregion PROMPT_BASE_TOKENS

# #region DICT_TOKENS_PER_ENTRY [TYPE Constant]
# @BRIEF Estimated tokens per dictionary entry in the glossary section.
DICT_TOKENS_PER_ENTRY = 20
# #endregion DICT_TOKENS_PER_ENTRY

# #region DICT_TOKENS_MAX [TYPE Constant]
# @BRIEF Cap for dictionary tokens in estimation to avoid overestimating.
DICT_TOKENS_MAX = 5000
# #endregion DICT_TOKENS_MAX

# #region CHARS_PER_TOKEN_MIXED [TYPE Constant]
# @BRIEF Characters per token for mixed Russian/English text (empirical ~2.2 for mixed, ~2 for Russian, ~4 for English).
CHARS_PER_TOKEN_MIXED = 2.2
# #endregion CHARS_PER_TOKEN_MIXED

# #region MIN_MAX_TOKENS [TYPE Constant]
# @BRIEF Minimum max_tokens to allow (4096 ensures reasonable output even for small batches).
MIN_MAX_TOKENS = 4096
# #endregion MIN_MAX_TOKENS

# #region MAX_OUTPUT_HEADROOM [TYPE Constant]
# @BRIEF Extra headroom added to max_output_needed beyond the estimate (1000 tokens buffer).
MAX_OUTPUT_HEADROOM = 1000
# #endregion MAX_OUTPUT_HEADROOM


# region _count_rows_that_fit [TYPE Function]
# @BRIEF Count how many rows fit within the available input budget.
# @PRE: input_per_row is non-empty; available_budget > 0.
# @POST: Returns (safe_count, total_input_tokens).
def _count_rows_that_fit(
    input_per_row: list[int],
    available_budget: int,
) -> tuple[int, int]:
    """Count consecutive rows that fit within available_budget.

    Returns:
        (safe_count, total_input_tokens): Number of rows and their total tokens.
    """
    running_total = 0
    safe_size = 0
    for tokens in input_per_row:
        if running_total + tokens + REASONING_OVERHEAD < available_budget:
            running_total += tokens
            safe_size += 1
        else:
            break
    safe_size = max(safe_size, 1)
    return safe_size, running_total
# endregion _count_rows_that_fit


# region estimate_token_budget [C:3] [TYPE Function]
# @BRIEF Estimate token budget for a batch of source rows and return safe batch parameters.
# @PRE: source_rows is a list of dicts (can be empty). target_languages is a non-empty list.
# @POST: Returns dict with batch_size_adjusted, estimated_input_tokens, estimated_output_tokens, max_output_needed, warning.
# @RATIONALE: Uses character-count heuristics (chars/2.2 for mixed text) since exact tokenization
#             depends on the LLM model. Estimates are intentionally conservative to prevent truncation.
# @REJECTED: Using tiktoken or similar tokenizer — would introduce a heavy dependency and still
#            not match DeepSeek's tokenizer exactly.
def estimate_token_budget(
    source_rows: list[dict],
    target_languages: list[str],
    source_column: str = "source_text",
    context_columns: list[str] | None = None,
    dictionary_entries: list | None = None,
    batch_size: int | None = None,
    context_window: int = DEFAULT_CONTEXT_WINDOW,
    max_output_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
) -> dict:
    """Estimate token budget and return safe batch parameters.

    Args:
        source_rows: List of row dicts with source text.
        target_languages: List of target language codes.
        source_column: Key for the source text in each row dict.
        context_columns: Optional list of keys for context columns.
        dictionary_entries: Optional list of dictionary entries for glossary.
        batch_size: Desired batch size. If None, auto-calculate max safe size.
        context_window: Model context window (default 64000 for DeepSeek v4 Flash).
        max_output_tokens: Hard max output tokens limit (default 8192).

    Returns:
        dict with:
            batch_size_adjusted: Safe batch size (may be less than requested).
            estimated_input_tokens: Total estimated input tokens for the batch.
            estimated_output_tokens: Total estimated output tokens.
            max_output_needed: Recommended max_tokens for this batch.
            warning: str | None if batch was reduced.
    """
    if not target_languages:
        target_languages = ["en"]

    num_languages = len(target_languages)

    # 1. Estimate tokens per row
    input_per_row = []
    limit = batch_size if batch_size else len(source_rows)
    for i, row in enumerate(source_rows):
        if i >= limit:
            break
        text = str(row.get(source_column, "") or "")
        # Use ~2.2 chars per token for mixed Russian/English text
        estimated_tokens = max(1, int(len(text) / CHARS_PER_TOKEN_MIXED))
        # Add context columns
        if context_columns:
            for col in context_columns:
                val = str(row.get(col, "") or "")
                estimated_tokens += max(1, int(len(val) / CHARS_PER_TOKEN_MIXED))
        input_per_row.append(estimated_tokens)

    # 2. Calculate dictionary tokens
    dict_tokens = 0
    if dictionary_entries:
        dict_tokens = min(
            len(dictionary_entries) * DICT_TOKENS_PER_ENTRY,
            DICT_TOKENS_MAX,
        )

    prompt_tokens = PROMPT_BASE_TOKENS + dict_tokens
    available_input_budget = context_window - max_output_tokens

    # 3. Calculate safe batch size using shared helper
    safe_size, input_total = _count_rows_that_fit(input_per_row, available_input_budget)
    estimated_input = prompt_tokens + input_total

    # If batch_size was specified and we reduced it, recalculate
    if batch_size and safe_size < batch_size:
        # Ensure at minimum one row even for huge content
        safe_size = max(safe_size, 1)
        _, truncated_total = _count_rows_that_fit(
            input_per_row[:batch_size], available_input_budget,
        )
        estimated_input = prompt_tokens + truncated_total

    # 4. Estimate output tokens
    estimated_output = (
        safe_size * num_languages * OUTPUT_PER_ROW_PER_LANG + REASONING_OVERHEAD
    )

    # 5. Calculate recommended max_tokens
    max_output_needed = min(
        estimated_output + MAX_OUTPUT_HEADROOM,
        context_window - estimated_input,
    )
    max_output_needed = max(max_output_needed, MIN_MAX_TOKENS)
    max_output_needed = min(max_output_needed, max_output_tokens)

    # 6. Generate warning if batch was reduced
    warning = None
    if batch_size and safe_size < batch_size:
        total_estimated = estimated_input + max_output_needed
        warning = (
            f"Reduced batch size from {batch_size} to {safe_size} "
            f"(estimated {total_estimated} tokens vs {context_window} window)"
        )

    return {
        "batch_size_adjusted": safe_size,
        "estimated_input_tokens": estimated_input,
        "estimated_output_tokens": estimated_output,
        "max_output_needed": max_output_needed,
        "warning": warning,
    }
# endregion estimate_token_budget
# #endregion estimate_token_budget