"""SAGE Audit Logger.

Provides structured audit logging for tracking LLM calls, RAG queries,
validation results, and security events.
"""

import logging
import uuid
from dataclasses import dataclass, field
from datetime import UTC, datetime
from enum import Enum
from typing import Any

from sage.logging.formatters import AuditFormatter
from sage.logging.handlers import RotatingAuditHandler


class AuditEventType(str, Enum):
    """Types of auditable events in SAGE."""

    # LLM Events
    LLM_REQUEST = "llm.request"
    LLM_RESPONSE = "llm.response"
    LLM_ERROR = "llm.error"

    # RAG Events
    RAG_QUERY = "rag.query"
    RAG_RETRIEVAL = "rag.retrieval"
    RAG_RESPONSE = "rag.response"

    # Validation Events
    VALIDATION_GROUNDING = "validation.grounding"
    VALIDATION_CITATION = "validation.citation"
    VALIDATION_SAFETY = "validation.safety"

    # Security Events
    AUTH_SUCCESS = "auth.success"
    AUTH_FAILURE = "auth.failure"
    ACCESS_DENIED = "access.denied"
    RATE_LIMIT = "rate.limit"

    # Configuration Events
    CONFIG_CHANGE = "config.change"
    CONFIG_LOAD = "config.load"

    # System Events
    SYSTEM_START = "system.start"
    SYSTEM_STOP = "system.stop"
    SYSTEM_ERROR = "system.error"

    # Data Events
    DOCUMENT_INGESTED = "document.ingested"
    DOCUMENT_DELETED = "document.deleted"
    INDEX_UPDATED = "index.updated"


@dataclass
class AuditEvent:
    """Structured audit event."""

    event_type: AuditEventType
    message: str
    data: dict[str, Any] = field(default_factory=dict)
    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: datetime = field(default_factory=lambda: datetime.now(UTC))
    user_id: str | None = None
    session_id: str | None = None
    request_id: str | None = None
    duration_ms: float | None = None
    success: bool = True
    error: str | None = None

    def to_dict(self) -> dict[str, Any]:
        """Convert event to dictionary for logging."""
        result: dict[str, Any] = {
            "event_type": self.event_type.value,
            "event_id": self.event_id,
            "timestamp": self.timestamp.isoformat(),
            "message": self.message,
            "success": self.success,
        }

        if self.user_id:
            result["user_id"] = self.user_id
        if self.session_id:
            result["session_id"] = self.session_id
        if self.request_id:
            result["request_id"] = self.request_id
        if self.duration_ms is not None:
            result["duration_ms"] = self.duration_ms
        if self.error:
            result["error"] = self.error
        if self.data:
            result["data"] = self.data

        return result


