# #region preview_prompt_helpers [C:2] [TYPE Module] [SEMANTICS token, budget, metadata, estimate]
# @BRIEF Token estimation metadata and budget helpers for preview prompt building.
# @RELATION DEPENDS_ON -> [estimate_token_budget]
# @RELATION DEPENDS_ON -> [TokenEstimator]

from typing import Any

from ...core.logger import logger
from ...models.translate import TranslationJob
from ._token_budget import DEFAULT_CONTEXT_WINDOW, estimate_token_budget
from .preview_token_estimator import TokenEstimator


# #region estimate_token_budget_for_rows [C:2] [TYPE Function] [SEMANTICS token, budget, estimate, truncate]
# @BRIEF Check token budget and optionally truncate source rows for preview.
# @POST Returns adjusted source_rows, actual_row_count, and token_budget dict.
def estimate_token_budget_for_rows(
    source_rows: list[dict[str, Any]],
    target_languages: list[str],
    job: TranslationJob,
    provider_model: str | None,
) -> dict[str, Any]:
    """Check token budget and optionally truncate rows for preview."""
    actual_row_count = len(source_rows)
    token_budget_dict = estimate_token_budget(
        source_rows=source_rows,
        target_languages=target_languages,
        source_column=job.translation_column,
        context_columns=job.context_columns,
        batch_size=actual_row_count,
        provider_info=provider_model,
    )

    adjusted_size = token_budget_dict["batch_size_adjusted"]
    if adjusted_size < actual_row_count:
        logger.explore(
            f"Reducing preview from {actual_row_count} to {adjusted_size} rows",
            {
                "estimated_input": token_budget_dict["estimated_input_tokens"],
                "estimated_output": token_budget_dict["estimated_output_tokens"],
                "context_window": DEFAULT_CONTEXT_WINDOW,
            },
        )
        source_rows = source_rows[:adjusted_size]
        actual_row_count = len(source_rows)
        token_budget_dict = estimate_token_budget(
            source_rows=source_rows,
            target_languages=target_languages,
            source_column=job.translation_column,
            context_columns=job.context_columns,
            batch_size=actual_row_count,
            provider_info=provider_model,
        )

    return {
        "source_rows": source_rows,
        "actual_row_count": actual_row_count,
        "token_budget": token_budget_dict,
    }
# #endregion estimate_token_budget_for_rows


# #region compute_build_token_metadata [C:1] [TYPE Function] [SEMANTICS token, metadata, estimate]
# @BRIEF Compute token estimation metadata for prompt builder return dict.
# @POST Returns prompt, row_meta, target_languages, cost estimates, and cost_warning.
def compute_build_token_metadata(
    prompt: str,
    row_meta: list[dict[str, Any]],
    target_languages: list[str],
    num_languages: int,
    dictionary_section: str,
    actual_row_count: int,
    sample_size: int,
) -> dict[str, Any]:
    """Compute token estimation metadata for the prompt builder return dict."""
    sample_prompt_tokens = TokenEstimator.estimate_prompt_tokens(prompt)
    sample_output_tokens = TokenEstimator.estimate_output_tokens(actual_row_count, num_languages)
    sample_total_tokens = sample_prompt_tokens + sample_output_tokens
    sample_cost = TokenEstimator.estimate_cost(sample_total_tokens)

    total_est_rows = sample_size * 10
    total_est_tokens = (
        TokenEstimator.estimate_prompt_tokens(
            prompt.replace(str(actual_row_count), str(total_est_rows))
        )
        + TokenEstimator.estimate_output_tokens(total_est_rows, num_languages)
    )
    total_est_cost = TokenEstimator.estimate_cost(total_est_tokens)

    cost_warning = TokenEstimator.check_cost_warning(
        sample_size, num_languages, sample_total_tokens, sample_cost,
    )

    return {
        "prompt": prompt,
        "row_meta": row_meta,
        "target_languages": target_languages,
        "num_languages": num_languages,
        "dictionary_section": dictionary_section,
        "sample_prompt_tokens": sample_prompt_tokens,
        "sample_output_tokens": sample_output_tokens,
        "sample_total_tokens": sample_total_tokens,
        "sample_cost": sample_cost,
        "total_est_rows": total_est_rows,
        "total_est_tokens": total_est_tokens,
        "total_est_cost": total_est_cost,
        "cost_warning": cost_warning,
        "actual_row_count": actual_row_count,
    }
# #endregion compute_build_token_metadata
# #endregion preview_prompt_helpers
