from __future__ import annotations

import pytest

from leshy.core.models import (
    AttackPayload,
    EvalResult,
    TargetResponse,
    TestCaseConfig,
    TestCaseResult,
    TestSuiteResult,
    TokenUsage,
    Verdict,
)
from leshy.loggers.console import ConsoleLogger


def _make_payload(content: str = "test") -> AttackPayload:
    return AttackPayload(
        id="p1", content=content, attack_name="test.attack", category="test"
    )


def _make_response(content: str = "ok") -> TargetResponse:
    return TargetResponse(content=content)


def _make_eval(verdict: Verdict = Verdict.RESILIENT) -> EvalResult:
    return EvalResult(verdict=verdict, evaluator_name="test", payload_id="p1")


class TestProgressBar:
    @pytest.mark.asyncio
    async def test_progress_starts_in_non_verbose(self) -> None:
        logger = ConsoleLogger(verbose=False)
        case = TestCaseConfig(name="Test", attacks=[])
        await logger.on_case_start(case)
        assert logger._progress is not None
        assert logger._task_id is not None
        # Clean up
        await logger.on_case_complete(TestCaseResult(name="Test"))

    @pytest.mark.asyncio
    async def test_progress_not_started_in_verbose(self) -> None:
        logger = ConsoleLogger(verbose=True)
        case = TestCaseConfig(name="Test", attacks=[])
        await logger.on_case_start(case)
        assert logger._progress is None

    @pytest.mark.asyncio
    async def test_progress_count_increments(self) -> None:
        logger = ConsoleLogger(verbose=False)
        case = TestCaseConfig(name="Test", attacks=[])
        await logger.on_case_start(case)

        payload = _make_payload()
        response = _make_response()
        await logger.on_payload_result(payload, response, [_make_eval()])
        assert logger._payload_count == 1

        await logger.on_payload_result(payload, response, [_make_eval()])
        assert logger._payload_count == 2

        await logger.on_case_complete(TestCaseResult(name="Test"))
        assert logger._progress is None

    @pytest.mark.asyncio
    async def test_progress_stopped_on_case_complete(self) -> None:
        logger = ConsoleLogger(verbose=False)
        case = TestCaseConfig(name="Test", attacks=[])
        await logger.on_case_start(case)
        assert logger._progress is not None
        await logger.on_case_complete(TestCaseResult(name="Test"))
        assert logger._progress is None
        assert logger._task_id is None


class TestTokenUsageDisplay:
    @pytest.mark.asyncio
    async def test_token_usage_printed_when_nonzero(
        self, capsys: pytest.CaptureFixture[str],
    ) -> None:
        logger = ConsoleLogger(verbose=False)
        result = TestSuiteResult(
            suite_name="Test",
            token_usage=TokenUsage(
                prompt_tokens=100, completion_tokens=50, total_tokens=150
            ),
        )
        await logger.on_suite_complete(result)
        assert result.token_usage.total_tokens == 150

    @pytest.mark.asyncio
    async def test_no_crash_with_zero_tokens(self) -> None:
        logger = ConsoleLogger(verbose=False)
        result = TestSuiteResult(
            suite_name="Test",
            token_usage=TokenUsage(),
        )
        await logger.on_suite_complete(result)
