"""Tests for nsct.orchestration.context_budget – Stage 20."""

import pytest

from nsct.orchestration.context_budget import (
    ContextBudgetConfig,
    ContextBudgetExhaustedError,
    ContextBudgetTracker,
)

# ---------------------------------------------------------------------------
# Helper – build a tracker with given overrides; min-values from Field(ge/le)
# ---------------------------------------------------------------------------

_DEFAULTS = {
    "planner_max_tokens": 12_000,
    "claim_extraction_max_tokens": 16_000,
    "contradiction_max_tokens": 24_000,
    "synthesis_max_tokens": 48_000,
}


def _tracker(**overrides):
    """Return a fresh ContextBudgetTracker with optional budget overrides."""
    merged = {**_DEFAULTS, **overrides}
    return ContextBudgetTracker(ContextBudgetConfig(**merged))


# ---------------------------------------------------------------------------
# 1) Default-Konfiguration validiert sich selbst
# ---------------------------------------------------------------------------


class TestDefaultConfig:
    """ContextBudgetConfig self-validation."""

    def test_defaults_sane(self):
        cfg = ContextBudgetConfig()
        # Defaults lie inside their [ge, le] windows
        assert cfg.planner_max_tokens == 12_000
        assert cfg.claim_extraction_max_tokens == 16_000
        assert cfg.contradiction_max_tokens == 24_000
        assert cfg.synthesis_max_tokens == 48_000
        # total is the sum
        assert cfg.total_max_tokens == 12_000 + 16_000 + 24_000 + 48_000

    def test_frozen(self):
        """Config is immutable – direct assignment is blocked."""
        cfg = ContextBudgetConfig()
        # Pydantic v2 frozen allows model_copy but not __setattr__
        with pytest.raises(Exception):
            cfg.planner_max_tokens = 999

    def test_total_max_tokens(self):
        cfg = ContextBudgetConfig()
        assert cfg.total_max_tokens == 100_000

    def test_get_limit_valid_key(self):
        cfg = ContextBudgetConfig()
        assert cfg.get_limit("planner") == 12_000
        assert cfg.get_limit("synthesis") == 48_000

    def test_get_limit_invalid_key_raises_value_error(self):
        cfg = ContextBudgetConfig()
        with pytest.raises(ValueError, match="Unknown stage key"):
            cfg.get_limit("nonexistent")

    def test_stage_keys(self):
        cfg = ContextBudgetConfig()
        keys = cfg.stage_keys
        assert keys == [
            "planner_max_tokens",
            "claim_extraction_max_tokens",
            "contradiction_max_tokens",
            "synthesis_max_tokens",
        ]


# ---------------------------------------------------------------------------
# 2) track_tokens akkumuliert korrekt über mehrere calls hinweg
# ---------------------------------------------------------------------------


class TestTrackAccumulation:
    """Accumulation of tracked tokens across calls."""

    def test_accumulates(self):
        tr = _tracker()
        tr.track_tokens("planner", 100)
        tr.track_tokens("planner", 200)
        assert tr._usage["planner"] == 300

    def test_multiple_stages_independent(self):
        tr = _tracker()
        tr.track_tokens("planner", 50)
        tr.track_tokens("claim_extraction", 30)
        assert tr._usage["planner"] == 50
        assert tr._usage["claim_extraction"] == 30
        assert tr._usage["claim_extraction"] != tr._usage["planner"]

    def test_total_usage_property(self):
        tr = _tracker()
        tr.track_tokens("planner", 100)
        tr.track_tokens("synthesis", 200)
        assert tr.total_usage == 300

    def test_usage_starts_zero(self):
        tr = _tracker()
        assert tr._usage.get("planner", 0) == 0


# ---------------------------------------------------------------------------
# 3) ContextBudgetExhaustedError wird bei Überschreitung geworfen
# ---------------------------------------------------------------------------


class TestExhaustedError:
    """Raising ContextBudgetExhaustedError on budget overrun."""

    def test_exact_limit_succeeds(self):
        """Hitting the limit exactly should still be allowed."""
        tr = _tracker(planner_max_tokens=8_000)  # minimum allowed
        tr.track_tokens("planner", 8_000)
        assert tr._usage["planner"] == 8_000

    def test_over_limit_raises(self):
        tr = _tracker(planner_max_tokens=8_000)
        tr.track_tokens("planner", 4_000)
        with pytest.raises(ContextBudgetExhaustedError) as exc_info:
            tr.track_tokens("planner", 4_001)  # 4000 + 4001 = 8001 > 8000
        assert "planner" in str(exc_info.value)
        assert exc_info.value.used_tokens == 4_000
        assert exc_info.value.limit_tokens == 8_000

    def test_over_limit_error_not_accumulated(self):
        """The counter must NOT be incremented on an overrun."""
        tr = _tracker(planner_max_tokens=8_000)
        tr.track_tokens("planner", 3_000)
        with pytest.raises(ContextBudgetExhaustedError):
            tr.track_tokens("planner", 5_001)
        assert tr._usage["planner"] == 3_000  # unchanged

    def test_error_attributes(self):
        tr = _tracker(planner_max_tokens=8_000)
        with pytest.raises(ContextBudgetExhaustedError) as exc_info:
            tr.track_tokens("planner", 8_001)
        err = exc_info.value
        assert err.stage_name == "planner"
        assert err.used_tokens == 0
        assert err.limit_tokens == 8_000


