"""
Starvex Observability Models - Data models for tracing and metrics
"""

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional
from datetime import datetime
import uuid


@dataclass
class TraceEvent:
    """Represents a single trace event in the observability pipeline"""

    trace_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    input_text: str = ""
    output_text: Optional[str] = None
    verdict: str = "PASSED"
    status: str = "success"  # success, blocked, flagged
    confidence_score: float = 1.0

    # User tracking
    user_id: Optional[str] = None
    session_id: Optional[str] = None

    # Timing
    latency_ms: float = 0.0
    processing_time_ms: float = 0.0
    guardrail_latency_ms: float = 0.0
    agent_latency_ms: float = 0.0

    # Token tracking
    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    model_name: Optional[str] = None

    # Block details
    blocked_reason: Optional[str] = None
    blocked_content: Optional[str] = None
    rules_triggered: List[Dict[str, Any]] = field(default_factory=list)

    # Metadata
    checks: List[Dict[str, Any]] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    environment: str = "production"
    sdk_version: str = ""
    timestamp: datetime = field(default_factory=datetime.utcnow)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for API submission"""
        return {
            "trace_id": self.trace_id,
            "input_text": self.input_text,
            "output_text": self.output_text,
            "verdict": self.verdict,
            "status": self.status,
            "confidence_score": self.confidence_score,
            "user_id_external": self.user_id,
            "session_id": self.session_id,
            "latency_ms": self.latency_ms,
            "processing_time_ms": self.processing_time_ms,
            "guardrail_latency_ms": self.guardrail_latency_ms,
            "agent_latency_ms": self.agent_latency_ms,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "total_tokens": self.total_tokens,
            "model_name": self.model_name,
            "blocked_reason": self.blocked_reason,
            "blocked_content": self.blocked_content,
            "rules_triggered": self.rules_triggered,
            "checks": self.checks,
            "metadata": self.metadata,
            "environment": self.environment,
            "sdk_version": self.sdk_version,
        }


@dataclass
class MetricsSnapshot:
    """Aggregated metrics snapshot for a time period"""

    total_calls: int = 0
    successful_calls: int = 0
    blocked_calls: int = 0
    flagged_calls: int = 0
    error_calls: int = 0

    # Latency stats
    avg_latency_ms: float = 0.0
    p50_latency_ms: float = 0.0
    p95_latency_ms: float = 0.0
    p99_latency_ms: float = 0.0

    # Token stats
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_tokens: int = 0

    # Provider breakdown
    by_provider: Dict[str, Dict[str, Any]] = field(default_factory=dict)

    # Time range
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None

    @property
    def success_rate(self) -> float:
        """Calculate success rate"""
        if self.total_calls == 0:
            return 0.0
        return self.successful_calls / self.total_calls

    @property
    def block_rate(self) -> float:
        """Calculate block rate"""
        if self.total_calls == 0:
            return 0.0
        return self.blocked_calls / self.total_calls

    @property
    def error_rate(self) -> float:
        """Calculate error rate"""
        if self.total_calls == 0:
            return 0.0
        return self.error_calls / self.total_calls

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "total_calls": self.total_calls,
            "successful_calls": self.successful_calls,
            "blocked_calls": self.blocked_calls,
            "flagged_calls": self.flagged_calls,
            "error_calls": self.error_calls,
            "success_rate": self.success_rate,
            "block_rate": self.block_rate,
            "error_rate": self.error_rate,
            "avg_latency_ms": self.avg_latency_ms,
            "p50_latency_ms": self.p50_latency_ms,
            "p95_latency_ms": self.p95_latency_ms,
            "p99_latency_ms": self.p99_latency_ms,
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "total_tokens": self.total_tokens,
            "by_provider": self.by_provider,
        }
