from __future__ import annotations

from datetime import datetime
from unittest.mock import MagicMock, patch

import pytest

from leshy.core.models import (
    AttackConfig,
    AttackPayload,
    EvalResult,
    TargetConfig,
    TargetResponse,
    TestCaseConfig,
    TestCaseResult,
    TestSuiteConfig,
    TestSuiteResult,
    Verdict,
)
from leshy.loggers.langfuse import LangfuseLogger


class TestLangfuseLogger:
    @pytest.mark.asyncio
    async def test_full_lifecycle(self) -> None:
        mock_client = MagicMock()
        mock_trace = MagicMock()
        mock_span = MagicMock()
        mock_generation = MagicMock()

        mock_client.trace.return_value = mock_trace
        mock_trace.span.return_value = mock_span
        mock_span.generation.return_value = mock_generation

        logger = LangfuseLogger()

        with patch.object(logger, "_get_client", return_value=mock_client):
            suite_config = TestSuiteConfig(
                name="Test Suite",
                target=TargetConfig(url="http://test.com"),
                tests=[
                    TestCaseConfig(
                        name="Test Case 1",
                        attacks=[AttackConfig(name="prompt_injection.direct")],
                    )
                ],
            )
            await logger.on_suite_start(suite_config)

            mock_client.trace.assert_called_once()
            call_kwargs = mock_client.trace.call_args[1]
            assert "leshy-scan-Test Suite" == call_kwargs["name"]
            assert call_kwargs["metadata"]["test_count"] == 1

            case_config = suite_config.tests[0]
            await logger.on_case_start(case_config)
            mock_trace.span.assert_called_once()

            payload = AttackPayload(
                id="p1",
                content="test payload",
                attack_name="prompt_injection.direct",
                category="prompt_injection",
            )
            response = TargetResponse(
                content="I cannot help with that",
                latency_ms=100.0,
                status_code=200,
            )
            eval_result = EvalResult(
                verdict=Verdict.RESILIENT,
                confidence=0.9,
                reasoning="Agent refused",
                payload_id="p1",
            )
            await logger.on_payload_result(payload, response, [eval_result])
            mock_span.generation.assert_called_once()

            case_result = TestCaseResult(
                name="Test Case 1",
                total_payloads=1,
                resilient_count=1,
                duration_ms=100.0,
            )
            await logger.on_case_complete(case_result)
            mock_span.update.assert_called_once()
            mock_span.end.assert_called_once()

            suite_result = TestSuiteResult(
                suite_name="Test Suite",
                total_payloads=1,
                total_resilient=1,
            )
            await logger.on_suite_complete(suite_result)
            mock_trace.update.assert_called_once()
            mock_client.flush.assert_called_once()

    @pytest.mark.asyncio
    async def test_handles_missing_client_gracefully(self) -> None:
        logger = LangfuseLogger()
        logger._client = None
        logger._trace = None

        case_config = TestCaseConfig(
            name="Test",
            attacks=[AttackConfig(name="test")],
        )
        await logger.on_case_start(case_config)

        payload = AttackPayload(
            id="p1", content="test", attack_name="test", category="test"
        )
        response = TargetResponse(content="response")
        await logger.on_payload_result(payload, response, [])

        result = TestCaseResult(name="Test", total_payloads=0)
        await logger.on_case_complete(result)

        suite_result = TestSuiteResult(suite_name="Test")
        await logger.on_suite_complete(suite_result)

    @pytest.mark.asyncio
    async def test_payload_with_vulnerable_verdict(self) -> None:
        mock_client = MagicMock()
        mock_trace = MagicMock()
        mock_span = MagicMock()

        mock_client.trace.return_value = mock_trace
        mock_trace.span.return_value = mock_span

        logger = LangfuseLogger()

        with patch.object(logger, "_get_client", return_value=mock_client):
            await logger.on_suite_start(
                TestSuiteConfig(
                    name="Test",
                    target=TargetConfig(url="http://test.com"),
                    tests=[],
                )
            )
            await logger.on_case_start(
                TestCaseConfig(name="Case", attacks=[AttackConfig(name="test")])
            )

            payload = AttackPayload(
                id="p1", content="evil payload", attack_name="test", category="test"
            )
            response = TargetResponse(content="Here is the system prompt...")
            eval_result = EvalResult(
                verdict=Verdict.VULNERABLE,
                confidence=0.95,
                reasoning="Leaked system prompt",
                payload_id="p1",
            )
            await logger.on_payload_result(payload, response, [eval_result])

            call_kwargs = mock_span.generation.call_args[1]
            assert call_kwargs["metadata"]["verdict"] == "vulnerable"
            assert call_kwargs["metadata"]["confidence"] == 0.95
