"""
Core Pydantic models for .kurral artifact schema
Copied from kurral-cli for independent operation
"""

import json
import hashlib
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import Any, Optional, List, Union
from uuid import UUID, uuid4

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator


class ReplayLevel(str, Enum):
    """
    Replay classification levels for A/B testing
    
    A: Deterministic replay - everything matches, return artifact outputs directly
    B: Non-deterministic replay - something changed, re-execute LLM with cached tool calls
    """

    A = "A"  # Deterministic replay (determinism score below threshold)
    B = "B"  # Non-deterministic replay (determinism score above threshold or tools changed)


class LLMParameters(BaseModel):
    """Detailed LLM parameters for determinism tracking (null values excluded)"""
    
    temperature: float = Field(..., ge=0.0, le=2.0)
    top_p: Optional[float] = Field(None, ge=0.0, le=1.0)
    top_k: Optional[int] = Field(None, ge=0)
    max_tokens: Optional[int] = Field(None, ge=1)
    frequency_penalty: Optional[float] = Field(None, ge=-2.0, le=2.0)
    presence_penalty: Optional[float] = Field(None, ge=-2.0, le=2.0)
    seed: Optional[int] = Field(None, description="Random seed for determinism")


class TokenUsage(BaseModel):
    """Comprehensive token usage and cost metrics"""
    
    # Provider counters (from API response)
    prompt_tokens: int = Field(default=0, ge=0, description="Input tokens sent to LLM")
    completion_tokens: int = Field(default=0, ge=0, description="Output tokens generated by LLM")
    total_tokens: int = Field(default=0, ge=0, description="Total tokens (prompt + completion)")
    
    # Caching metrics (if provider supports it)
    cached_tokens: Optional[int] = Field(None, ge=0, description="Tokens served from cache")
    cache_creation_tokens: Optional[int] = Field(None, ge=0, description="Tokens used to populate cache")
    cache_read_tokens: Optional[int] = Field(None, ge=0, description="Tokens read from cache")
    
    # Cache efficiency
    cache_hit_rate: Optional[float] = Field(None, ge=0.0, le=1.0, description="Ratio of cached to total tokens")
    
    # Cost calculation
    estimated_cost_usd: Optional[float] = Field(None, ge=0.0, description="Estimated cost in USD")
    
    # Additional metrics
    reasoning_tokens: Optional[int] = Field(None, ge=0, description="Tokens used for reasoning (if applicable)")
    audio_tokens: Optional[int] = Field(None, ge=0, description="Audio tokens (if applicable)")


class ModelConfig(BaseModel):
    """Enhanced LLM model configuration with detailed parameters"""

    model_name: str = Field(..., description="Model identifier (e.g., gpt-4-0613)")
    model_version: Optional[str] = Field(None, description="Specific version hash if available")
    provider: str = Field(..., description="Provider (openai, anthropic, etc.)")
    parameters: LLMParameters = Field(..., description="Detailed model parameters")
    stop_sequences: Optional[list[str]] = None
    additional_params: Optional[dict[str, Any]] = Field(default=None, description="Additional provider-specific params")

    def cache_key(self) -> str:
        """Generate cache key from model config"""
        data = self.model_dump(exclude={"additional_params"})
        return hashlib.sha256(json.dumps(data, sort_keys=True).encode()).hexdigest()