class AuditLogger:
    """Audit logger for SAGE operations.

    Provides structured logging of security-relevant events including
    LLM API calls, RAG queries, validation results, and configuration changes.

    Usage:
        audit = AuditLogger()

        # Log a RAG query
        audit.log_query(
            query="What is criticality safety?",
            sources=["DOC-001", "DOC-002"],
            response="Criticality safety is...",
        )

        # Log an LLM call
        audit.log_llm_call(
            model="gpt-4",
            input_tokens=500,
            output_tokens=200,
            latency_ms=1500,
        )
    """

    def __init__(
        self,
        log_path: str = "logs/audit.jsonl",
        max_size: str = "10MB",
        retention_days: int = 90,
        service_name: str = "sage",
        environment: str = "development",
    ):
        """Initialize audit logger.

        Args:
            log_path: Path to audit log file.
            max_size: Maximum file size before rotation.
            retention_days: Days to retain old log files.
            service_name: Service name for log entries.
            environment: Deployment environment.
        """
        self._logger = logging.getLogger("sage.audit")
        self._logger.setLevel(logging.INFO)
        self._logger.handlers.clear()
        self._logger.propagate = False

        # Configure handler with audit formatter
        handler = RotatingAuditHandler(
            filename=log_path,
            max_bytes=max_size,
            retention_days=retention_days,
        )
        handler.setFormatter(
            AuditFormatter(
                service_name=service_name,
                environment=environment,
            )
        )
        self._logger.addHandler(handler)

        self._service_name = service_name
        self._environment = environment

    def log_event(self, event: AuditEvent) -> None:
        """Log an audit event.

        Args:
            event: Audit event to log.
        """
        level = logging.INFO if event.success else logging.WARNING
        self._logger.log(
            level,
            event.message,
            extra={"audit_data": event.to_dict()},
        )

    def log_llm_call(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        latency_ms: float,
        provider: str = "openai",
        success: bool = True,
        error: str | None = None,
        request_id: str | None = None,
    ) -> None:
        """Log an LLM API call.

        Args:
            model: Model identifier.
            input_tokens: Number of input tokens.
            output_tokens: Number of output tokens.
            latency_ms: Request latency in milliseconds.
            provider: LLM provider name.
            success: Whether the call succeeded.
            error: Error message if call failed.
            request_id: Request identifier for correlation.
        """
        event = AuditEvent(
            event_type=AuditEventType.LLM_RESPONSE if success else AuditEventType.LLM_ERROR,
            message=f"LLM call to {provider}/{model}",
            data={
                "provider": provider,
                "model": model,
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "total_tokens": input_tokens + output_tokens,
            },
            duration_ms=latency_ms,
            success=success,
            error=error,
            request_id=request_id,
        )
        self.log_event(event)

    def log_query(
        self,
        query: str,
        sources: list[str],
        response: str,
        latency_ms: float | None = None,
        grounded: bool | None = None,
        citations_valid: bool | None = None,
        user_id: str | None = None,
        session_id: str | None = None,
    ) -> None:
        """Log a RAG query and response.

        Args:
            query: User query text.
            sources: List of source document IDs used.
            response: Generated response text.
            latency_ms: Total query latency in milliseconds.
            grounded: Whether response is grounded in sources.
            citations_valid: Whether citations are valid.
            user_id: User identifier.
            session_id: Session identifier.
        """
        event = AuditEvent(
            event_type=AuditEventType.RAG_RESPONSE,
            message="RAG query processed",
            data={
                "query": query[:500],  # Truncate long queries
                "query_length": len(query),
                "sources": sources,
                "source_count": len(sources),
                "response_length": len(response),
                "grounded": grounded,
                "citations_valid": citations_valid,
            },
            duration_ms=latency_ms,
            user_id=user_id,
            session_id=session_id,
        )
        self.log_event(event)

    def log_validation(
        self,
        validation_type: str,
        passed: bool,
        details: dict[str, Any] | None = None,
        request_id: str | None = None,
    ) -> None:
        """Log a validation result.

        Args:
            validation_type: Type of validation (grounding, citation, safety).
            passed: Whether validation passed.
            details: Additional validation details.
            request_id: Request identifier for correlation.
        """
        event_types = {
            "grounding": AuditEventType.VALIDATION_GROUNDING,
            "citation": AuditEventType.VALIDATION_CITATION,
            "safety": AuditEventType.VALIDATION_SAFETY,
        }

        event = AuditEvent(
            event_type=event_types.get(validation_type, AuditEventType.VALIDATION_GROUNDING),
            message=f"Validation {validation_type}: {'PASSED' if passed else 'FAILED'}",
            data={"validation_type": validation_type, "details": details or {}},
            success=passed,
            request_id=request_id,
        )
        self.log_event(event)

    def log_security_event(
        self,
        event_type: AuditEventType,
        message: str,
        user_id: str | None = None,
        ip_address: str | None = None,
        details: dict[str, Any] | None = None,
    ) -> None:
        """Log a security-related event.

        Args:
            event_type: Type of security event.
            message: Event description.
            user_id: User identifier.
            ip_address: Client IP address.
            details: Additional event details.
        """
        data = details or {}
        if ip_address:
            data["ip_address"] = ip_address

        event = AuditEvent(
            event_type=event_type,
            message=message,
            data=data,
            user_id=user_id,
            success=event_type not in {
                AuditEventType.AUTH_FAILURE,
                AuditEventType.ACCESS_DENIED,
                AuditEventType.RATE_LIMIT,
            },
        )
        self.log_event(event)

    def log_document_event(
        self,
        event_type: AuditEventType,
        document_id: str,
        document_type: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Log a document-related event.

        Args:
            event_type: Type of document event.
            document_id: Document identifier.
            document_type: Type of document.
            metadata: Additional document metadata.
        """
        event = AuditEvent(
            event_type=event_type,
            message=f"Document event: {event_type.value}",
            data={
                "document_id": document_id,
                "document_type": document_type,
                **(metadata or {}),
            },
        )
        self.log_event(event)

    def log_error(
        self,
        error: Exception,
        context: str | None = None,
        request_id: str | None = None,
    ) -> None:
        """Log an error event.

        Args:
            error: Exception that occurred.
            context: Additional context about where error occurred.
            request_id: Request identifier for correlation.
        """
        event = AuditEvent(
            event_type=AuditEventType.SYSTEM_ERROR,
            message=f"Error: {type(error).__name__}: {str(error)}",
            data={
                "error_type": type(error).__name__,
                "context": context,
            },
            success=False,
            error=str(error),
            request_id=request_id,
        )
        self.log_event(event)


# Global audit logger instance
_audit_logger: AuditLogger | None = None


def get_audit_logger(
    log_path: str = "logs/audit.jsonl",
    **kwargs: Any,
) -> AuditLogger:
    """Get or create the global audit logger.

    Args:
        log_path: Path to audit log file.
        **kwargs: Additional arguments for AuditLogger.

    Returns:
        AuditLogger: Global audit logger instance.
    """
    global _audit_logger
    if _audit_logger is None:
        _audit_logger = AuditLogger(log_path=log_path, **kwargs)
    return _audit_logger
