"""Tests for the universal --budget N token-cap flag."""

from __future__ import annotations

import json
import os
import sys
from pathlib import Path

import pytest

sys.path.insert(0, str(Path(__file__).parent))
from conftest import (
    git_init,
    index_in_process,
    invoke_cli,
)

# ---------------------------------------------------------------------------
# Unit tests for budget_truncate()
# ---------------------------------------------------------------------------


class TestBudgetTruncate:
    """Unit tests for the budget_truncate() plain-text function."""

    def test_no_budget_returns_unchanged(self):
        """budget=0 returns text unchanged."""
        from roam.output.formatter import budget_truncate

        text = "Hello world\nSecond line\n"
        assert budget_truncate(text, 0) == text

    def test_negative_budget_returns_unchanged(self):
        """Negative budget returns text unchanged."""
        from roam.output.formatter import budget_truncate

        text = "Hello world\n"
        assert budget_truncate(text, -1) == text

    def test_short_text_within_budget(self):
        """Text shorter than budget is returned unchanged."""
        from roam.output.formatter import budget_truncate

        text = "Short text"
        # 10 chars / 4 = 2.5 tokens -> budget=100 is plenty
        assert budget_truncate(text, 100) == text

    def test_long_text_truncated(self):
        """Text longer than budget is truncated with metadata."""
        from roam.output.formatter import budget_truncate

        # Create text that is ~200 tokens (800 chars)
        text = "x" * 800
        result = budget_truncate(text, 50)  # budget = 50 tokens = 200 chars

        assert len(result) < len(text)
        assert "truncated" in result
        assert "budget: 50 tokens" in result
        assert "full output: ~200 tokens" in result

    def test_truncation_at_line_boundary(self):
        """Truncation prefers breaking at a line boundary."""
        from roam.output.formatter import budget_truncate

        # 10 lines of 30 chars each = 300 chars = ~75 tokens
        lines = ["A" * 28 + "\n" for _ in range(10)]
        text = "".join(lines)
        result = budget_truncate(text, 25)  # 25 tokens = 100 chars

        # Should end cleanly (last char before truncation notice should be \n or A)
        before_notice = result.split("\n\n... truncated")[0]
        assert before_notice.endswith("\n") or before_notice.endswith("A")

    def test_truncation_metadata_format(self):
        """Truncation notice includes budget and full output tokens."""
        from roam.output.formatter import budget_truncate

        text = "word " * 500  # 2500 chars = ~625 tokens
        result = budget_truncate(text, 100)

        assert "budget: 100 tokens" in result
        assert "full output: ~625 tokens" in result


# ---------------------------------------------------------------------------
# Unit tests for budget_truncate_json()
# ---------------------------------------------------------------------------


class TestBudgetTruncateJson:
    """Unit tests for the budget_truncate_json() function."""

    def test_no_budget_returns_unchanged(self):
        """budget=0 returns data unchanged."""
        from roam.output.formatter import budget_truncate_json

        data = {"command": "test", "summary": {"verdict": "ok"}, "items": [1, 2, 3]}
        result = budget_truncate_json(data, 0)
        assert result == data

    def test_small_data_within_budget(self):
        """Data that fits within budget is returned unchanged."""
        from roam.output.formatter import budget_truncate_json

        data = {"command": "test", "summary": {"verdict": "ok"}}
        result = budget_truncate_json(data, 10000)
        assert result == data

    def test_preserves_envelope_fields(self):
        """Envelope fields (command, summary, schema, etc.) are always preserved."""
        from roam.output.formatter import budget_truncate_json

        data = {
            "command": "health",
            "summary": {"verdict": "healthy", "score": 90},
            "schema": "roam-envelope-v1",
            "schema_version": "1.0.0",
            "version": "11.0.0",
            "project": "test",
            "_meta": {"timestamp": "2026-01-01T00:00:00Z"},
            # Large payload that forces truncation
            "items": [{"name": f"item_{i}", "data": "x" * 100} for i in range(100)],
        }

        # Very tight budget: ~50 tokens = 200 chars
        result = budget_truncate_json(data, 50)

        assert result["command"] == "health"
        assert "verdict" in result["summary"]
        assert result["summary"]["verdict"] == "healthy"

    def test_truncates_list_fields(self):
        """List fields in payload are truncated to fit budget."""
        from roam.output.formatter import budget_truncate_json

        big_list = [{"name": f"item_{i}", "value": "x" * 50} for i in range(50)]
        data = {
            "command": "test",
            "summary": {"verdict": "ok"},
            "items": big_list,
        }

        # Budget that can hold envelope + a few items but not all 50
        result = budget_truncate_json(data, 200)

        assert "items" in result
        assert len(result["items"]) < 50

    def test_adds_truncation_metadata(self):
        """Truncated output includes truncation metadata in summary."""
        from roam.output.formatter import budget_truncate_json

        data = {
            "command": "test",
            "summary": {"verdict": "ok"},
            "big_list": [{"name": f"item_{i}", "data": "x" * 100} for i in range(100)],
        }

        result = budget_truncate_json(data, 100)

        assert result["summary"]["truncated"] is True
        assert result["summary"]["budget_tokens"] == 100
        assert "full_output_tokens" in result["summary"]

    def test_does_not_mutate_original(self):
        """Original data dict is not mutated."""
        from roam.output.formatter import budget_truncate_json

        data = {
            "command": "test",
            "summary": {"verdict": "ok"},
            "items": [{"name": f"item_{i}"} for i in range(50)],
        }
        original_len = len(data["items"])

        budget_truncate_json(data, 50)

        assert len(data["items"]) == original_len
        assert "truncated" not in data["summary"]

    def test_drops_non_preserved_keys_when_very_tight(self):
        """With very tight budget, non-preserved keys are dropped entirely."""
        from roam.output.formatter import budget_truncate_json

        data = {
            "command": "test",
            "summary": {"verdict": "ok"},
            "huge_payload_a": [{"data": "x" * 200} for _ in range(50)],
            "huge_payload_b": [{"data": "y" * 200} for _ in range(50)],
        }

        # Very tight budget: should drop payload keys
        result = budget_truncate_json(data, 20)

        assert result["command"] == "test"
        assert result["summary"]["truncated"] is True


