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("langsmith")
class LangSmithLogger(Logger):
    """Export traces to LangSmith for observability."""

    name = "langsmith"

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._client: Any = None
        self._suite_run: Any = None
        self._case_run: Any = None
        self._project_name = kwargs.get("project", "leshy-security-scans")

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

    async def on_suite_start(self, config: TestSuiteConfig) -> None:
        try:
            self._client = self._get_client()
            self._suite_run = self._client.create_run(
                name=f"leshy-scan-{config.name}",
                run_type="chain",
                inputs={
                    "suite_name": config.name,
                    "target_type": config.target.type,
                    "target_url": config.target.url,
                    "test_count": len(config.tests),
                },
                tags=["leshy", "security-scan"],
                project_name=self._project_name,
            )
        except ImportError:
            raise
        except Exception as e:
            log.warning("Failed to initialize LangSmith run: %s", e)
            self._client = None

    async def on_case_start(self, case: TestCaseConfig) -> None:
        if self._client is None or self._suite_run is None:
            return
        try:
            self._case_run = self._client.create_run(
                name=f"test-case-{case.name}",
                run_type="chain",
                inputs={
                    "case_name": case.name,
                    "multi_turn": case.multi_turn,
                    "max_turns": case.max_turns,
                    "attacks": [a.name for a in case.attacks],
                    "tags": case.tags,
                },
                parent_run_id=self._suite_run.id,
                project_name=self._project_name,
            )
        except Exception as e:
            log.debug("Failed to create LangSmith case run: %s", e)

    async def on_payload_result(
        self,
        payload: AttackPayload,
        response: TargetResponse,
        results: list[EvalResult],
    ) -> None:
        if self._client is None or self._case_run is None:
            return
        try:
            verdict = results[0].verdict.value if results else "unknown"
            payload_run = self._client.create_run(
                name=f"payload-{payload.id}",
                run_type="llm",
                inputs={"payload": payload.content},
                parent_run_id=self._case_run.id,
                project_name=self._project_name,
                extra={
                    "metadata": {
                        "attack_name": payload.attack_name,
                        "category": payload.category,
                        "turn": payload.turn,
                    }
                },
            )
            self._client.update_run(
                payload_run.id,
                outputs={
                    "response": response.content,
                    "verdict": verdict,
                    "confidence": results[0].confidence if results else 0.0,
                    "reasoning": results[0].reasoning if results else "",
                },
                extra={
                    "metadata": {
                        "latency_ms": response.latency_ms,
                        "status_code": response.status_code,
                    }
                },
                end_time=payload_run.end_time,
            )
        except Exception as e:
            log.debug("Failed to log payload to LangSmith: %s", e)

    async def on_case_complete(self, result: TestCaseResult) -> None:
        if self._client is None or self._case_run is None:
            return
        try:
            self._client.update_run(
                self._case_run.id,
                outputs={
                    "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,
                },
                end_time=self._case_run.end_time,
            )
        except Exception as e:
            log.debug("Failed to end LangSmith case run: %s", e)
        self._case_run = None

    async def on_suite_complete(self, result: TestSuiteResult) -> None:
        if self._client is None or self._suite_run is None:
            return
        try:
            self._client.update_run(
                self._suite_run.id,
                outputs={
                    "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,
                },
                end_time=self._suite_run.end_time,
            )
        except Exception as e:
            log.debug("Failed to end LangSmith suite run: %s", e)
        self._suite_run = None
