"""Data models for LLM requests and responses with provenance tracking."""

from datetime import datetime
from typing import Any, Dict, List, Optional
from uuid import UUID, uuid4

from pydantic import BaseModel, Field


class ProvenanceMetadata(BaseModel):
    """Metadata for tracking request provenance."""
    
    workflow_id: Optional[str] = None
    chunk_id: Optional[str] = None
    section_id: Optional[str] = None
    source_id: Optional[str] = None
    character_range: Optional[tuple[int, int]] = None
    caller_id: Optional[str] = None
    tags: Dict[str, str] = Field(default_factory=dict)


class TokenUsage(BaseModel):
    """Token usage information."""
    
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    
    @property
    def cost_estimate(self) -> float:
        """Rough cost estimate (customize per model)."""
        # Example: Claude Sonnet pricing
        prompt_cost = self.prompt_tokens * 0.003 / 1000
        completion_cost = self.completion_tokens * 0.015 / 1000
        return prompt_cost + completion_cost


class LLMRequest(BaseModel):
    """Normalized LLM request with provenance."""
    
    id: UUID = Field(default_factory=uuid4)
    model: str
    messages: List[Dict[str, str]]
    temperature: float = 0.0
    max_tokens: int = 4000
    top_p: Optional[float] = None
    stop_sequences: Optional[List[str]] = None
    
    # Provenance
    metadata: ProvenanceMetadata = Field(default_factory=ProvenanceMetadata)
    
    # Additional LiteLLM params
    extra_params: Dict[str, Any] = Field(default_factory=dict)
    
    created_at: datetime = Field(default_factory=datetime.utcnow)
    
    def to_litellm_params(self) -> Dict[str, Any]:
        """Convert to LiteLLM completion parameters."""
        params = {
            "model": self.model,
            "messages": self.messages,
            "temperature": self.temperature,
            "max_tokens": self.max_tokens,
        }
        
        if self.top_p is not None:
            params["top_p"] = self.top_p
        if self.stop_sequences:
            params["stop"] = self.stop_sequences
            
        params.update(self.extra_params)
        return params
    
    def normalize_for_cache_key(self) -> str:
        """Normalize request for deterministic cache key generation."""
        # Sort messages, normalize whitespace
        normalized_messages = []
        for msg in self.messages:
            normalized_messages.append({
                "role": msg["role"],
                "content": " ".join(msg["content"].split())
            })
        
        # Create deterministic string
        parts = [
            f"model:{self.model}",
            f"temperature:{self.temperature}",
            f"max_tokens:{self.max_tokens}",
            f"messages:{normalized_messages}"
        ]
        
        if self.top_p is not None:
            parts.append(f"top_p:{self.top_p}")
        if self.stop_sequences:
            parts.append(f"stop:{sorted(self.stop_sequences)}")
        
        return "|".join(parts)


class LLMResponse(BaseModel):
    """LLM response with provenance and metadata."""
    
    id: UUID = Field(default_factory=uuid4)
    request_id: UUID
    content: str
    
    # Token tracking
    usage: Optional[TokenUsage] = None
    
    # Cache information
    from_cache: bool = False
    cache_key: Optional[str] = None
    
    # Provider metadata
    provider: str
    model: str
    provider_request_id: Optional[str] = None
    latency_ms: Optional[float] = None
    
    # Provenance (copied from request)
    metadata: ProvenanceMetadata
    
    created_at: datetime = Field(default_factory=datetime.utcnow)
    
    @property
    def cost_estimate(self) -> float:
        """Get estimated cost of this request."""
        if self.usage:
            return self.usage.cost_estimate
        return 0.0


class CompletionBuilder:
    """Fluent builder for LLM requests."""
    
    def __init__(self, client: 'LedgerClient'):
        self._client = client
        self._model: str = "openai/gpt-4o"
        self._messages: List[Dict[str, str]] = []
        self._temperature: float = 0.0
        self._max_tokens: int = 4000
        self._metadata = ProvenanceMetadata()
        self._extra_params: Dict[str, Any] = {}
    
    def model(self, model: str) -> 'CompletionBuilder':
        """Set the model."""
        self._model = model
        return self
    
    def system(self, content: str) -> 'CompletionBuilder':
        """Add system message."""
        self._messages.insert(0, {"role": "system", "content": content})
        return self
    
    def user(self, content: str) -> 'CompletionBuilder':
        """Add user message."""
        self._messages.append({"role": "user", "content": content})
        return self
    
    def assistant(self, content: str) -> 'CompletionBuilder':
        """Add assistant message (for multi-turn)."""
        self._messages.append({"role": "assistant", "content": content})
        return self
    
    def temperature(self, temp: float) -> 'CompletionBuilder':
        """Set temperature."""
        self._temperature = temp
        return self
    
    def max_tokens(self, tokens: int) -> 'CompletionBuilder':
        """Set max tokens."""
        self._max_tokens = tokens
        return self
    
    def with_metadata(self, **kwargs) -> 'CompletionBuilder':
        """Add provenance metadata."""
        for key, value in kwargs.items():
            if hasattr(self._metadata, key):
                setattr(self._metadata, key, value)
        return self
    
    def extra(self, **kwargs) -> 'CompletionBuilder':
        """Add extra LiteLLM parameters."""
        self._extra_params.update(kwargs)
        return self
    
    def build(self) -> LLMRequest:
        """Build the request."""
        return LLMRequest(
            model=self._model,
            messages=self._messages,
            temperature=self._temperature,
            max_tokens=self._max_tokens,
            metadata=self._metadata,
            extra_params=self._extra_params
        )
    
    def execute(self, use_cache: bool = True) -> LLMResponse:
        """Build and execute the request."""
        request = self.build()
        return self._client.complete(request, use_cache=use_cache)
    
    async def execute_async(self, use_cache: bool = True) -> LLMResponse:
        """Build and execute the request asynchronously."""
        request = self.build()
        return await self._client.complete_async(request, use_cache=use_cache)
