# #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 DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_OUTPUT_TOKENS, estimate_token_budget


# 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 TestTokenBudget [TYPE Class]
# @BRIEF Test suite for estimate_token_budget.
class TestTokenBudget:

    # region test_small_rows_fit_at_requested_size [TYPE Function]
    # @BRIEF Short text rows fill the requested batch_size without reduction.
    def test_small_rows_fit_at_requested_size(self):
        """Short text of ~50 chars should fit 50-row batch easily."""
        rows = [_make_row("Hello world, this is a short text for translation.")] * 50
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru"],
            batch_size=50,
        )
        assert result["batch_size_adjusted"] == 50
        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
# #endregion TestTokenBudget
