from __future__ import annotations

import logging
from typing import Any

from leshy.core.models import (
    AttackPayload,
    EvalResult,
    TargetResponse,
    TestCaseConfig,
    TestCaseResult,
    TestSuiteConfig,
    TestSuiteResult,
)
from leshy.core.registry import logger_registry
from leshy.loggers.base import Logger

log = logging.getLogger(__name__)


@logger_registry.register("langfuse")
class LangfuseLogger(Logger):
    """Export traces to Langfuse for observability."""

    name = "langfuse"

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._client: Any = None
        self._trace: Any = None
        self._current_span: Any = None

    def _get_client(self) -> Any:
        try:
            from langfuse import Langfuse
        except ImportError:
            raise ImportError(
                "The 'langfuse' package is required for Langfuse logging. "
                "Install it with: pip install 'leshy[langfuse]'"
            )
        return Langfuse()

    async def on_suite_start(self, config: TestSuiteConfig) -> None:
        try:
            self._client = self._get_client()
            self._trace = self._client.trace(
                name=f"leshy-scan-{config.name}",
                metadata={
                    "suite_name": config.name,
                    "target_type": config.target.type,
                    "target_url": config.target.url,
                    "test_count": len(config.tests),
                },
                tags=["leshy", "security-scan"],
            )
        except ImportError:
            raise
        except Exception as e:
            log.warning("Failed to initialize Langfuse trace: %s", e)
            self._client = None

    async def on_case_start(self, case: TestCaseConfig) -> None:
        if self._trace is None:
            return
        try:
            self._current_span = self._trace.span(
                name=f"test-case-{case.name}",
                metadata={
                    "multi_turn": case.multi_turn,
                    "max_turns": case.max_turns,
                    "attacks": [a.name for a in case.attacks],
                    "tags": case.tags,
                },
            )
        except Exception as e:
            log.debug("Failed to create Langfuse span: %s", e)

    async def on_payload_result(
        self,
        payload: AttackPayload,
        response: TargetResponse,
        results: list[EvalResult],
    ) -> None:
        if self._current_span is None:
            return
        try:
            verdict = results[0].verdict.value if results else "unknown"
            self._current_span.generation(
                name=f"payload-{payload.id}",
                model="target-agent",
                input=payload.content,
                output=response.content,
                metadata={
                    "attack_name": payload.attack_name,
                    "category": payload.category,
                    "turn": payload.turn,
                    "verdict": verdict,
                    "confidence": results[0].confidence if results else 0.0,
                    "reasoning": results[0].reasoning if results else "",
                    "latency_ms": response.latency_ms,
                    "status_code": response.status_code,
                },
            )
        except Exception as e:
            log.debug("Failed to log payload to Langfuse: %s", e)

    async def on_case_complete(self, result: TestCaseResult) -> None:
        if self._current_span is None:
            return
        try:
            self._current_span.update(
                output={
                    "total_payloads": result.total_payloads,
                    "vulnerable_count": result.vulnerable_count,
                    "resilient_count": result.resilient_count,
                    "inconclusive_count": result.inconclusive_count,
                    "error_count": result.error_count,
                    "duration_ms": result.duration_ms,
                },
            )
            self._current_span.end()
        except Exception as e:
            log.debug("Failed to end Langfuse span: %s", e)
        self._current_span = None

    async def on_suite_complete(self, result: TestSuiteResult) -> None:
        if self._trace is None:
            return
        try:
            self._trace.update(
                output={
                    "total_payloads": result.total_payloads,
                    "total_vulnerable": result.total_vulnerable,
                    "total_resilient": result.total_resilient,
                    "total_inconclusive": result.total_inconclusive,
                    "total_errors": result.total_errors,
                },
            )
            self._client.flush()
        except Exception as e:
            log.debug("Failed to flush Langfuse: %s", e)