# ---------------------------------------------------------------------------
# Unit tests for estimate_tokens()
# ---------------------------------------------------------------------------


class TestEstimateTokens:
    """Unit tests for the estimate_tokens() helper."""

    def test_basic_estimate(self):
        """4 chars = 1 token."""
        from roam.output.formatter import estimate_tokens

        assert estimate_tokens("abcd") == 1

    def test_empty_string(self):
        """Empty string returns 1 (minimum)."""
        from roam.output.formatter import estimate_tokens

        assert estimate_tokens("") == 1

    def test_longer_text(self):
        """400 chars = 100 tokens."""
        from roam.output.formatter import estimate_tokens

        assert estimate_tokens("x" * 400) == 100


# ---------------------------------------------------------------------------
# Integration: json_envelope with budget
# ---------------------------------------------------------------------------


class TestJsonEnvelopeWithBudget:
    """Test that json_envelope() applies budget truncation."""

    def test_no_budget_default(self):
        """Default budget=0 does not truncate."""
        from roam.output.formatter import json_envelope

        result = json_envelope(
            "test",
            summary={"verdict": "ok"},
            items=[{"name": f"item_{i}"} for i in range(50)],
        )

        assert len(result.get("items", [])) == 50
        assert "truncated" not in result["summary"]

    def test_budget_truncates_envelope(self):
        """budget > 0 truncates large envelopes."""
        from roam.output.formatter import json_envelope

        result = json_envelope(
            "test",
            summary={"verdict": "ok"},
            budget=100,
            items=[{"name": f"item_{i}", "data": "x" * 100} for i in range(100)],
        )

        # Should have fewer items than the original 100
        if "items" in result:
            assert len(result["items"]) < 100
        assert result["summary"]["truncated"] is True
        assert result["summary"]["budget_tokens"] == 100

    def test_small_envelope_not_truncated(self):
        """Small envelope within budget is not modified."""
        from roam.output.formatter import json_envelope

        result = json_envelope(
            "test",
            summary={"verdict": "ok"},
            budget=10000,
            items=["a", "b", "c"],
        )

        assert result.get("items") == ["a", "b", "c"]
        assert "truncated" not in result["summary"]


# ---------------------------------------------------------------------------
# CLI integration: --budget flag accessible via ctx.obj
# ---------------------------------------------------------------------------


class TestBudgetCLIFlag:
    """Test that --budget is available as a CLI flag."""

    def test_budget_flag_default(self):
        """--budget defaults to 0 (unlimited)."""
        from click.testing import CliRunner

        from roam.cli import cli

        runner = CliRunner()
        # Use --help on a subcommand to verify the flag parses
        result = runner.invoke(cli, ["--budget", "0", "--help"])
        assert result.exit_code == 0

    def test_budget_flag_parses(self):
        """--budget N is accepted and stored in ctx.obj."""
        import click
        from click.testing import CliRunner

        captured = {}

        @click.command()
        @click.pass_context
        def check_budget(ctx):
            captured["budget"] = ctx.obj.get("budget", 0)
            click.echo("ok")

        from roam.cli import cli

        # Temporarily add our test command
        # Instead, just verify the flag is accepted by the CLI group
        result = runner = CliRunner()
        result = runner.invoke(cli, ["--budget", "500", "--help"])
        assert result.exit_code == 0

    def test_budget_flag_accepted(self):
        """--budget N is accepted by the CLI group without error."""
        from click.testing import CliRunner

        from roam.cli import cli

        runner = CliRunner()
        # Pass --budget with a valid value; --help to avoid running a command
        result = runner.invoke(cli, ["--budget", "500", "--help"])
        assert result.exit_code == 0
        # No error about unrecognized option
        assert "Error" not in result.output