class ResolvedPrompt(BaseModel):
    """Enhanced prompt template with hashing for replay fidelity"""

    template: str = Field(..., description="Original prompt template")
    template_id: Optional[str] = Field(None, description="Template identifier if versioned")
    template_hash: Optional[str] = Field(None, description="SHA256 hash of template")
    variables: dict[str, Any] = Field(
        default_factory=dict, description="Template variables used"
    )
    variables_hash: Optional[str] = Field(None, description="SHA256 hash of variables")
    final_text: str = Field(..., description="Final resolved prompt text")
    final_text_hash: Optional[str] = Field(None, description="SHA256 hash of final text")
    system_prompt: Optional[str] = Field(None, description="System prompt if separate")
    messages: Optional[list[dict[str, str]]] = Field(
        None, description="Chat messages if applicable"
    )

    @model_validator(mode="after")
    def compute_hashes(self) -> "ResolvedPrompt":
        """Auto-compute hashes if not provided"""
        if self.template_hash is None and self.template:
            self.template_hash = hashlib.sha256(self.template.encode()).hexdigest()
        
        if self.variables_hash is None and self.variables:
            try:
                # Try to serialize variables, filtering out non-serializable items
                serializable_vars = self._make_serializable(self.variables)
                vars_str = json.dumps(serializable_vars, sort_keys=True)
                self.variables_hash = hashlib.sha256(vars_str.encode()).hexdigest()
            except (TypeError, ValueError):
                # If still can't serialize, hash the string representation
                vars_str = str(self.variables)
                self.variables_hash = hashlib.sha256(vars_str.encode()).hexdigest()
        
        if self.final_text_hash is None and self.final_text:
            self.final_text_hash = hashlib.sha256(self.final_text.encode()).hexdigest()
        return self
    
    @staticmethod
    def _make_serializable(obj: Any, max_depth: int = 3, current_depth: int = 0) -> Any:
        """Make an object JSON serializable by filtering out problematic types"""
        if current_depth > max_depth:
            return "<max_depth>"
        
        if obj is None or isinstance(obj, (bool, int, float, str)):
            return obj
        
        if isinstance(obj, (list, tuple)):
            return [ResolvedPrompt._make_serializable(item, max_depth, current_depth + 1) for item in obj]
        
        if isinstance(obj, dict):
            result = {}
            for k, v in obj.items():
                # Skip callbacks and other non-serializable objects
                if k == "callbacks" or callable(v):
                    continue
                try:
                    result[str(k)] = ResolvedPrompt._make_serializable(v, max_depth, current_depth + 1)
                except:
                    result[str(k)] = f"<{type(v).__name__}>"
            return result
        
        # Try JSON serialization test
        try:
            json.dumps(obj)
            return obj
        except (TypeError, ValueError):
            return f"<{type(obj).__name__}>"

    def cache_key(self) -> str:
        """Generate cache key from prompt"""
        return self.final_text_hash or hashlib.sha256(self.final_text.encode()).hexdigest()


class EffectType(str, Enum):
    """Types of side effects a tool can have"""
    HTTP = "http"
    DB_WRITE = "db_write"
    EMAIL = "email"
    FS = "fs"
    MCP = "mcp"
    OTHER = "other"


class ToolCallStatus(str, Enum):
    """Status of tool call execution"""
    OK = "ok"
    ERROR = "error"
    TIMEOUT = "timeout"


