# #region TestTokenBudget [C:3] [TYPE Module] [SEMANTICS test, token, budget, estimation, batch, translate]
# @BRIEF Verify estimate_token_budget contracts — safe batch sizing, auto-reduction, warning generation.
# @RELATION BINDS_TO -> [estimate_token_budget]
# @TEST_EDGE  empty_rows — empty source rows returns batch_size_adjusted=1
# @TEST_EDGE  small_rows — short text fits in a single batch at requested size
# @TEST_EDGE  large_rows — long text causes batch size reduction
# @TEST_EDGE  multi_language — more target languages increases output estimate
# @TEST_EDGE  context_columns — context columns increase input token estimate
# @TEST_EDGE  dictionary_entries — glossary entries increase input token estimate
# @TEST_EDGE  auto_calc — no batch_size specified, auto-calculates max safe
# @TEST_EDGE  exact_fit — exactly fits context window, batch_adjusted == requested
# @TEST_EDGE  conservative_min — even with huge rows, batch_size_adjusted >= 1
# @TEST_INVARIANT  batch_size_adjusted >= 1 always
# @TEST_INVARIANT  max_output_needed between MIN_MAX_TOKENS(4096) and max_output_tokens(8192)
# @TEST_INVARIANT  warning is None when batch fits, str when reduced

from src.plugins.translate._token_budget import (
    CJK_RATIO,
    DEFAULT_CONTEXT_WINDOW,
    DEFAULT_MAX_OUTPUT_TOKENS,
    OTHER_RATIO,
    estimate_token_budget,
    _compute_max_rows_by_output,
    _estimate_tokens_for_text,
)


# region _make_row [TYPE Function]
# @BRIEF Create a test source row dict.
def _make_row(text: str, **context) -> dict:
    row = {"source_text": text, "row_index": "0"}
    row.update(context)
    return row
# endregion _make_row


# region _make_cjk_row [TYPE Function]
# @BRIEF Create a test source row with CJK text.
def _make_cjk_row() -> dict:
    return {"source_text": "你好世界这是一个测试消息", "row_index": "0"}
# endregion _make_cjk_row


