"""
ModelCallReceipt: Receipt for LLM/model inference calls.

Tracks every call to an AI model with:
- Token counts (input/output)
- Latency measurements
- Model identification
- Cost attribution

This provides the audit trail required by Claim 1 for "all AI operations."
"""
from __future__ import annotations

from typing import List, Optional

from receipts.compat.pyd import Field, model_validator
from receipts.base import BaseReceipt, Domain


class ModelCallReceipt(BaseReceipt):
    """
    Receipt documenting a single LLM/model inference call.

    Every model call in the system emits this receipt, creating a complete
    audit trail of AI operations with token counts and cost attribution.
    """

    receipt_type: str = Field(default="ModelCallReceipt")
    domain: str = Field(default=Domain.GOVERNANCE.value)

    # Model identification
    model_id: str = Field(
        ...,
        description="Model identifier (e.g., 'claude-3-opus-20240229')",
    )
    provider: str = Field(
        ...,
        description="Provider name (e.g., 'anthropic', 'openai', 'local')",
    )

    # Call context
    episode_id: Optional[str] = Field(
        default=None,
        description="Episode this call belongs to",
    )
    intent_hash: Optional[str] = Field(
        default=None,
        description="Hash of the intent/prompt (for deduplication)",
    )

    # Token metrics
    input_tokens: int = Field(
        ...,
        ge=0,
        description="Number of input tokens",
    )
    output_tokens: int = Field(
        ...,
        ge=0,
        description="Number of output tokens",
    )
    total_tokens: int = Field(
        ...,
        ge=0,
        description="Total tokens (input + output)",
    )

    # Timing
    latency_ms: int = Field(
        ...,
        ge=0,
        description="End-to-end latency in milliseconds",
    )
    time_to_first_token_ms: Optional[int] = Field(
        default=None,
        ge=0,
        description="Time to first token for streaming responses",
    )

    # Cost attribution
    estimated_cost_usd: Optional[float] = Field(
        default=None,
        ge=0,
        description="Estimated cost in USD",
    )

    # Response metadata
    finish_reason: str = Field(
        ...,
        description="Why the model stopped generating (provider-specific: stop, length, end_turn, max_tokens, etc.)",
    )
    tool_calls: Optional[List[str]] = Field(
        default=None,
        description="List of tool call IDs if tool_use",
    )

    # Privacy-preserving content hashes (for integrations)
    input_hash: Optional[str] = Field(
        default=None,
        description="SHA256 hash of input content (privacy-preserving)",
    )
    output_hash: Optional[str] = Field(
        default=None,
        description="SHA256 hash of output content (privacy-preserving)",
    )
    message_count: Optional[int] = Field(
        default=None,
        ge=0,
        description="Number of messages in the conversation",
    )

    # Integration metadata
    integration_source: Optional[str] = Field(
        default=None,
        description="Which integration emitted this receipt (e.g., 'assay.integrations.openai')",
    )

    # OpenTelemetry bridge (for correlating receipts with OTel traces)
    otel_trace_id: Optional[str] = Field(
        default=None,
        description="OpenTelemetry trace ID for correlating with observability systems",
    )
    otel_span_id: Optional[str] = Field(
        default=None,
        description="OpenTelemetry span ID for correlating with observability systems",
    )

    # Error tracking
    error: Optional[str] = Field(
        default=None,
        description="Error message if call failed",
    )
    retry_count: int = Field(
        default=0,
        ge=0,
        description="Number of retries before this response",
    )

    @model_validator(mode="after")
    def _validate_total_tokens(self) -> "ModelCallReceipt":
        """Validate total_tokens matches input + output."""
        expected = self.input_tokens + self.output_tokens
        if self.total_tokens != expected:
            raise ValueError(
                f"total_tokens ({self.total_tokens}) must equal "
                f"input_tokens ({self.input_tokens}) + output_tokens ({self.output_tokens}) = {expected}"
            )
        return self


def create_model_call_receipt(
    model_id: str,
    provider: str,
    input_tokens: int,
    output_tokens: int,
    latency_ms: int,
    finish_reason: str = "stop",
    episode_id: Optional[str] = None,
    **kwargs,
) -> ModelCallReceipt:
    """Factory function to create a ModelCallReceipt."""
    import uuid
    return ModelCallReceipt(
        receipt_id=f"mcr_{uuid.uuid4().hex[:16]}",
        model_id=model_id,
        provider=provider,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        total_tokens=input_tokens + output_tokens,
        latency_ms=latency_ms,
        finish_reason=finish_reason,
        episode_id=episode_id,
        **kwargs,
    )


__all__ = ["ModelCallReceipt", "create_model_call_receipt"]