# ---------------------------------------------------------------------------
# 4) get_usage() liefert korrekte Pct und Bytes
# ---------------------------------------------------------------------------


class TestGetUsage:
    """get_usage() returns correct pct and bytes."""

    def test_tracks_tokens_and_bytes(self):
        tr = _tracker()
        tr.track_tokens("planner", 100)
        usage = tr.get_usage()
        p = usage["planner"]
        assert p["usage_tokens"] == 100.0
        assert p["limit_tokens"] == 12_000.0
        # 100 / 12000 * 100 = 0.8333… %
        assert p["usage_pct"] == 0.83
        # 100 * 4.0 = 400.0 bytes
        assert p["usage_bytes"] == 400.0

    def test_unknown_stage_shows_zero(self):
        """Stages that have never been tracked still appear with zeros."""
        tr = _tracker()
        tr.track_tokens("planner", 100)
        usage = tr.get_usage()
        assert usage["synthesis"]["usage_tokens"] == 0.0
        assert usage["synthesis"]["usage_pct"] == 0.0
        assert usage["synthesis"]["usage_bytes"] == 0.0

    def test_usage_pct_rounded(self):
        tr = _tracker(planner_max_tokens=8_000)
        tr.track_tokens("planner", 1)
        usage = tr.get_usage()
        # 1/8000 * 100 = 0.0125 → rounded to 0.01
        assert usage["planner"]["usage_pct"] == 0.01


# ---------------------------------------------------------------------------
# 5) reset_stage und reset_all funktionieren
# ---------------------------------------------------------------------------


class TestReset:
    """reset_stage and reset_all work correctly."""

    def test_reset_stage(self):
        tr = _tracker()
        tr.track_tokens("planner", 100)
        tr.reset_stage("planner")
        assert tr._usage["planner"] == 0

    def test_reset_stage_unknown_ignored(self):
        """Resetting an unknown stage is a no-op, not an error."""
        tr = _tracker()
        tr.reset_stage("nonexistent")
        assert tr._usage == {}

    def test_reset_all(self):
        tr = _tracker()
        tr.track_tokens("planner", 100)
        tr.track_tokens("synthesis", 200)
        tr.reset_all()
        assert tr._usage == {}
        assert tr.total_usage == 0

    def test_after_reset_usage_reflects_zero(self):
        tr = _tracker()
        tr.track_tokens("planner", 100)
        tr.reset_all()
        usage = tr.get_usage()
        assert usage["planner"]["usage_tokens"] == 0.0
        assert usage["planner"]["usage_pct"] == 0.0


# ---------------------------------------------------------------------------
# 6) is_exhausted() erkennt exhaustion richtig
# ---------------------------------------------------------------------------


class TestIsExhausted:
    """is_exhausted() detects exhaustion correctly."""

    def test_not_exhausted_under_budget(self):
        tr = _tracker()
        tr.track_tokens("planner", 100)
        assert tr.is_exhausted("planner") is False

    def test_exhausted_at_limit(self):
        tr = _tracker(planner_max_tokens=8_000)
        tr.track_tokens("planner", 8_000)
        assert tr.is_exhausted("planner") is True

    def test_is_exhausted_any_stage(self):
        """No argument → checks all stages."""
        tr = _tracker()
        tr.track_tokens("synthesis", 48_000)  # at limit
        assert tr.is_exhausted() is True

    def test_is_exhausted_specific_false(self):
        tr = _tracker(planner_max_tokens=8_000, synthesis_max_tokens=32_000)
        tr.track_tokens("synthesis", 1_000)
        # synthesis is under limit
        assert tr.is_exhausted("synthesis") is False

    def test_all_stages_clean(self):
        tr = _tracker()
        tr.track_tokens("planner", 1)
        assert tr.is_exhausted() is False


# ---------------------------------------------------------------------------
# 7) Ungültige Stage-Namen werfen ValueError
# ---------------------------------------------------------------------------


class TestInvalidStageNames:
    """Invalid stage names raise ValueError."""

    def track_tokens_bad_stage(self):
        tr = _tracker()
        with pytest.raises(ValueError, match="Unknown stage"):
            tr.track_tokens("nonexistent", 100)

    def get_limit_bad_stage(self):
        cfg = ContextBudgetConfig()
        with pytest.raises(ValueError, match="Unknown stage key"):
            cfg.get_limit("nonexistent")


# ---------------------------------------------------------------------------
# 8) Token = 0 oder negativ wirft ValueError
# ---------------------------------------------------------------------------


class TestInvalidTokens:
    """Zero or negative token counts raise ValueError."""

    def test_zero_tokens_raises(self):
        tr = _tracker()
        with pytest.raises(ValueError, match="tokens must be positive"):
            tr.track_tokens("planner", 0)

    def test_negative_tokens_raises(self):
        tr = _tracker()
        with pytest.raises(ValueError, match="tokens must be positive"):
            tr.track_tokens("planner", -42)