# region TestTokenBudget [TYPE Class]
# @BRIEF Test suite for estimate_token_budget and related token estimation functions.
class TestTokenBudget:

    # region test_cjk_ratio_estimate [TYPE Function]
    # @BRIEF Verify CJK token estimation uses the conservative CJK_RATIO.
    def test_cjk_ratio_estimate(self):
        """12 CJK chars / 1.0 = 12 tokens (conservative), plus 1 for empty context."""
        tokens = _estimate_tokens_for_text("你好世界这是一个测试消息")
        expected = int(12 / CJK_RATIO)
        assert tokens == expected, f"CJK estimate {tokens} != expected {expected}"
    # endregion test_cjk_ratio_estimate

    # region test_mixed_text_estimate [TYPE Function]
    # @BRIEF Verify mixed CJK+Ltn text estimation.
    def test_mixed_text_estimate(self):
        """Mixed text: CJK at CJK_RATIO, Latin at OTHER_RATIO."""
        text = "你好世界! Hello world, this is a test!"
        tokens = _estimate_tokens_for_text(text)
        # 4 CJK chars / 1.0 = 4, 30 non-CJK / 1.8 = 16, total = 20
        assert tokens == 20, f"Expected 20, got {tokens}"
    # endregion test_mixed_text_estimate

    # region test_empty_text_estimate [TYPE Function]
    # @BRIEF Empty text returns 1 token minimum.
    def test_empty_text_estimate(self):
        assert _estimate_tokens_for_text("") == 1
        assert _estimate_tokens_for_text(None) == 1
    # endregion test_empty_text_estimate

    # region test_compute_max_rows_by_output_single_lang [TYPE Function]
    # @BRIEF Single target language: compute max rows from output budget.
    def test_compute_max_rows_by_output_single_lang(self):
        """With max_output_tokens=8192 and 1 language, should return limited rows."""
        max_rows = _compute_max_rows_by_output(max_output_tokens=8192, num_languages=1)
        assert max_rows >= 1
        assert max_rows < 100  # Should be reasonable
    # endregion test_compute_max_rows_by_output_single_lang

    # region test_compute_max_rows_by_output_multi_lang [TYPE Function]
    # @BRIEF Multiple target languages reduce max rows from output budget.
    def test_compute_max_rows_by_output_multi_lang(self):
        """More languages = fewer rows that fit in the same output budget."""
        single = _compute_max_rows_by_output(max_output_tokens=16384, num_languages=1)
        multi = _compute_max_rows_by_output(max_output_tokens=16384, num_languages=4)
        assert multi <= single
        assert multi >= 1
    # endregion test_compute_max_rows_by_output_multi_lang

    # region test_compute_max_rows_by_output_small_budget [TYPE Function]
    # @BRIEF Very small output budget returns at least 1 row.
    def test_compute_max_rows_by_output_small_budget(self):
        """Even with minimal budget, at least 1 row fits."""
        max_rows = _compute_max_rows_by_output(max_output_tokens=4096, num_languages=2)
        assert max_rows >= 1
    # endregion test_compute_max_rows_by_output_small_budget

    # region test_max_rows_by_output_in_return_dict [TYPE Function]
    # @BRIEF estimate_token_budget returns max_rows_by_output field.
    def test_max_rows_by_output_in_return_dict(self):
        """The return dict includes max_rows_by_output as a positive int."""
        rows = [_make_row("Short text.")] * 10
        result = estimate_token_budget(rows, ["ru"], batch_size=10)
        assert "max_rows_by_output" in result
        assert isinstance(result["max_rows_by_output"], int)
        assert result["max_rows_by_output"] >= 1
    # endregion test_max_rows_by_output_in_return_dict

    # region test_max_rows_by_output_changes_with_languages [TYPE Function]
    # @BRIEF max_rows_by_output decreases with more target languages.
    def test_max_rows_by_output_changes_with_languages(self):
        """More target languages = smaller max_rows_by_output."""
        rows = [_make_row("Hello world")] * 10
        single = estimate_token_budget(rows, ["ru"], batch_size=10)
        multi = estimate_token_budget(rows, ["ru", "en", "fr", "de"], batch_size=10)
        assert multi["max_rows_by_output"] <= single["max_rows_by_output"]
    # endregion test_max_rows_by_output_changes_with_languages

    # region test_small_rows_fit_at_requested_size [TYPE Function]
    # @BRIEF Short text rows fill the requested batch_size without reduction.
    # @RATIONALE OUTPUT_PER_ROW_PER_LANG=200 means 1 lang × 200 + 50 overhead = 250 tok/row.
    #            With DEFAULT_MAX_OUTPUT_TOKENS=16384 and overhead 5000, max rows ≈ 45.
    def test_small_rows_fit_at_requested_size(self):
        """Short text of ~50 chars should fill the maximum output-limited batch size."""
        rows = [_make_row("Hello world, this is a short text for translation.")] * 45
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru"],
            batch_size=45,
        )
        assert result["batch_size_adjusted"] == 45
        assert result["estimated_input_tokens"] > 0
        assert result["estimated_output_tokens"] > 0
        assert result["max_output_needed"] >= 4096
        assert result["max_output_needed"] <= DEFAULT_MAX_OUTPUT_TOKENS
        assert result["warning"] is None
    # endregion test_small_rows_fit_at_requested_size

    # region test_large_rows_reduce_batch_size [TYPE Function]
    # @BRIEF Very long text rows force batch size reduction to fit context window.
    def test_large_rows_reduce_batch_size(self):
        """~10000 char rows should cause batch reduction from 50 to a smaller number."""
        long_text = "X" * 10000  # ~4545 tokens each at 2.2 chars/token
        rows = [_make_row(long_text)] * 50
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru", "en"],
            batch_size=50,
        )
        # Each row is ~4545 tokens input, plus ~2400 output (2 langs * 60 * 20 + 2000)
        # With 50 rows: ~227K tokens — way over 64K. Should reduce significantly.
        assert result["batch_size_adjusted"] < 50
        assert result["batch_size_adjusted"] >= 1
        assert result["warning"] is not None
        assert "Reduced batch size" in result["warning"]
    # endregion test_large_rows_reduce_batch_size

    # region test_multi_language_increases_output_estimate [TYPE Function]
    # @BRIEF More target languages increase the estimated output tokens.
    def test_multi_language_increases_output_estimate(self):
        """4 target languages should have higher output estimate than 1."""
        rows = [_make_row("Short text.")] * 10

        single_lang = estimate_token_budget(rows, ["ru"], batch_size=10)
        multi_lang = estimate_token_budget(rows, ["ru", "en", "fr", "de"], batch_size=10)

        assert multi_lang["estimated_output_tokens"] > single_lang["estimated_output_tokens"]
        # Both should fit without reduction
        assert single_lang["batch_size_adjusted"] == 10
        assert multi_lang["batch_size_adjusted"] == 10
    # endregion test_multi_language_increases_output_estimate

    # region test_context_columns_increase_input [TYPE Function]
    # @BRIEF Context columns add to the input token estimate.
    def test_context_columns_increase_input(self):
        """Rows with context columns should have higher input tokens."""
        rows = [_make_row("Short text.")] * 10

        no_context = estimate_token_budget(rows, ["ru"], batch_size=10)
        with_context = estimate_token_budget(
            rows, ["ru"], batch_size=10,
            context_columns=["description", "category"],
        )

        assert with_context["estimated_input_tokens"] > no_context["estimated_input_tokens"]
    # endregion test_context_columns_increase_input

    # region test_dictionary_entries_increase_input [TYPE Function]
    # @BRIEF Dictionary entries add to the input token estimate.
    def test_dictionary_entries_increase_input(self):
        """Dictionary entries should increase estimated_input_tokens."""
        rows = [_make_row("Short text.")] * 10
        dict_entries = [{"source_term": f"term_{i}", "target_term": f"trans_{i}"} for i in range(100)]

        no_dict = estimate_token_budget(rows, ["ru"], batch_size=10)
        with_dict = estimate_token_budget(rows, ["ru"], batch_size=10, dictionary_entries=dict_entries)

        assert with_dict["estimated_input_tokens"] > no_dict["estimated_input_tokens"]
    # endregion test_dictionary_entries_increase_input

    # region test_auto_calc_finds_safe_size [TYPE Function]
    # @BRIEF Without batch_size specified, auto-calculate max safe batch.
    def test_auto_calc_finds_safe_size(self):
        """Auto-calculation should find the max rows that fit context window."""
        rows = [_make_row("A" * 500)] * 100  # ~227 tokens/row

        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru", "en"],
            batch_size=None,  # auto-calc
        )

        assert result["batch_size_adjusted"] >= 1
        assert result["batch_size_adjusted"] <= 100
        # Verify it found a reasonable size
        assert result["estimated_input_tokens"] + result["max_output_needed"] <= DEFAULT_CONTEXT_WINDOW
    # endregion test_auto_calc_finds_safe_size

    # region test_empty_rows_returns_minimum [TYPE Function]
    # @BRIEF Empty source rows returns batch_size_adjusted=1 and minimum estimates.
    def test_empty_rows_returns_minimum(self):
        """Empty rows list should return batch_size_adjusted=1 with min estimates."""
        result = estimate_token_budget(
            source_rows=[],
            target_languages=["ru"],
            batch_size=10,
        )
        assert result["batch_size_adjusted"] == 1
        assert result["estimated_input_tokens"] >= 300  # at least prompt base
        assert result["estimated_output_tokens"] >= 2000  # at least reasoning overhead
    # endregion test_empty_rows_returns_minimum

    # region test_exact_fit_no_warning [TYPE Function]
    # @BRIEF When batch fits exactly, no warning is generated.
    def test_exact_fit_no_warning(self):
        """Small batch that fits easily should have no warning."""
        rows = [_make_row("Short text.")] * 5
        result = estimate_token_budget(rows, ["ru"], batch_size=5)
        assert result["warning"] is None
        assert result["batch_size_adjusted"] == 5
    # endregion test_exact_fit_no_warning

    # region test_huge_rows_still_return_min_one [TYPE Function]
    # @BRIEF Even with massive rows, batch_size_adjusted is at least 1.
    def test_huge_rows_still_return_min_one(self):
        """Massive rows that exceed context window still return batch_size_adjusted >= 1."""
        huge_text = "X" * 500000  # ~227K tokens — exceeds context window
        rows = [_make_row(huge_text)] * 10
        result = estimate_token_budget(rows, ["ru", "en", "fr"], batch_size=10)
        assert result["batch_size_adjusted"] >= 1
    # endregion test_huge_rows_still_return_min_one

    # region test_max_output_needed_bounds [TYPE Function]
    # @BRIEF max_output_needed stays within reasonable bounds.
    def test_max_output_needed_bounds(self):
        """max_output_needed is between MIN_MAX_TOKENS and max_output_tokens."""
        rows = [_make_row("Short text.")] * 5
        result = estimate_token_budget(rows, ["ru", "en"], batch_size=5)
        assert 4096 <= result["max_output_needed"] <= DEFAULT_MAX_OUTPUT_TOKENS
    # endregion test_max_output_needed_bounds

    # region test_single_row_never_reduced [TYPE Function]
    # @BRIEF Batch of 1 should never be reduced (batch_size_adjusted >= requested when requested=1).
    def test_single_row_never_reduced(self):
        """Even a very long single row should still fit (batch_size_adjusted=1)."""
        long_text = "X" * 100000  # large but still fits context if alone
        rows = [_make_row(long_text)]
        result = estimate_token_budget(rows, ["ru", "en", "fr", "de"], batch_size=1)
        assert result["batch_size_adjusted"] == 1
    # endregion test_single_row_never_reduced

    # region test_no_target_languages_defaults_to_en [TYPE Function]
    # @BRIEF When target_languages is empty, defaults to ["en"].
    def test_no_target_languages_defaults_to_en(self):
        """Empty target_languages defaults to ['en'] without error."""
        rows = [_make_row("Short text.")] * 5
        result = estimate_token_budget(rows, [], batch_size=5)
        assert result["batch_size_adjusted"] == 5
        assert result["estimated_output_tokens"] > 0
    # endregion test_no_target_languages_defaults_to_en

