"""Tests for Feature 18: Token Usage & Cost Display.

Tests the cost calculation logic, cost table model selection,
and context window percentage thresholds.
"""

import pytest


# ---------------------------------------------------------------------------
# Cost table and model selection
# ---------------------------------------------------------------------------

class TestCostTable:
    """Test cost-per-1K-token pricing by model family."""

    COST_TABLE = {
        "haiku": {"input": 0.001, "output": 0.005},
        "sonnet": {"input": 0.003, "output": 0.015},
        "opus": {"input": 0.015, "output": 0.075},
    }

    def _select_cost_key(self, model_name: str) -> str:
        model_lower = model_name.lower()
        for fam in ("haiku", "sonnet", "opus"):
            if fam in model_lower:
                return fam
        return "sonnet"  # default

    def test_selects_haiku(self):
        assert self._select_cost_key("claude-haiku-4-5") == "haiku"
        assert self._select_cost_key("HAIKU-latest") == "haiku"

    def test_selects_sonnet(self):
        assert self._select_cost_key("claude-sonnet-4-6") == "sonnet"
        assert self._select_cost_key("sonnet") == "sonnet"

    def test_selects_opus(self):
        assert self._select_cost_key("claude-opus-4-6") == "opus"
        assert self._select_cost_key("Opus-Next") == "opus"

    def test_unknown_model_defaults_to_sonnet(self):
        assert self._select_cost_key("gpt-4") == "sonnet"
        assert self._select_cost_key("unknown-model") == "sonnet"
        assert self._select_cost_key("") == "sonnet"


# ---------------------------------------------------------------------------
# Cost estimation
# ---------------------------------------------------------------------------

class TestCostEstimation:
    """Test cost calculation: (input_tokens/1000 * rate) + (output_tokens/1000 * rate)."""

    COST_TABLE = {
        "haiku": {"input": 0.001, "output": 0.005},
        "sonnet": {"input": 0.003, "output": 0.015},
        "opus": {"input": 0.015, "output": 0.075},
    }

    def _calc_cost(self, input_tokens: int, output_tokens: int, model: str) -> float:
        costs = self.COST_TABLE[model]
        return input_tokens / 1000 * costs["input"] + output_tokens / 1000 * costs["output"]

    def test_zero_tokens_zero_cost(self):
        assert self._calc_cost(0, 0, "sonnet") == 0.0

    def test_haiku_cost(self):
        # 10K input + 5K output: 10*0.001 + 5*0.005 = 0.01 + 0.025 = 0.035
        cost = self._calc_cost(10_000, 5_000, "haiku")
        assert abs(cost - 0.035) < 1e-6

    def test_sonnet_cost(self):
        # 10K input + 5K output: 10*0.003 + 5*0.015 = 0.03 + 0.075 = 0.105
        cost = self._calc_cost(10_000, 5_000, "sonnet")
        assert abs(cost - 0.105) < 1e-6

    def test_opus_cost(self):
        # 10K input + 5K output: 10*0.015 + 5*0.075 = 0.15 + 0.375 = 0.525
        cost = self._calc_cost(10_000, 5_000, "opus")
        assert abs(cost - 0.525) < 1e-6

    def test_large_session_cost(self):
        # 100K input + 50K output with opus
        cost = self._calc_cost(100_000, 50_000, "opus")
        expected = 100 * 0.015 + 50 * 0.075  # 1.5 + 3.75 = 5.25
        assert abs(cost - expected) < 1e-6


# ---------------------------------------------------------------------------
# Context window % thresholds
# ---------------------------------------------------------------------------

