"""
vgovern SDK - Type Definitions

This module defines the core data structures for the SDK.
All types are designed to be serializable and framework-agnostic.
"""

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


class StepType(str, Enum):
    """Enumeration of possible step types in agent execution."""
    LLM_CALL = "llm_call"
    TOOL_CALL = "tool_call"
    TOOL_RESULT = "tool_result"
    OBSERVATION = "observation"

    FINAL_RESPONSE = "final_response"
    UNKNOWN = "unknown"


class TraceStatus(str, Enum):
    """Enumeration of possible trace statuses."""
    RUNNING = "running"
    HALTED = "halted"
    COMPLETED = "completed"
    ERROR = "error"


class GovernanceDecision(str, Enum):
    """Governance decision returned by policy engine."""
    ALLOW = "ALLOW"
    HALT = "HALT"


@dataclass(frozen=True)
class GovernanceResponse:
    """Response from policy engine governance check."""
    decision: GovernanceDecision
    reason: Optional[str] = None
    policy_id: Optional[str] = None


@dataclass
class TokenUsage:
    """Token usage metrics for LLM calls."""
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0

    def __post_init__(self):
        if self.total_tokens == 0:
            object.__setattr__(self, "total_tokens", self.prompt_tokens + self.completion_tokens)


@dataclass
class StepData:
    """Captured data for a single execution step."""
    step_id: str
    trace_id: UUID
    timestamp: datetime
    step_type: StepType
    model: Optional[str] = None
    prompt: Optional[str] = None
    response: Optional[str] = None
    token_usage: Optional[TokenUsage] = None
    latency_ms: Optional[int] = None
    estimated_cost_usd: Optional[float] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class TraceMetadata:
    """Metadata for a trace execution."""
    agent_id: str
    session_id: Optional[str] = None
    environment: str = "production"
    user_id: Optional[str] = None
    tags: Dict[str, str] = field(default_factory=dict)


@dataclass
class TraceInfo:
    """Information about an active trace."""
    trace_id: UUID
    agent_id: str
    start_time: datetime
    status: TraceStatus = TraceStatus.RUNNING
    session_id: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = field(default_factory=dict)


@dataclass
class SDKConfig:
    """Configuration for the Governance SDK."""
    backend_url: str
    api_key: str
    agent_id: str
    environment: str = "production"
    timeout_ms: int = 1000
    max_retries: int = 2
    enable_policy_check: bool = True
    enable_async_reporting: bool = True
    redact_sensitive_data: bool = True
    fail_safe_mode: bool = True
    enable_console_trace: bool = True


class GovernanceHaltError(Exception):
    """
    Raised when the policy engine returns a HALT decision.
    
    This is the ONLY exception that should stop agent execution.
    All other SDK errors are handled silently in fail-safe mode.
    """

    def __init__(self, reason: str, policy_id: Optional[str] = None):
        self.reason = reason
        self.policy_id = policy_id
        super().__init__(f"Governance HALT: {reason}")

    def to_dict(self) -> Dict[str, Any]:
        return {
            "error_type": "GovernanceHaltError",
            "reason": self.reason,
            "policy_id": self.policy_id,
        }


@dataclass
class StepNode:
    """Tree node for execution step with parent-child relationships."""
    step_id: str
    trace_id: UUID
    name: str
    step_type: StepType
    parent_id: Optional[str] = None
    children: List['StepNode'] = field(default_factory=list)
    input_data: Optional[Dict[str, Any]] = None
    output_data: Optional[Dict[str, Any]] = None
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None
    duration_ms: Optional[int] = None
    policy_hits: Optional[List[Dict[str, Any]]] = None
    estimated_cost_usd: Optional[float] = None

    def add_child(self, child: 'StepNode') -> None:
        """Add a child step to this node."""
        self.children.append(child)

    def is_complete(self) -> bool:
        """Check if this step has been completed (has end_time)."""
        return self.end_time is not None

    def get_duration_display(self) -> str:
        """Get human-readable duration string."""
        if self.duration_ms is None:
            return "0.00s"
        seconds = self.duration_ms / 1000.0
        return f"{seconds:.2f}s"

    def to_dict(self) -> Dict[str, Any]:
        """Convert step node to dictionary representation."""
        return {
            "step_id": self.step_id,
            "trace_id": str(self.trace_id),
            "name": self.name,
            "step_type": self.step_type.value,
            "parent_id": self.parent_id,
            "children": [child.to_dict() for child in self.children],
            "input_data": self.input_data,
            "output_data": self.output_data,
            "start_time": self.start_time.isoformat() if self.start_time else None,
            "end_time": self.end_time.isoformat() if self.end_time else None,
            "duration_ms": self.duration_ms,
            "policy_hits": self.policy_hits,
        }

    @classmethod
    def from_step_data(cls, step_data: StepData, name: Optional[str] = None) -> 'StepNode':
        """Create StepNode from legacy StepData."""
        return cls(
            step_id=step_data.step_id,
            trace_id=step_data.trace_id,
            name=name or step_data.step_type.value,
            step_type=step_data.step_type,
            start_time=step_data.timestamp,
            end_time=step_data.timestamp,
            duration_ms=step_data.latency_ms,
            input_data={"prompt": step_data.prompt, "model": step_data.model} if step_data.prompt else None,
            output_data={"response": step_data.response} if step_data.response else None,
        )

@dataclass
class Trace:
    """Complete trace with execution tree structure."""
    trace_id: UUID
    start_time: datetime
    end_time: Optional[datetime] = None
    root_steps: List[StepNode] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    user_input: Optional[str] = None

    def add_root_step(self, step: StepNode) -> None:
        """Add a root-level step to the trace."""
        self.root_steps.append(step)

    def get_total_duration_ms(self) -> int:
        """Calculate total trace duration in milliseconds."""
        if not self.end_time or not self.start_time:
            return 0
        delta = self.end_time - self.start_time
        return int(delta.total_seconds() * 1000)

    def is_complete(self) -> bool:
        """Check if trace is complete (has end_time)."""
        return self.end_time is not None

    def to_dict(self) -> Dict[str, Any]:
        """Convert trace to dictionary representation."""
        return {
            "trace_id": str(self.trace_id),
            "start_time": self.start_time.isoformat() if self.start_time else None,
            "end_time": self.end_time.isoformat() if self.end_time else None,
            "total_duration_ms": self.get_total_duration_ms(),
            "root_steps": [step.to_dict() for step in self.root_steps],
            "metadata": self.metadata,
            "user_input": self.user_input,
        }

    @classmethod
    def from_trace_info(cls, trace_info: TraceInfo, steps: List[StepData] = None) -> 'Trace':
        """Create Trace from legacy TraceInfo."""
        trace = cls(
            trace_id=trace_info.trace_id,
            start_time=trace_info.start_time,
            metadata=trace_info.metadata or {},
            user_input=trace_info.metadata.get("user_input") if trace_info.metadata else None,
        )
        if steps:
            for step_data in steps:
                step_node = StepNode.from_step_data(step_data)
                trace.add_root_step(step_node)
        return trace
