"""
Base types used throughout the Risicare SDK.

These types form the foundation of the data model. All other types
build upon these primitives.

Design Principles:
1. Immutability: All core types are frozen dataclasses
2. Validation: Explicit validation functions for IDs
3. Serialization: JSON-compatible with custom encoders
4. Extensibility: Metadata dictionaries for custom attributes
5. W3C Compatibility: Trace/Span IDs follow W3C Trace Context format
"""

from __future__ import annotations

import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Dict, Optional


# =============================================================================
# Type Aliases
# =============================================================================

TraceID = str    # 32-character hex string (W3C compatible)
SpanID = str     # 16-character hex string (W3C compatible)
SessionID = str  # UUID string
AgentID = str    # User-defined or auto-generated
Timestamp = datetime
Metadata = Dict[str, Any]


# =============================================================================
# ID Generation
# =============================================================================

def generate_trace_id() -> TraceID:
    """
    Generate a W3C Trace Context compatible trace ID.

    Format: 32 lowercase hex characters (16 bytes)
    Example: "4bf92f3577b34da6a3ce929d0e0e4736"

    Returns:
        A unique trace ID string.
    """
    return uuid.uuid4().hex


def generate_span_id() -> SpanID:
    """
    Generate a W3C Trace Context compatible span ID.

    Format: 16 lowercase hex characters (8 bytes)
    Example: "00f067aa0ba902b7"

    Returns:
        A unique span ID string.
    """
    return uuid.uuid4().hex[:16]


def generate_session_id() -> SessionID:
    """
    Generate a session ID.

    Format: UUID string with hyphens
    Example: "550e8400-e29b-41d4-a716-446655440000"

    Returns:
        A unique session ID string.
    """
    return str(uuid.uuid4())


def generate_agent_id(prefix: str = "agent") -> AgentID:
    """
    Generate an agent ID with optional prefix.

    Format: "{prefix}_{short_uuid}"
    Example: "agent_a1b2c3d4"

    Args:
        prefix: Optional prefix for the agent ID.

    Returns:
        A unique agent ID string.
    """
    return f"{prefix}_{uuid.uuid4().hex[:8]}"


def utc_now() -> Timestamp:
    """
    Get current UTC timestamp.

    Returns:
        Current datetime in UTC timezone.
    """
    return datetime.now(timezone.utc)


# =============================================================================
# Validation
# =============================================================================

def validate_trace_id(trace_id: str) -> bool:
    """
    Validate a trace ID format.

    W3C Trace Context specifies trace IDs as 32-character lowercase hex strings.

    Args:
        trace_id: The trace ID to validate.

    Returns:
        True if valid, False otherwise.
    """
    if not isinstance(trace_id, str):
        return False
    if len(trace_id) != 32:
        return False
    # Must be lowercase hex (W3C Trace Context requirement)
    if trace_id != trace_id.lower():
        return False
    try:
        int(trace_id, 16)
        return True
    except ValueError:
        return False


def validate_span_id(span_id: str) -> bool:
    """
    Validate a span ID format.

    W3C Trace Context specifies span IDs as 16-character lowercase hex strings.

    Args:
        span_id: The span ID to validate.

    Returns:
        True if valid, False otherwise.
    """
    if not isinstance(span_id, str):
        return False
    if len(span_id) != 16:
        return False
    # Must be lowercase hex (W3C Trace Context requirement)
    if span_id != span_id.lower():
        return False
    try:
        int(span_id, 16)
        return True
    except ValueError:
        return False


# =============================================================================
# Enumerations
# =============================================================================

class SpanKind(Enum):
    """
    The kind of span, indicating its role in the trace.

    Based on OpenTelemetry SpanKind with agent-specific extensions.
    """
    # Standard kinds
    INTERNAL = "internal"      # Default, internal operation
    SERVER = "server"          # Server-side of RPC
    CLIENT = "client"          # Client-side of RPC
    PRODUCER = "producer"      # Message producer
    CONSUMER = "consumer"      # Message consumer

    # Agent-specific kinds
    AGENT = "agent"            # Agent lifecycle span
    LLM_CALL = "llm_call"      # LLM API call
    TOOL_CALL = "tool_call"    # Tool/function execution
    RETRIEVAL = "retrieval"    # RAG retrieval operation
    DECISION = "decision"      # Agent decision point (generic)
    DELEGATION = "delegation"  # Delegation to another agent
    COORDINATION = "coordination"  # Multi-agent coordination
    MESSAGE = "message"        # Inter-agent message

    # Phase-specific kinds (distinguish THINK vs DECIDE vs OBSERVE in queries)
    THINK = "think"            # Reasoning/planning phase
    DECIDE = "decide"          # Decision-making phase
    OBSERVE = "observe"        # Environment observation phase
    REFLECT = "reflect"        # Self-evaluation phase


