# #region AdaptiveBatchSizer [C:3] [TYPE Module] [SEMANTICS translate, batch, sizing, token-budget]
# @defgroup Translate Module group.
# @BRIEF Adaptive batch sizing for LLM translation — splits source rows into variable-sized
#         batches based on actual content length and token budget estimates.
#        Output budget is the PRIMARY constraint; input budget is secondary.
# @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).
#            Safety factor applied to input budget to account for tokenizer variance.
#            Output-aware row cap computed from estimate_token_budget result.
#            Provider-level context_window/max_output_tokens take priority over PROVIDER_DEFAULTS.
# @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 (
    INPUT_SAFETY_FACTOR,
    PROMPT_BASE_TOKENS,
    estimate_token_budget,
)
from ._utils import estimate_row_tokens


# #region AdaptiveBatchSizer [C:3] [TYPE Class]
# @defgroup Translate Module group.
# @BRIEF Split source rows into auto-sized batches based on token budget estimates.
class AdaptiveBatchSizer:
    def __init__(self, db: Session, config_manager: ConfigManager) -> None:
        self.db = db
        self.config_manager = config_manager

    # #region resolve_provider_config [C:2] [TYPE Function] [SEMANTICS llm, provider, model, token-config]
# @ingroup Translate
    # @BRIEF Resolve provider token config (model name + token limits) for budget estimation.
    #        DB values (context_window, max_output_tokens) take priority over PROVIDER_DEFAULTS.
    # @POST Returns dict with model, context_window, max_output_tokens. None values = use fallback.
    # @SIDE_EFFECT DB query to LLM provider table.
    def resolve_provider_config(self, job: TranslationJob) -> dict:
        """Resolve provider token config for token budget estimation."""
        if not job.provider_id:
            return {"model": None, "context_window": None, "max_output_tokens": None}
        try:
            token_config = LLMProviderService(self.db).get_provider_token_config(job.provider_id)
            return token_config
        except Exception:
            return {"model": None, "context_window": None, "max_output_tokens": None}
    # #endregion resolve_provider_config

    # #region auto_size_batches [C:3] [TYPE Function] [SEMANTICS translate, batch, sizing]
# @ingroup Translate
    # @BRIEF Split source rows into variable-sized batches based on content length.
    #        Output budget (max_rows_by_output) is the PRIMARY row-count constraint.
    # @PRE source_rows is non-empty. job has valid config.
    # @POST Returns list of batches. Each batch fits within estimated input token budget
    #        with INPUT_SAFETY_FACTOR headroom AND within output row cap.
    # @SIDE_EFFECT DB query to resolve provider config. 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]]:
        with belief_scope("AdaptiveBatchSizer.auto_size_batches"):
            if not source_rows:
                return []

            # Resolve provider config (DB values > PROVIDER_DEFAULTS)
            token_config = self.resolve_provider_config(job)
            if provider_info is not None and token_config["model"] is None:
                token_config["model"] = provider_info

            # 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 (uses provider config from DB if available)
            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=token_config["model"],
                context_window=token_config["context_window"],
                max_output_tokens=token_config["max_output_tokens"],
            )

            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 with safety factor
            available_input = budget.get("available_input_budget")
            if available_input is not None:
                raw_budget = available_input - PROMPT_BASE_TOKENS
                per_batch_budget = int(raw_budget * INPUT_SAFETY_FACTOR)
            else:
                estimated_input = budget.get("estimated_input_tokens", 50000)
                raw_budget = estimated_input - PROMPT_BASE_TOKENS
                per_batch_budget = int(raw_budget * INPUT_SAFETY_FACTOR)

            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. Compute max rows per batch — OUTPUT BUDGET IS PRIMARY CONSTRAINT
            max_rows_by_output = budget.get("max_rows_by_output", 20)

            # Absolute hard cap as safety net (not the primary constraint)
            absolute_hard_cap = 50
            max_rows_hard_cap = min(max_rows_by_output, absolute_hard_cap)

            # Respect job.batch_size as user override
            if job.batch_size:
                max_rows_hard_cap = min(max_rows_hard_cap, job.batch_size)

            # 5. Greedy batch splitting with output constraint
            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)

            # 6. Log stats
            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,
                    "safety_factor": INPUT_SAFETY_FACTOR,
                    "max_rows_by_output": max_rows_by_output,
                    "max_rows_hard_cap": max_rows_hard_cap,
                })

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