"""
MLX Domain Receipts

Receipt types for the MLX local inference subsystem (Apple Silicon).
All receipts are Ed25519-signed, JCS-canonical, and stored in the standard
receipt ledger for unified audit trails.

---
feature:
  id: feature.mlx_receipts
  name: MLX Domain Receipts
  kind: module
  summary: Receipt models for MLX local inference events
  owners: ["@codex"]
  stage: production
  tags: [receipts, mlx, local, inference, p1]
  depends_on: [feature.receipt_writer]
  inputs: [MLX inference events]
  outputs: [Signed receipts in ledger]
  invariants:
    - All receipts are Ed25519-signed
    - Receipt IDs are globally unique
    - All local inference is $0 cost
  risk_axes:
    unsigned_receipt: 0.0  # Enforced by writer
    missing_metrics: 0.01
---
"""

from __future__ import annotations

import logging
from datetime import datetime, timezone
from typing import Any, Dict, List, Literal, Optional

from pydantic import Field

from receipts.base import BaseReceipt, Domain
from receipts.compat.pyd import ConfigDictLike

logger = logging.getLogger(__name__)


# =============================================================================
# Receipt Models
# =============================================================================

class MLXInferenceReceipt(BaseReceipt):
    """
    Receipt emitted for every MLX local inference call.

    Captures metrics for local ($0) inference on Apple Silicon using MLX.
    Used for performance tracking, debugging, and audit trails.
    """

    receipt_type: str = Field(default="MLXInferenceReceipt", alias="type")
    domain: str = Field(default="mlx")

    # Request identifiers
    request_id: str = Field(description="Unique ID for this MLX request")
    task_id: Optional[str] = Field(
        default=None,
        description="Task ID if this inference is part of a larger task"
    )

    # Model selection
    model: str = Field(description="MLX model used (e.g., 'llama3.2:3b')")
    tier: str = Field(
        description="Model tier selected (reflex/reason/code)"
    )

    # Task context
    glyphs: List[str] = Field(
        default_factory=list,
        description="Task glyphs used for tier selection"
    )
    prompt_preview: str = Field(
        default="",
        description="First 200 chars of the prompt"
    )
    response_preview: str = Field(
        default="",
        description="First 200 chars of the response"
    )

    # Metrics
    tokens_in: int = Field(
        default=0,
        description="Input tokens consumed"
    )
    tokens_out: int = Field(
        default=0,
        description="Output tokens generated"
    )
    duration_ms: float = Field(
        default=0.0,
        description="Total inference duration in milliseconds"
    )

    # Cost (always $0 for local)
    cost_usd: float = Field(
        default=0.0,
        description="Cost in USD (always $0 for local inference)"
    )

    # Success/error
    success: bool = Field(
        default=True,
        description="Whether inference succeeded"
    )
    error: Optional[str] = Field(
        default=None,
        description="Error message if failed"
    )

    # System context
    ollama_host: Optional[str] = Field(
        default=None,
        description="Ollama server endpoint used"
    )

    model_config = ConfigDictLike(json_schema_extra={
        "example": {
            "request_id": "mlx_abc123",
            "task_id": "task_xyz789",
            "model": "llama3.2:3b",
            "tier": "reason",
            "glyphs": ["🧠", "⚡"],
            "prompt_preview": "Explain quantum computing in simple terms...",
            "response_preview": "Quantum computing uses quantum bits or qubits...",
            "tokens_in": 45,
            "tokens_out": 234,
            "duration_ms": 1823.5,
            "cost_usd": 0.0,
            "success": True,
            "ollama_host": "http://localhost:11434"
        }
    })


class MLXModelLoadReceipt(BaseReceipt):
    """
    Receipt emitted when an MLX model is loaded into memory.

    Tracks model loading events for performance analysis and debugging.
    """

    receipt_type: str = Field(default="MLXModelLoadReceipt", alias="type")
    domain: str = Field(default="mlx")

    # Model info
    model: str = Field(description="Model name/identifier")
    tier: str = Field(description="Tier this model serves")

    # Load metrics
    load_duration_ms: float = Field(
        default=0.0,
        description="Time taken to load model into memory"
    )
    model_size_mb: Optional[float] = Field(
        default=None,
        description="Model size in megabytes"
    )

    # Success/error
    success: bool = Field(
        default=True,
        description="Whether model loaded successfully"
    )
    error: Optional[str] = Field(
        default=None,
        description="Error message if load failed"
    )

    model_config = ConfigDictLike(json_schema_extra={
        "example": {
            "model": "llama3.2:3b",
            "tier": "reason",
            "load_duration_ms": 2345.6,
            "model_size_mb": 2048.0,
            "success": True
        }
    })


