# #region Test.TokenBudget [C:3] [TYPE Module] [SEMANTICS test, translate, token, budget, estimation]
# @BRIEF Tests for _token_budget.py — token estimation, batch sizing, output-aware constraints.
# @RELATION BINDS_TO -> [_token_budget]

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "src"))

import pytest

from src.plugins.translate._token_budget import (
    _estimate_tokens_for_text,
    _count_rows_that_fit,
    _calculate_output_tokens,
    _compute_max_rows_by_output,
    _apply_output_aware_batch_sizing,
    _build_warning,
    estimate_token_budget,
)


class TestEstimateTokensForText:
    """_estimate_tokens_for_text — CJK-aware heuristic token counting."""

    def test_empty_string(self):
        """Empty text returns 1."""
        assert _estimate_tokens_for_text("") == 1

    def test_ascii_only(self):
        """Pure ASCII text ~1.8 chars/token."""
        result = _estimate_tokens_for_text("hello world")
        assert result >= 1
        assert isinstance(result, int)

    def test_cjk_only(self):
        """Pure CJK text ~1.0 chars/token."""
        result = _estimate_tokens_for_text("你好世界")
        assert result >= 1

    def test_mixed_text(self):
        """Mixed CJK + ASCII."""
        result = _estimate_tokens_for_text("hello 你好 world 世界")
        assert result >= 1

    def test_cjk_ranges(self):
        """CJK includes Unicode ranges. Fullwidth forms."""
        result = _estimate_tokens_for_text("\uff01\uff02")
        assert result >= 1

    def test_long_text(self):
        """Long text estimate is proportional."""
        short = _estimate_tokens_for_text("hello")
        long_ = _estimate_tokens_for_text("hello " * 100)
        assert long_ > short


class TestCountRowsThatFit:
    """_count_rows_that_fit — consecutive rows within budget."""

    def test_all_rows_fit(self):
        """All rows fit within budget."""
        safe, total = _count_rows_that_fit([10, 20, 30], 5000)
        assert safe == 3
        assert total == 60

    def test_partial_fit(self):
        """Only first 2 rows fit."""
        safe, total = _count_rows_that_fit([10, 20, 5000], 2100)
        assert safe == 2
        assert total == 30

    def test_no_rows_fit(self):
        """First row exceeds budget."""
        safe, total = _count_rows_that_fit([5000], 1000)
        assert safe == 0
        assert total == 0

    def test_empty_input(self):
        """Empty list returns (0, 0)."""
        safe, total = _count_rows_that_fit([], 5000)
        assert safe == 0
        assert total == 0

    def test_single_row_just_fits(self):
        """Single row just fits (with reasoning overhead)."""
        budget = 200 + 2001  # row tokens + REASONING_OVERHEAD + 1
        safe, total = _count_rows_that_fit([200], budget)
        assert safe == 1
        assert total == 200


class TestCalculateOutputTokens:
    """_calculate_output_tokens — output budget estimation."""

    def test_basic(self):
        """Basic calculation returns positive int."""
        result = _calculate_output_tokens(
            safe_size=5, num_languages=2
        )
        assert result >= 1
        assert isinstance(result, int)

    def test_zero_safe_size(self):
        """Zero safe_size returns only overhead."""
        result = _calculate_output_tokens(safe_size=0, num_languages=2)
        assert result >= 2000  # REASONING_OVERHEAD

    def test_large_batch(self):
        """Large batch produces larger estimate."""
        small = _calculate_output_tokens(1, 1)
        large = _calculate_output_tokens(50, 5)
        assert large > small


class TestComputeMaxRowsByOutput:
    """_compute_max_rows_by_output — output-constrained row limit."""

    def test_basic(self):
        """Positive result for normal params."""
        result = _compute_max_rows_by_output(max_output_tokens=16384, num_languages=2)
        assert result >= 1

    def test_small_output_budget(self):
        """Tiny output budget returns 1."""
        result = _compute_max_rows_by_output(max_output_tokens=100, num_languages=2)
        assert result >= 1

    def test_zero_per_row(self):
        """Edge: zero per_row returns fallback 20."""
        result = _compute_max_rows_by_output(max_output_tokens=16384, num_languages=0)
        assert result >= 1