class TestContextWindowPercentage:
    """Test context window percentage calculation and warning thresholds."""

    CONTEXT_WINDOW = 200_000

    def _ctx_pct(self, total_tokens: int) -> float:
        return min(100.0, total_tokens / self.CONTEXT_WINDOW * 100)

    def test_zero_tokens(self):
        assert self._ctx_pct(0) == 0.0

    def test_half_context(self):
        assert self._ctx_pct(100_000) == 50.0

    def test_full_context(self):
        assert self._ctx_pct(200_000) == 100.0

    def test_over_context_capped_at_100(self):
        assert self._ctx_pct(300_000) == 100.0

    def test_below_40_no_display(self):
        """Below 40% the ctx indicator should not be shown."""
        pct = self._ctx_pct(50_000)  # 25%
        assert pct < 40

    def test_40_to_79_shows_indicator(self):
        """Between 40-79% shows ctx indicator but no warning."""
        pct = self._ctx_pct(100_000)  # 50%
        assert 40 <= pct < 80

    def test_80_plus_shows_warning(self):
        """At 80%+ shows yellow warning to consider /compress."""
        pct = self._ctx_pct(160_000)  # 80%
        assert pct >= 80

    def test_exact_threshold_80(self):
        pct = self._ctx_pct(160_000)
        assert pct == 80.0


# ---------------------------------------------------------------------------
# Turn summary formatting
# ---------------------------------------------------------------------------

class TestTurnSummaryParts:
    """Test the parts array assembly for the turn summary line."""

    def _build_parts(
        self,
        turn_count: int,
        tool_count: int,
        session_tokens: dict,
        model: str = "sonnet",
        synthesis_reason: str = "",
    ) -> tuple[list[str], float]:
        COST_TABLE = {
            "haiku": {"input": 0.001, "output": 0.005},
            "sonnet": {"input": 0.003, "output": 0.015},
            "opus": {"input": 0.015, "output": 0.075},
        }
        costs = COST_TABLE[model]
        ctx_pct = 0.0
        parts = [f"{turn_count} turns", f"{tool_count} tools"]
        if session_tokens["total"] > 0:
            parts.append(f"{session_tokens['total']:,} tokens")
            est_cost = (
                session_tokens["input"] / 1000 * costs["input"]
                + session_tokens["output"] / 1000 * costs["output"]
            )
            if est_cost >= 0.01:
                parts.append(f"${est_cost:.2f}")
            elif est_cost > 0:
                parts.append(f"${est_cost:.4f}")
            ctx_pct = min(100.0, session_tokens["total"] / 200_000 * 100)
            if ctx_pct >= 40:
                parts.append(f"ctx {ctx_pct:.0f}%")
        if synthesis_reason:
            parts.append(f"synthesized ({synthesis_reason})")
        return parts, ctx_pct

    def test_basic_parts(self):
        parts, _ = self._build_parts(3, 5, {"input": 0, "output": 0, "total": 0})
        assert parts == ["3 turns", "5 tools"]

    def test_with_tokens_and_cost(self):
        tokens = {"input": 10_000, "output": 5_000, "total": 15_000}
        parts, ctx_pct = self._build_parts(2, 3, tokens, model="sonnet")
        assert "15,000 tokens" in parts
        assert "$0.10" in parts  # 10*0.003 + 5*0.015 = 0.105 → "$0.10" (f:.2f)
        assert ctx_pct < 40  # 7.5%, no ctx indicator

    def test_context_warning_at_80_percent(self):
        tokens = {"input": 100_000, "output": 80_000, "total": 180_000}
        parts, ctx_pct = self._build_parts(5, 10, tokens, model="haiku")
        assert ctx_pct >= 80
        assert any("ctx" in p for p in parts)

    def test_very_small_cost_uses_4_decimals(self):
        tokens = {"input": 100, "output": 50, "total": 150}
        parts, _ = self._build_parts(1, 1, tokens, model="haiku")
        # 0.1*0.001 + 0.05*0.005 = 0.0001 + 0.00025 = 0.00035
        cost_parts = [p for p in parts if p.startswith("$")]
        assert len(cost_parts) == 1
        assert cost_parts[0].startswith("$0.000")

    def test_synthesis_reason_appended(self):
        tokens = {"input": 1000, "output": 500, "total": 1500}
        parts, _ = self._build_parts(3, 2, tokens, synthesis_reason="budget_exhausted")
        assert "synthesized (budget_exhausted)" in parts