class MLXTierSelectionReceipt(BaseReceipt):
    """
    Receipt emitted when MLX automatically selects a tier based on glyphs.

    Tracks tier selection logic for debugging and optimization.
    """

    receipt_type: str = Field(default="MLXTierSelectionReceipt", alias="type")
    domain: str = Field(default="mlx")

    # Input
    glyphs: List[str] = Field(
        default_factory=list,
        description="Task glyphs analyzed"
    )
    explicit_tier: Optional[str] = Field(
        default=None,
        description="Explicitly requested tier (if any)"
    )

    # Output
    selected_tier: str = Field(
        description="Tier selected by auto-selection logic"
    )
    selected_model: str = Field(
        description="Model selected for the tier"
    )

    # Reasoning
    selection_reason: str = Field(
        default="",
        description="Human-readable reason for tier selection"
    )

    # Fallback tracking
    preferred_model_available: bool = Field(
        default=True,
        description="Whether preferred model for tier was available"
    )
    fallback_applied: bool = Field(
        default=False,
        description="Whether fallback to different tier was needed"
    )

    model_config = ConfigDictLike(json_schema_extra={
        "example": {
            "glyphs": ["💻", "🖥️", "debug"],
            "explicit_tier": None,
            "selected_tier": "code",
            "selected_model": "qwen2.5-coder:7b",
            "selection_reason": "Code-related glyphs detected",
            "preferred_model_available": True,
            "fallback_applied": False
        }
    })


# =============================================================================
# Convenience constructors
# =============================================================================

def create_mlx_inference_receipt(
    request_id: str,
    model: str,
    tier: str,
    tokens_in: int,
    tokens_out: int,
    duration_ms: float,
    prompt_preview: str = "",
    response_preview: str = "",
    task_id: Optional[str] = None,
    glyphs: Optional[List[str]] = None,
    success: bool = True,
    error: Optional[str] = None,
    **kwargs
) -> MLXInferenceReceipt:
    """Create an MLX inference receipt."""
    return MLXInferenceReceipt(
        receipt_id=f"mlx_inf_{request_id}",
        request_id=request_id,
        task_id=task_id,
        model=model,
        tier=tier,
        glyphs=glyphs or [],
        prompt_preview=prompt_preview[:200],
        response_preview=response_preview[:200],
        tokens_in=tokens_in,
        tokens_out=tokens_out,
        duration_ms=duration_ms,
        cost_usd=0.0,  # Always $0 for local
        success=success,
        error=error,
        ollama_host="http://localhost:11434",
        **kwargs
    )


def create_mlx_model_load_receipt(
    model: str,
    tier: str,
    load_duration_ms: float,
    model_size_mb: Optional[float] = None,
    success: bool = True,
    error: Optional[str] = None,
    **kwargs
) -> MLXModelLoadReceipt:
    """Create an MLX model load receipt."""
    import uuid
    return MLXModelLoadReceipt(
        receipt_id=f"mlx_load_{uuid.uuid4().hex[:12]}",
        model=model,
        tier=tier,
        load_duration_ms=load_duration_ms,
        model_size_mb=model_size_mb,
        success=success,
        error=error,
        **kwargs
    )


def create_mlx_tier_selection_receipt(
    glyphs: List[str],
    selected_tier: str,
    selected_model: str,
    selection_reason: str = "",
    explicit_tier: Optional[str] = None,
    preferred_model_available: bool = True,
    fallback_applied: bool = False,
    **kwargs
) -> MLXTierSelectionReceipt:
    """Create an MLX tier selection receipt."""
    import uuid
    return MLXTierSelectionReceipt(
        receipt_id=f"mlx_tier_{uuid.uuid4().hex[:12]}",
        glyphs=glyphs,
        explicit_tier=explicit_tier,
        selected_tier=selected_tier,
        selected_model=selected_model,
        selection_reason=selection_reason,
        preferred_model_available=preferred_model_available,
        fallback_applied=fallback_applied,
        **kwargs
    )