class TestApplyOutputAwareBatchSizing:
    """_apply_output_aware_batch_sizing — reduce batch until output fits."""

    def test_no_reduction_needed(self):
        """Batch fits without reduction."""
        result = _apply_output_aware_batch_sizing(
            safe_size=10, num_languages=1, max_output_tokens=50000
        )
        assert result == 10

    def test_full_reduction(self):
        """Batch reduced when output exceeds budget."""
        result = _apply_output_aware_batch_sizing(
            safe_size=10, num_languages=3, max_output_tokens=2000
        )
        assert result < 10

    def test_zero_size(self):
        """Zero safe_size returns 0."""
        result = _apply_output_aware_batch_sizing(
            safe_size=0, num_languages=1, max_output_tokens=50000
        )
        assert result == 0


class TestBuildWarning:
    """_build_warning — warning message generation."""

    def test_no_warning_when_size_ok(self):
        """No warning when batch size matches safe size."""
        assert _build_warning(10, 10, 20, 64000, 1000, 500, None) is None

    def test_reduced_batch_warning(self):
        """Warning when safe_size < requested batch_size."""
        w = _build_warning(10, 5, 20, 64000, 1000, 500, None)
        assert w is not None
        assert "Reduced" in w

    def test_auto_calc_warning(self):
        """Warning when auto-calculated batch smaller than total rows."""
        w = _build_warning(None, 5, 20, 64000, 1000, 500, None)
        assert w is not None
        assert "Auto-calculated" in w

    def test_dict_warning_appended(self):
        """Dictionary warning is appended."""
        w = _build_warning(10, 5, 20, 64000, 1000, 500, "Dict entries capped")
        assert w is not None
        assert "Dict entries capped" in w

    def test_dict_warning_only(self):
        """Dictionary warning alone when no batch reduction."""
        w = _build_warning(None, 20, 20, 64000, 1000, 500, "Dict capped")
        assert w == "Dict capped"


class TestEstimateTokenBudget:
    """estimate_token_budget — main budget estimation entry point."""

    def test_empty_source_rows(self):
        """Empty source rows returns sensible defaults."""
        result = estimate_token_budget(
            source_rows=[],
            target_languages=["ru"],
        )
        assert result["batch_size_adjusted"] >= 1
        assert "estimated_input_tokens" in result
        assert "estimated_output_tokens" in result
        assert "max_output_needed" in result

    def test_single_row_single_lang(self):
        """Single row, single language."""
        result = estimate_token_budget(
            source_rows=[{"source_text": "hello world"}],
            target_languages=["ru"],
        )
        assert result["batch_size_adjusted"] >= 1

    def test_multiple_rows(self):
        """Multiple rows produce batch estimate."""
        rows = [{"source_text": f"row {i}"} for i in range(5)]
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru", "de"],
            max_output_tokens=32000,
        )
        assert result["batch_size_adjusted"] >= 1

    def test_with_context_columns(self):
        """Context columns add to token estimate."""
        rows = [{"source_text": "hello", "description": "a long description here"}]
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru"],
            source_column="source_text",
            context_columns=["description"],
        )
        assert result["batch_size_adjusted"] >= 1

    def test_with_dictionary_entries(self):
        """Dictionary entries add tokens."""
        rows = [{"source_text": "hello"}]
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru"],
            dictionary_entries=[{"id": 1}, {"id": 2}],
        )
        assert result["batch_size_adjusted"] >= 1

    def test_with_provider_info(self):
        """Provider info resolves context window."""
        rows = [{"source_text": "hello"}]
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru"],
            provider_info="gpt-4o-mini",
        )
        assert result["batch_size_adjusted"] >= 1

    def test_explicit_context_window(self):
        """Explicit context_window overrides provider default."""
        rows = [{"source_text": "hello"}]
        result = estimate_token_budget(
            source_rows=rows,
            target_languages=["ru"],
            context_window=128000,
            max_output_tokens=4096,
        )
        assert result["batch_size_adjusted"] >= 1
        assert result["available_input_budget"] == 128000 - 4096

    def test_large_batch_reduced(self):
        """Batch size reduced when too large."""
        many_rows = [{"source_text": "x" * 5000} for _ in range(100)]
        result = estimate_token_budget(
            source_rows=many_rows,
            target_languages=["ru", "de", "fr"],
            batch_size=50,
        )
        assert result["batch_size_adjusted"] >= 1

    def test_no_target_languages_defaults(self):
        """Empty target_languages defaults to ['en']."""
        result = estimate_token_budget(
            source_rows=[{"source_text": "hello"}],
            target_languages=[],
        )
        assert result["batch_size_adjusted"] >= 1

    def test_provider_not_found_fallback(self):
        """Unknown provider falls back to defaults."""
        result = estimate_token_budget(
            source_rows=[{"source_text": "hello"}],
            target_languages=["ru"],
            provider_info="unknown-model-42",
        )
        assert result["batch_size_adjusted"] >= 1