# #region AdaptiveBatchSizer [C:3] [TYPE Module] [SEMANTICS translate, batch, sizing, token-budget]
# @BRIEF Adaptive batch sizing for LLM translation — splits source rows into variable-sized
#         batches based on actual content length and token budget estimates.
# @LAYER Domain
# @RELATION DEPENDS_ON -> [estimate_token_budget]
# @RELATION DEPENDS_ON -> [LLMProviderService]
# @RELATION DEPENDS_ON -> [TranslationJob]
# @RATIONALE Extracted from TranslationExecutor to comply with INV_7 (module < 400 lines).
#            Fixed batch_size of 50 wastes LLM context for short rows and overflows for
#            long rows. Variable sizing maximizes throughput while preventing truncation.
# @REJECTED Fixed batch_size of 50 — causes truncation on long-content rows.
#           Single monolithic batch — would lose all progress on any failure.


from sqlalchemy.orm import Session

from ...core.config_manager import ConfigManager
from ...core.logger import belief_scope, logger
from ...models.translate import TranslationJob
from ...services.llm_provider import LLMProviderService
from ._token_budget import (
    JSON_OVERHEAD_PER_ROW,
    MAX_OUTPUT_HEADROOM,
    OUTPUT_PER_ROW_PER_LANG,
    PROMPT_BASE_TOKENS,
    REASONING_OVERHEAD,
    estimate_token_budget,
)
from ._utils import estimate_row_tokens