# endregion TestTokenBudget


# region TestProviderDefaults [TYPE Class]
# @BRIEF Test suite for PROVIDER_DEFAULTS resolution in estimate_token_budget.
class TestProviderDefaults:

    # region test_qwen_flash_provider_defaults [TYPE Function]
    # @BRIEF qwen-flash resolves to context_window=131072, max_output_tokens=32768.
    def test_qwen_flash_provider_defaults(self):
        """qwen-flash model name triggers PROVIDER_DEFAULTS with known caps."""
        rows = [_make_row("Short text.")] * 5
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru"],
            provider_info="qwen-flash",
            batch_size=5,
        )
        # qwen-flash has context_window=131072, max_output_tokens=32768
        assert result["max_output_tokens"] == 32768
        # available_input_budget = context_window - max_output_tokens
        assert result["available_input_budget"] == 131072 - 32768
        assert result["warning"] is None
    # endregion test_qwen_flash_provider_defaults

    # region test_qwen_flash_overridden_by_explicit [TYPE Function]
    # @BRIEF Explicit context_window/max_output_tokens override provider defaults.
    def test_qwen_flash_overridden_by_explicit(self):
        """When both provider_info and explicit caps given, explicit wins."""
        rows = [_make_row("Short text.")] * 5
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru"],
            provider_info="qwen-flash",
            context_window=99999,
            max_output_tokens=20000,
            batch_size=5,
        )
        assert result["max_output_tokens"] == 20000
        assert result["available_input_budget"] == 99999 - 20000
    # endregion test_qwen_flash_overridden_by_explicit