# ---------------------------------------------------------------------------
# Integration: health command with --budget
# ---------------------------------------------------------------------------


@pytest.fixture
def health_project(tmp_path, monkeypatch):
    """Minimal indexed project for health command testing."""
    proj = tmp_path / "repo"
    proj.mkdir()
    (proj / ".gitignore").write_text(".roam/\n")

    src = proj / "src"
    src.mkdir()
    (src / "app.py").write_text("def main():\n    print('hello')\n\ndef helper():\n    return main()\n")
    (src / "utils.py").write_text("def format_name(name):\n    return name.title()\n")

    git_init(proj)
    monkeypatch.chdir(proj)
    out, rc = index_in_process(proj)
    assert rc == 0, f"index failed: {out}"

    return proj


class TestHealthWithBudget:
    """Test health command respects --budget flag."""

    def test_health_no_budget(self, health_project, monkeypatch):
        """Health without budget produces full output."""
        from click.testing import CliRunner

        monkeypatch.chdir(health_project)
        runner = CliRunner()
        result = invoke_cli(runner, ["--detail", "health"], cwd=health_project, json_mode=True)
        assert result.exit_code == 0
        data = json.loads(result.output)
        assert "truncated" not in data["summary"]

    def test_health_with_budget(self, health_project, monkeypatch):
        """Health with tight budget truncates output."""
        from click.testing import CliRunner

        from roam.cli import cli

        monkeypatch.chdir(health_project)
        runner = CliRunner()
        # Use a very small budget to force truncation
        old_cwd = os.getcwd()
        try:
            os.chdir(str(health_project))
            result = runner.invoke(
                cli,
                ["--json", "--budget", "50", "health"],
                catch_exceptions=False,
            )
        finally:
            os.chdir(old_cwd)

        assert result.exit_code == 0
        data = json.loads(result.output)
        assert data["command"] == "health"
        assert data["summary"]["truncated"] is True
        assert data["summary"]["budget_tokens"] == 50

    def test_health_large_budget_no_truncation(self, health_project, monkeypatch):
        """Health with large budget does not truncate."""
        from click.testing import CliRunner

        from roam.cli import cli

        monkeypatch.chdir(health_project)
        runner = CliRunner()
        old_cwd = os.getcwd()
        try:
            os.chdir(str(health_project))
            result = runner.invoke(
                cli,
                ["--json", "--budget", "100000", "--detail", "health"],
                catch_exceptions=False,
            )
        finally:
            os.chdir(old_cwd)

        assert result.exit_code == 0
        data = json.loads(result.output)
        assert "truncated" not in data["summary"]


# ---------------------------------------------------------------------------
# Integration: budget_truncate for text output
# ---------------------------------------------------------------------------


class TestBudgetTruncateTextIntegration:
    """Test budget_truncate with realistic command output."""

    def test_realistic_health_output(self):
        """Simulated health text output truncates correctly."""
        from roam.output.formatter import budget_truncate

        # Simulate a large health output
        lines = ["VERDICT: Healthy codebase (85/100)"]
        lines.append("")
        lines.append("=== Cycles ===")
        for i in range(20):
            lines.append(f"  cycle {i}: a -> b -> c -> a")
        lines.append("")
        lines.append("=== God Components ===")
        for i in range(30):
            lines.append(f"  {i}. BigClass (degree={50 + i})")

        text = "\n".join(lines)
        result = budget_truncate(text, 100)  # 100 tokens = 400 chars

        # Should still start with VERDICT
        assert result.startswith("VERDICT:")
        assert "truncated" in result


# ---------------------------------------------------------------------------
# MCP integration: _apply_budget helper
# ---------------------------------------------------------------------------


class TestMCPApplyBudget:
    """Test the MCP _apply_budget helper."""

    def test_apply_budget_zero(self):
        """budget=0 returns data unchanged."""
        from roam.mcp_server import _apply_budget

        data = {"command": "test", "summary": {"verdict": "ok"}}
        assert _apply_budget(data, 0) is data

    def test_apply_budget_truncates(self):
        """budget > 0 applies truncation on large data."""
        from roam.mcp_server import _apply_budget

        data = {
            "command": "test",
            "summary": {"verdict": "ok"},
            "large_list": [{"data": "x" * 200} for _ in range(100)],
        }

        result = _apply_budget(data, 100)
        assert result["summary"]["truncated"] is True