# #region AdaptiveBatchSizer [C:3] [TYPE Class]
# @BRIEF Split source rows into auto-sized batches based on token budget estimates.
class AdaptiveBatchSizer:
    """Split source rows into auto-sized batches based on token budget estimates.

    Each batch is sized so that its total estimated tokens fit within the
    available context window (input budget), accounting for prompt overhead,
    dictionary entries, and output tokens.
    """

    def __init__(self, db: Session, config_manager: ConfigManager) -> None:
        self.db = db
        self.config_manager = config_manager

    # #region resolve_provider_model [C:2] [TYPE Function] [SEMANTICS llm, provider, model]
    # @BRIEF Resolve the LLM provider model name for token budget estimation.
    # @POST Returns model name string or None if resolution fails.
    # @SIDE_EFFECT DB query to LLM provider table.
    def resolve_provider_model(self, job: TranslationJob) -> str | None:
        """Resolve the provider model name for token budget estimation."""
        if not job.provider_id:
            return None
        try:
            p_svc = LLMProviderService(self.db)
            p = p_svc.get_provider(job.provider_id)
            if p:
                return p.default_model or "gpt-4o-mini"
        except Exception:
            pass
        return None
    # #endregion resolve_provider_model

    # #region auto_size_batches [C:3] [TYPE Function] [SEMANTICS translate, batch, sizing]
    # @BRIEF Split source rows into variable-sized batches based on content length.
    # @PRE source_rows is non-empty. job has valid config.
    # @POST Returns list of batches, each batch is a list of row dicts.
    #        Each batch fits within the estimated token budget for its rows.
    # @SIDE_EFFECT DB query to resolve provider model. Logs batch statistics.
    def auto_size_batches(
        self,
        job: TranslationJob,
        source_rows: list[dict],
        target_languages: list[str],
        provider_info: str | None = None,
    ) -> list[list[dict]]:
        """Split source rows into auto-sized batches based on content length.

        Each batch is sized so that its total estimated tokens fit within the
        available context window (input budget), accounting for prompt overhead,
        dictionary entries, and output tokens.
        """
        with belief_scope("AdaptiveBatchSizer.auto_size_batches"):
            if not source_rows:
                return []

            if provider_info is None:
                provider_info = self.resolve_provider_model(job)

            # 1. Estimate per-row token counts
            row_tokens: list[int] = []
            for row in source_rows:
                source_text = row.get("source_text", "")
                source_data = row.get("source_data")
                tokens = estimate_row_tokens(source_text, source_data, job)
                row_tokens.append(tokens)

            # 2. Get budget recommendation
            budget = estimate_token_budget(
                source_rows=source_rows,
                target_languages=target_languages,
                source_column=job.translation_column or "source_text",
                context_columns=job.context_columns,
                batch_size=len(source_rows),
                provider_info=provider_info,
            )

            recommended = budget.get("batch_size_adjusted", 0)

            # Fallback: if budget calculation fails, use fixed size
            if recommended <= 0:
                fallback_size = job.batch_size or 50
                logger.explore("Token budget returned zero — falling back to fixed batch size", {
                    "fallback_size": fallback_size,
                    "total_rows": len(source_rows),
                })
                return [
                    source_rows[i:i + fallback_size]
                    for i in range(0, len(source_rows), fallback_size)
                ]

            # 3. Compute per-batch row-content budget
            # Use the ACTUAL available input capacity (context_window - max_output_tokens),
            # NOT the sum of the first N rows. This prevents long rows from being
            # incorrectly placed in 1-row batches when the context window has plenty of room.
            available_input = budget.get("available_input_budget")
            if available_input is not None:
                # New-style: use actual available input capacity
                per_batch_budget = available_input - PROMPT_BASE_TOKENS
            else:
                # Fallback for tests: use estimated_input (sum of first N rows)
                estimated_input = budget.get("estimated_input_tokens", 50000)
                per_batch_budget = estimated_input - PROMPT_BASE_TOKENS

            if per_batch_budget <= 0:
                fallback_size = job.batch_size or 50
                logger.explore("Per-batch budget collapsed — falling back to fixed batch size", {
                    "per_batch_budget": per_batch_budget,
                    "fallback_size": fallback_size,
                    "total_rows": len(source_rows),
                })
                return [
                    source_rows[i:i + fallback_size]
                    for i in range(0, len(source_rows), fallback_size)
                ]

            # 4. Greedy batch splitting
            # Compute max rows per batch from OUTPUT constraint.
            # This prevents output truncation (finish_reason=length) when batching
            # many short rows within the large input budget.
            num_languages = len(target_languages)
            max_output_tokens_val = budget.get("max_output_tokens")
            if max_output_tokens_val is not None:
                output_per_row = num_languages * OUTPUT_PER_ROW_PER_LANG + JSON_OVERHEAD_PER_ROW
                available_output = max_output_tokens_val - REASONING_OVERHEAD - MAX_OUTPUT_HEADROOM
                max_rows_by_output = max(available_output // output_per_row, 1) if output_per_row > 0 else 20
                max_rows_hard_cap = max_rows_by_output
            else:
                # Fallback for tests: old formula
                max_rows_hard_cap = max(recommended * 2, 20)

            # 5. Respect job.batch_size as the absolute maximum rows per batch.
            #    User-configured batch_size overrides model-based estimates to
            #    prevent LLM quality degradation on large batches.
            if job.batch_size:
                max_rows_hard_cap = min(max_rows_hard_cap, job.batch_size)

            batches: list[list[dict]] = []
            current_batch: list[dict] = []
            current_tokens = 0

            for i, row in enumerate(source_rows):
                rt = row_tokens[i]

                if rt > per_batch_budget:
                    if current_batch:
                        batches.append(current_batch)
                        current_batch = []
                        current_tokens = 0
                    logger.reason("Single row exceeds per-batch token budget — placing in own batch", {
                        "row_index": i,
                        "row_tokens": rt,
                        "per_batch_budget": per_batch_budget,
                    })
                    batches.append([row])
                    continue

                should_split = bool(
                    current_batch
                    and (current_tokens + rt > per_batch_budget
                         or len(current_batch) >= max_rows_hard_cap)
                )

                if should_split:
                    batches.append(current_batch)
                    current_batch = [row]
                    current_tokens = rt
                else:
                    current_batch.append(row)
                    current_tokens += rt

            if current_batch:
                batches.append(current_batch)

            # 5. Log adaptive batch statistics
            if batches:
                avg_size = sum(len(b) for b in batches) / max(1, len(batches))
                max_size = max(len(b) for b in batches)
                logger.reason("Auto-sized batches", {
                    "total_rows": len(source_rows),
                    "num_batches": len(batches),
                    "avg_batch_size": round(avg_size, 1),
                    "max_batch_size": max_size,
                    "recommended_batch_size": recommended,
                    "per_batch_budget": per_batch_budget,
                })

            return batches
    # #endregion auto_size_batches
# #endregion AdaptiveBatchSizer
# #endregion AdaptiveBatchSizer