class ToolCall(BaseModel):
    """Enhanced tool call record with stubbing and effect tracking"""

    # Identity and hierarchy
    id: str = Field(default_factory=lambda: str(uuid4()), description="Unique tool call ID")
    parent_id: Optional[str] = Field(None, description="Parent tool call ID for nested/chained calls")
    tool_name: str = Field(..., description="Name of the tool/function called")
    namespace: Optional[str] = Field(None, description="Tool namespace (e.g., mcp server)")
    agent_id: Optional[str] = Field(None, description="Agent that made the call")
    
    # Timing - explicit start and end for concurrency analysis
    start_time: datetime = Field(default_factory=datetime.utcnow, description="When execution started")
    end_time: Optional[datetime] = Field(None, description="When execution completed")
    latency_ms: Optional[int] = Field(None, description="Execution time in milliseconds")
    timestamp: datetime = Field(default_factory=datetime.utcnow, description="Legacy: use start_time")
    
    # Inputs/Outputs with hashing
    input: dict[str, Any] = Field(default_factory=dict, description="Input parameters")
    input_hash: Optional[str] = Field(None, description="SHA256 hash of input")
    output: dict[str, Any] = Field(default_factory=dict, description="Tool output/response")
    output_hash: Optional[str] = Field(None, description="SHA256 hash of output")
    summary: Optional[str] = Field(None, description="Human-readable summary of call")
    
    # Classification and tracking
    type: str = Field(default="tool", description="Operation type: tool, llm, agent, retrieval, etc.")
    effect_type: EffectType = Field(default=EffectType.OTHER, description="Type of side effect")
    status: ToolCallStatus = Field(default=ToolCallStatus.OK, description="Execution status")
    
    # Error handling
    error_flag: bool = Field(default=False, description="Whether call resulted in error")
    error_text: Optional[str] = Field(None, description="Error message if failed")
    error_type: Optional[str] = Field(None, description="Error type/class for categorization")
    error_stack: Optional[str] = Field(None, description="Full stack trace for debugging")
    
    # Caching and replay
    cache_key: str = Field(..., description="Content-addressed cache key")
    stubbed_in_replay: bool = Field(default=False, description="Was this stubbed in replay?")
    
    # Extensibility
    metadata: dict[str, Any] = Field(
        default_factory=dict, 
        description="Extensible metadata: cost, tokens, provider info, custom attributes"
    )
    
    # Legacy support (for backward compatibility)
    inputs: Optional[dict[str, Any]] = Field(None, description="Legacy: use 'input' instead")
    outputs: Optional[dict[str, Any]] = Field(None, description="Legacy: use 'output' instead")
    error: Optional[str] = Field(None, description="Legacy: use 'error_text' instead")
    duration_ms: Optional[int] = Field(None, description="Legacy: use 'latency_ms' instead")

    @staticmethod
    def generate_cache_key(tool_name: str, inputs: dict[str, Any]) -> str:
        """Generate deterministic cache key from tool call"""
        data = {"tool": tool_name, "inputs": inputs}
        content = json.dumps(data, sort_keys=True)
        return f"sha256:{hashlib.sha256(content.encode()).hexdigest()}"

    @model_validator(mode="before")
    @classmethod
    def ensure_cache_key_and_hashes(cls, values: dict[str, Any]) -> dict[str, Any]:
        """Auto-generate cache key, hashes, and timing if not provided"""
        # Handle legacy field names
        if "inputs" in values and "input" not in values:
            values["input"] = values["inputs"]
        if "outputs" in values and "output" not in values:
            values["output"] = values["outputs"]
        if "error" in values and values["error"] and "error_text" not in values:
            values["error_text"] = values["error"]
            values["error_flag"] = True
        if "duration_ms" in values and "latency_ms" not in values:
            values["latency_ms"] = values["duration_ms"]
        
        # Calculate end_time from start_time + latency if not provided
        if "end_time" not in values and "start_time" in values and "latency_ms" in values:
            from datetime import timedelta
            start = values["start_time"]
            if isinstance(start, datetime) and values["latency_ms"]:
                values["end_time"] = start + timedelta(milliseconds=values["latency_ms"])
        
        # Calculate latency_ms from start_time and end_time if not provided
        if "latency_ms" not in values and "start_time" in values and "end_time" in values:
            start = values["start_time"]
            end = values["end_time"]
            if isinstance(start, datetime) and isinstance(end, datetime):
                values["latency_ms"] = int((end - start).total_seconds() * 1000)
            
        # Generate cache key
        if "cache_key" not in values or not values["cache_key"]:
            tool_name = values.get("tool_name", "unknown")
            input_data = values.get("input", values.get("inputs", {}))
            values["cache_key"] = cls.generate_cache_key(tool_name, input_data)
        
        # Generate hashes
        if "input_hash" not in values and values.get("input"):
            input_str = json.dumps(values["input"], sort_keys=True)
            values["input_hash"] = hashlib.sha256(input_str.encode()).hexdigest()
        
        if "output_hash" not in values and values.get("output"):
            output_str = json.dumps(values["output"], sort_keys=True)
            values["output_hash"] = hashlib.sha256(output_str.encode()).hexdigest()
        
        return values


class GraphVersion(BaseModel):
    """Graph and tool versioning for determinism tracking"""
    
    model_config = ConfigDict(exclude_none=True)
    
    graph_hash: Optional[str] = Field(None, description="SHA256 hash of compiled graph structure")
    graph_checksum: Optional[str] = Field(None, description="Checksum including node/edge configs")
    tool_schemas_hash: Optional[str] = Field(None, description="Combined hash of all tool schemas")
    tools: Optional[list[dict[str, Any]]] = Field(None, description="Individual tool definitions with hashes")
    policy_pack_hash: Optional[str] = Field(None, description="Hash of active policy pack")
    policy_version: Optional[str] = Field(None, description="Policy pack version string")
    preprocessors_hash: Optional[str] = Field(None, description="Hash of all preprocessors")
    postprocessors_hash: Optional[str] = Field(None, description="Hash of all postprocessors")