class SpanStatus(Enum):
    """The status of a span."""
    UNSET = "unset"    # Status not set
    OK = "ok"          # Operation completed successfully
    ERROR = "error"    # Operation failed


class TraceFlags(Enum):
    """W3C Trace Context flags."""
    SAMPLED = 0x01     # Trace is sampled
    RANDOM = 0x02      # Trace ID is random


class AgentRole(Enum):
    """The role of an agent in a multi-agent system."""
    # Primary roles
    ORCHESTRATOR = "orchestrator"  # Coordinates other agents
    WORKER = "worker"              # Executes assigned tasks
    SUPERVISOR = "supervisor"      # Monitors and validates
    SPECIALIST = "specialist"      # Domain expert

    # Communication roles
    ROUTER = "router"              # Routes messages/tasks
    AGGREGATOR = "aggregator"      # Aggregates results
    BROADCASTER = "broadcaster"    # Broadcasts to multiple agents

    # Specialized roles
    CRITIC = "critic"              # Reviews and critiques
    PLANNER = "planner"            # Creates execution plans
    EXECUTOR = "executor"          # Executes plans
    RETRIEVER = "retriever"        # Retrieves information
    VALIDATOR = "validator"        # Validates outputs


class MessageType(Enum):
    """Types of messages in agent communication."""
    # Control messages
    TASK = "task"                  # Task assignment
    RESULT = "result"              # Task result
    STATUS = "status"              # Status update
    ERROR = "error"                # Error notification

    # Communication messages
    QUERY = "query"                # Information request
    RESPONSE = "response"          # Information response
    BROADCAST = "broadcast"        # Broadcast to all
    DIRECT = "direct"              # Direct message

    # Coordination messages
    PROPOSAL = "proposal"          # Propose action
    VOTE = "vote"                  # Vote on proposal
    CONSENSUS = "consensus"        # Consensus reached
    CONFLICT = "conflict"          # Conflict detected

    # Lifecycle messages
    HEARTBEAT = "heartbeat"        # Agent alive signal
    SHUTDOWN = "shutdown"          # Shutdown signal
    HANDOFF = "handoff"            # Handoff to another agent


class SemanticPhase(Enum):
    """
    Semantic phases of agent execution.

    Based on the Think-Decide-Act-Reflect pattern common in agent architectures.
    """
    THINK = "think"            # Reasoning, planning
    DECIDE = "decide"          # Decision making
    ACT = "act"                # Action execution
    REFLECT = "reflect"        # Self-evaluation
    OBSERVE = "observe"        # Environment observation
    COMMUNICATE = "communicate"  # Inter-agent communication
    COORDINATE = "coordinate"  # Multi-agent coordination


class EvaluationType(Enum):
    """Types of evaluations that can be performed."""
    # Automated evaluations
    LLM_JUDGE = "llm_judge"        # LLM-based evaluation
    RULE_BASED = "rule_based"      # Rule-based checks
    SIMILARITY = "similarity"      # Embedding similarity
    REGEX = "regex"                # Regex pattern matching

    # Human evaluations
    HUMAN_LABEL = "human_label"    # Human annotation
    HUMAN_RATING = "human_rating"  # Human rating
    HUMAN_PREFERENCE = "human_preference"  # A/B preference

    # Statistical evaluations
    STATISTICAL = "statistical"    # Statistical tests
    AB_TEST = "ab_test"            # A/B test results


class FixType(Enum):
    """Types of fixes that can be applied."""
    PROMPT = "prompt"              # Prompt modification
    PARAMETER = "parameter"        # Parameter adjustment
    TOOL = "tool"                  # Tool configuration
    RETRY = "retry"                # Retry with backoff
    FALLBACK = "fallback"          # Fallback to alternative
    GUARD = "guard"                # Add guardrail
    ROUTING = "routing"            # Change routing logic


# =============================================================================
# Data Classes
# =============================================================================

@dataclass(frozen=True)
class TimeRange:
    """A time range with start and optional end."""
    start: Timestamp
    end: Optional[Timestamp] = None

    @property
    def duration_ms(self) -> Optional[float]:
        """Calculate duration in milliseconds."""
        if self.end is None:
            return None
        delta = self.end - self.start
        return delta.total_seconds() * 1000

    @property
    def is_complete(self) -> bool:
        """Check if the time range has an end time."""
        return self.end is not None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "start": self.start.isoformat(),
            "end": self.end.isoformat() if self.end else None,
            "duration_ms": self.duration_ms,
        }
