from __future__ import annotations

import json
import textwrap
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
import respx

from leshy.core.models import TokenUsage
from leshy.core.tokens import TokenTracker, token_tracker


class TestTokenTracker:
    def test_initial_state_is_zero(self) -> None:
        tracker = TokenTracker()
        usage = tracker.usage
        assert usage.prompt_tokens == 0
        assert usage.completion_tokens == 0
        assert usage.total_tokens == 0

    def test_add_accumulates(self) -> None:
        tracker = TokenTracker()
        tracker.add(100, 50)
        tracker.add(200, 75)
        usage = tracker.usage
        assert usage.prompt_tokens == 300
        assert usage.completion_tokens == 125
        assert usage.total_tokens == 425

    def test_reset_clears(self) -> None:
        tracker = TokenTracker()
        tracker.add(100, 50)
        tracker.reset()
        usage = tracker.usage
        assert usage.prompt_tokens == 0
        assert usage.completion_tokens == 0
        assert usage.total_tokens == 0

    def test_usage_returns_token_usage_model(self) -> None:
        tracker = TokenTracker()
        tracker.add(10, 20)
        usage = tracker.usage
        assert isinstance(usage, TokenUsage)


class TestTokenUsageModel:
    def test_defaults(self) -> None:
        usage = TokenUsage()
        assert usage.prompt_tokens == 0
        assert usage.completion_tokens == 0
        assert usage.total_tokens == 0

    def test_serialization(self) -> None:
        usage = TokenUsage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
        data = usage.model_dump()
        assert data == {
            "prompt_tokens": 100,
            "completion_tokens": 50,
            "total_tokens": 150,
        }


class TestRunnerTokenIntegration:
    """Test that runner resets/reads the token tracker."""

    def test_suite_result_includes_token_usage(
        self, tmp_path: Path
    ) -> None:
        import httpx
        import anyio

        from leshy.core.config import load_config
        from leshy.core.runner import Runner

        config_file = tmp_path / "config.yaml"
        config_file.write_text(textwrap.dedent("""\
            name: Token Test
            target:
              type: http
              url: https://token-test.example.com/chat
              body_template: '{"message": "{{payload}}"}'
              response_path: "$.response"
            loggers:
              - console
            tests:
              - name: Quick Test
                attacks:
                  - name: prompt_injection.direct
                    max_payloads: 1
        """))

        config = load_config(config_file)

        with respx.mock:
            respx.post("https://token-test.example.com/chat").respond(
                200,
                json={"response": "I cannot help with that."},
            )
            runner = Runner(config)
            result = anyio.run(runner.run)

        # Pattern evaluator doesn't use tokens, so should be zero
        assert result.token_usage.prompt_tokens == 0
        assert result.token_usage.completion_tokens == 0
        assert result.token_usage.total_tokens == 0

    def test_token_usage_in_json_output(self, tmp_path: Path) -> None:
        """Token usage should serialize into JSON results."""
        from leshy.core.models import TestSuiteResult

        result = TestSuiteResult(suite_name="test")
        result.token_usage = TokenUsage(
            prompt_tokens=500, completion_tokens=200, total_tokens=700
        )
        data = result.model_dump(mode="json")
        assert data["token_usage"]["prompt_tokens"] == 500
        assert data["token_usage"]["completion_tokens"] == 200
        assert data["token_usage"]["total_tokens"] == 700

    def test_runner_resets_tracker_between_runs(self, tmp_path: Path) -> None:
        import anyio

        from leshy.core.models import (
            TestSuiteConfig, TargetConfig, TestCaseConfig, AttackConfig,
        )
        from leshy.core.runner import Runner
        from leshy.core.tokens import token_tracker

        config = TestSuiteConfig(
            name="Reset Test",
            target=TargetConfig(
                type="http",
                url="https://reset-test.example.com/chat",
                body_template='{"message": "{{payload}}"}',
                response_path="$.response",
            ),
            loggers=["console"],
            tests=[
                TestCaseConfig(
                    name="Test",
                    attacks=[AttackConfig(name="prompt_injection.direct", max_payloads=1)],
                ),
            ],
        )

        # Seed the tracker with some tokens
        token_tracker.add(999, 999)

        with respx.mock:
            respx.post("https://reset-test.example.com/chat").respond(
                200, json={"response": "No."},
            )
            runner = Runner(config)
            result = anyio.run(runner.run)

        # Should be reset, not carry over the 999 + 999
        assert result.token_usage.prompt_tokens == 0
        assert result.token_usage.total_tokens == 0