class TimeEnvironment(BaseModel):
    """Extended time and environment snapshot with feature flags"""

    timestamp: datetime = Field(..., description="Execution timestamp (ISO 8601)")
    timezone: str = Field(default="UTC")
    wall_clock_time: str = Field(..., description="Human-readable time")
    locale: Optional[str] = Field(None, description="Locale setting (e.g., en_US)")
    environment_vars: Optional[dict[str, str]] = Field(
        None, description="Relevant env vars (sanitized)"
    )
    feature_flags: Optional[dict[str, Any]] = Field(
        None, description="Feature flags active during execution"
    )
    clock_freeze: Optional[bool] = Field(None, description="Was time frozen for testing? (omitted if false/null)")


class DeterminismReport(BaseModel):
    """Determinism analysis and scoring"""

    overall_score: float = Field(..., ge=0.0, le=1.0, description="Overall score (0.0-1.0)")
    breakdown: dict[str, float] = Field(..., description="Score breakdown by factor")
    missing_fields: list[str] = Field(default_factory=list)
    warnings: list[str] = Field(default_factory=list)

    @field_validator("breakdown")
    @classmethod
    def validate_breakdown(cls, v: dict[str, float]) -> dict[str, float]:
        """Ensure breakdown scores are in valid range"""
        for key, score in v.items():
            if not 0.0 <= score <= 1.0:
                raise ValueError(f"Score for {key} must be between 0.0 and 1.0")
        return v


class KurralArtifact(BaseModel):
    """
    Complete .kurral artifact - the core data structure
    """

    # Identity
    kurral_id: UUID = Field(default_factory=uuid4, description="Unique artifact ID")
    run_id: str = Field(..., description="Source trace/run ID (e.g., from LangSmith)")
    tenant_id: str = Field(..., description="Tenant/organization identifier")
    semantic_buckets: list[str] = Field(
        default_factory=list, description="Business logic categories"
    )
    environment: str = Field(default="production", description="Environment (prod/staging/dev)")

    # Metadata
    schema_version: str = Field(default="1.0.0", description="Artifact schema version")
    created_at: datetime = Field(default_factory=datetime.utcnow)
    created_by: Optional[str] = Field(None, description="User/service that created artifact")

    # Determinism (determined during replay, not during artifact generation)
    deterministic: bool = Field(False, description="Overall determinism flag - determined during replay")
    replay_level: Optional[ReplayLevel] = Field(None, description="Replay reliability level (A/B) - determined during replay")
    determinism_report: Optional[DeterminismReport] = Field(None, description="Detailed scoring - calculated during replay")

    # Execution Data
    inputs: dict[str, Any] = Field(..., description="Function/agent inputs")
    outputs: dict[str, Any] = Field(..., description="Function/agent outputs")
    error: Optional[str] = Field(None, description="Error message if execution failed")

    # LLM Configuration
    llm_config: ModelConfig = Field(..., description="Model configuration")
    resolved_prompt: ResolvedPrompt = Field(..., description="Prompt with variables")

    # Graph & Tool Versioning
    graph_version: Optional[GraphVersion] = Field(None, description="Graph and tool schema hashes")

    # Tool Calls
    tool_calls: list[ToolCall] = Field(default_factory=list, description="All tool invocations")

    # Time & Environment
    time_env: TimeEnvironment = Field(..., description="Time and environment snapshot")

    # Metrics
    duration_ms: int = Field(..., ge=0, description="Total execution time")
    cost_usd: Optional[float] = Field(None, ge=0.0, description="Estimated cost in USD")
    token_usage: Optional[TokenUsage] = Field(
        None, description="Comprehensive token usage and cost metrics"
    )

    # Storage
    object_storage_uri: Optional[str] = Field(
        None, description="R2/storage URI if stored externally"
    )
    tags: Optional[dict[str, str]] = Field(None, description="Custom tags")

    # Note: replay_level and deterministic are set during replay, not during artifact generation
    # No validation needed here - these fields are optional during generation

    def to_json(self, pretty: bool = False) -> str:
        """Serialize to JSON string"""
        indent = 2 if pretty else None
        return self.model_dump_json(indent=indent, exclude_none=True)

    def save(self, filepath: Union[str, Path]) -> None:
        """Save artifact to .kurral file"""
        import tempfile
        import shutil
        
        path = Path(filepath)
        path.parent.mkdir(parents=True, exist_ok=True)

        # Write to temp file first, then atomically move to final location
        # This prevents leaving empty files if serialization fails
        try:
            json_content = self.to_json(pretty=True)
            if not json_content or len(json_content.strip()) == 0:
                raise ValueError("Artifact serialization produced empty JSON")
            
            # Write to temp file
            temp_path = path.with_suffix(path.suffix + '.tmp')
            with open(temp_path, "w", encoding='utf-8') as f:
                f.write(json_content)
            
            # Atomically move temp file to final location
            shutil.move(str(temp_path), str(path))
        except Exception as e:
            # Clean up temp file if it exists
            temp_path = path.with_suffix(path.suffix + '.tmp')
            if temp_path.exists():
                temp_path.unlink()
            # Remove empty file if it was created
            if path.exists() and path.stat().st_size == 0:
                path.unlink()
            raise RuntimeError(f"Failed to save artifact to {filepath}: {e}") from e

    @classmethod
    def load(cls, filepath: Union[str, Path]) -> "KurralArtifact":
        """Load artifact from .kurral file"""
        with open(filepath, "r", encoding="utf-8") as f:
            data = json.load(f)
        return cls.model_validate(data)

    @classmethod
    def from_json(cls, json_str: str) -> "KurralArtifact":
        """Deserialize from JSON string"""
        return cls.model_validate_json(json_str)

    def cache_key(self) -> str:
        """Generate cache key for this artifact"""
        data = {
            "model": self.llm_config.cache_key(),
            "prompt": self.resolved_prompt.cache_key(),
            "inputs": self.inputs,
        }
        return hashlib.sha256(json.dumps(data, sort_keys=True).encode()).hexdigest()