# endregion TestProviderDefaults


# region TestOutputPerRow [TYPE Class]
# @BRIEF Test suite for OUTPUT_PER_ROW_PER_LANG=200 constraint in _compute_max_rows_by_output.
class TestOutputPerRow:

    # region test_output_per_row_lang_200 [TYPE Function]
    # @BRIEF _compute_max_rows_by_output uses OUTPUT_PER_ROW_PER_LANG=200.
    def test_output_per_row_lang_200(self):
        """Verify the max rows calculation uses 200 tok/row/lang."""
        from src.plugins.translate._token_budget import (
            OUTPUT_PER_ROW_PER_LANG,
            REASONING_OVERHEAD,
            MAX_OUTPUT_HEADROOM,
            OUTPUT_SAFETY_FACTOR,
        )

        num_languages = 1
        max_output_tokens = 16384

        max_rows = _compute_max_rows_by_output(
            max_output_tokens=max_output_tokens,
            num_languages=num_languages,
        )

        # Manual computation matching the formula:
        # overhead = REASONING_OVERHEAD + MAX_OUTPUT_HEADROOM = 2000 + 3000 = 5000
        # per_row = 1 * 200 + 50 = 250
        # available = int((16384 - 5000) * 0.55) = int(11384 * 0.55) = int(6261.2) = 6261
        # max(6261 // 250, 1) = max(25, 1) = 25
        overhead = REASONING_OVERHEAD + MAX_OUTPUT_HEADROOM
        per_row = num_languages * OUTPUT_PER_ROW_PER_LANG + 50
        available = int((max_output_tokens - overhead) * OUTPUT_SAFETY_FACTOR)
        expected = max(available // per_row, 1)

        assert max_rows == expected, (
            f"max_rows={max_rows} != computed={expected}. "
            f"Check that OUTPUT_PER_ROW_PER_LANG={OUTPUT_PER_ROW_PER_LANG} is used."
        )
    # endregion test_output_per_row_lang_200

    # region test_output_per_row_lang_200_multi [TYPE Function]
    # @BRIEF More languages reduce max rows proportionally (OUTPUT_PER_ROW_PER_LANG=200).
    def test_output_per_row_lang_200_multi(self):
        """4 languages produce fewer max_rows than 1 language."""
        single = _compute_max_rows_by_output(max_output_tokens=32768, num_languages=1)
        multi = _compute_max_rows_by_output(max_output_tokens=32768, num_languages=4)

        # 4 languages = 4x per-row cost → fewer rows fit
        assert multi < single
        assert multi >= 1
    # endregion test_output_per_row_lang_200_multi

# endregion TestOutputPerRow


# #region TestOutputSafetyFactor [C:3] [TYPE Class] [SEMANTICS test,token,budget,safety-factor,regression]
# @BRIEF Regression tests for OUTPUT_SAFETY_FACTOR=0.55 — verifies it prevents
#         finish_reason=length by being conservative (≤0.70) while not falling
#         below the realistic minimum for qwen-flash 4-language batch sizing.
# @RATIONALE The factor was lowered from 0.70→0.55 to be MORE conservative,
#            reducing max_rows_by_output from an overoptimistic 36 to a realistic 17.
#            The actual batch size was always ~17 due to INPUT budget constraints;
#            this fix just made the output estimate match reality, preventing
#            finish_reason=length truncation cascades.
class TestOutputSafetyFactor:

    # #region test_output_safety_factor_not_above_070 [C:3] [TYPE Function]
    # @BRIEF OUTPUT_SAFETY_FACTOR must be ≤ 0.70 — the conservative guard.
    # @RATIONALE Higher values (e.g. 0.75) would overestimate output capacity,
    #            causing finish_reason=length and LLM retry cascades.
    def test_output_safety_factor_not_above_070(self):
        from src.plugins.translate._token_budget import OUTPUT_SAFETY_FACTOR
        assert OUTPUT_SAFETY_FACTOR <= 0.70, (
            f"OUTPUT_SAFETY_FACTOR={OUTPUT_SAFETY_FACTOR}, must be ≤ 0.70"
        )
        assert OUTPUT_SAFETY_FACTOR >= 0.40, (
            f"OUTPUT_SAFETY_FACTOR={OUTPUT_SAFETY_FACTOR}, too conservative"
        )
    # #endregion test_output_safety_factor_not_above_070

    # #region test_max_rows_by_output_qwen_flash_4langs [C:3] [TYPE Function]
    # @BRIEF qwen-flash + 4 target languages: max_rows_by_output ≥ 14 and ≤ 24.
    # @RATIONALE Production: input-budget limits batches to ~17 rows anyway.
    #            Output budget estimate must be realistic (14-24 range).
    def test_max_rows_by_output_qwen_flash_4langs(self):
        from src.plugins.translate._token_budget import _compute_max_rows_by_output

        max_rows = _compute_max_rows_by_output(
            max_output_tokens=32768,  # qwen-flash
            num_languages=4,
        )

        assert 14 <= max_rows <= 24, (
            f"max_rows_by_output={max_rows}, expected 14-24 range."
        )
    # #endregion test_max_rows_by_output_qwen_flash_4langs

    # #region test_output_safety_factor_consistent_with_per_row [C:2] [TYPE Function]
    # @BRIEF max_rows_by_output is consistent with OUTPUT_PER_ROW_PER_LANG.
    def test_output_safety_factor_consistent_with_per_row(self):
        from src.plugins.translate._token_budget import (
            _compute_max_rows_by_output,
            OUTPUT_SAFETY_FACTOR,
            OUTPUT_PER_ROW_PER_LANG,
            REASONING_OVERHEAD,
            MAX_OUTPUT_HEADROOM,
            JSON_OVERHEAD_PER_ROW,
        )

        # Manual computation must match
        max_tok = 32768
        n_lang = 4
        overhead = REASONING_OVERHEAD + MAX_OUTPUT_HEADROOM
        per_row = n_lang * OUTPUT_PER_ROW_PER_LANG + JSON_OVERHEAD_PER_ROW
        available = int((max_tok - overhead) * OUTPUT_SAFETY_FACTOR)
        expected = max(available // per_row, 1)

        actual = _compute_max_rows_by_output(max_tok, n_lang)
        assert actual == expected, (
            f"actual={actual} != expected={expected}"
        )
    # #endregion test_output_safety_factor_consistent_with_per_row

    # #region test_single_lang_output_rows_above_20 [C:2] [TYPE Function]
    # @BRIEF Single target language with 0.55: max rows ≥ 20.
    def test_single_lang_output_rows_above_20(self):
        from src.plugins.translate._token_budget import _compute_max_rows_by_output

        max_rows = _compute_max_rows_by_output(
            max_output_tokens=16384,
            num_languages=1,
        )
        assert max_rows >= 20, f"Expected ≥20 rows for 1 lang, got {max_rows}"
    # #endregion test_single_lang_output_rows_above_20

# #endregion TestOutputSafetyFactor
# #endregion TestTokenBudget