class ReplayOverrides(BaseModel):
    """Optional overrides for replay execution"""

    prompt: Optional[str] = Field(None, description="Override prompt text")
    temperature: Optional[float] = Field(None, ge=0.0, le=2.0)
    max_tokens: Optional[int] = Field(None, ge=1)
    model_name: Optional[str] = None
    inputs: Optional[dict[str, Any]] = None


class ReplayResult(BaseModel):
    """Result of a replay execution"""

    kurral_id: UUID = Field(..., description="Original artifact ID")
    replay_timestamp: datetime = Field(default_factory=datetime.utcnow)
    outputs: dict[str, Any] = Field(..., description="Replay outputs")
    match: bool = Field(..., description="Whether outputs match original")
    diff: Optional[dict[str, Any]] = Field(None, description="Differences if any")
    tool_calls: list[ToolCall] = Field(default_factory=list)
    error: Optional[str] = None
    duration_ms: int = Field(..., ge=0)
    cache_hits: int = Field(default=0, ge=0)
    cache_misses: int = Field(default=0, ge=0)
    new_tool_calls: list[ToolCall] = Field(
        default_factory=list,
        description="Tool calls made during replay that were not in original artifact"
    )
    unused_tool_calls: list[ToolCall] = Field(
        default_factory=list,
        description="Tool calls from original artifact that were not used during replay"
    )
    stream: Optional[dict[str, Any]] = Field(
        None,
        description="Reconstructed output stream with items/full_text/stream_map",
    )
    graph_version: Optional["GraphVersion"] = Field(
        None, description="Pinned graph metadata used during replay"
    )
    llm_state: Optional["ReplayLLMState"] = Field(
        None, description="Hydrated LLM sampling state used during replay"
    )
    validation: Optional["ReplayValidation"] = Field(
        None, description="Replay validation metrics (hash + structural comparisons)"
    )
    replay_metadata: Optional["ReplayMetadata"] = Field(
        None, description="Metadata describing the replay execution and assertions"
    )


class ReplayLLMState(BaseModel):
    """Hydrated sampling parameters for replay"""

    model_name: str
    provider: str
    model_version: Optional[str] = None
    temperature: Optional[float] = None
    top_p: Optional[float] = None
    top_k: Optional[int] = None
    max_tokens: Optional[int] = None
    frequency_penalty: Optional[float] = None
    presence_penalty: Optional[float] = None
    seed: Optional[int] = None


class ReplayValidation(BaseModel):
    """Validation information for comparing original and replay outputs"""

    original_hash: str
    replay_hash: str
    hash_match: bool
    structural_match: bool
    diff: Optional[dict[str, Any]] = None


class AssertionResult(BaseModel):
    """Represents the outcome of a replay assertion"""

    name: str
    passed: bool
    details: Optional[dict[str, Any]] = None


class ReplayMetadata(BaseModel):
    """Metadata describing a replay execution"""

    replay_id: str
    record_ref: str
    replay_level: ReplayLevel
    assertion_results: List[AssertionResult] = Field(default_factory=list)

