"""Main LLM client wrapping LiteLLM with enhanced features."""

import time
from typing import Optional, List, Dict, Any

import litellm
from litellm import completion, acompletion

from .models import LLMRequest, LLMResponse, TokenUsage, ProvenanceMetadata, CompletionBuilder
from .cache import DeterministicCache, InMemoryCache, RedisCache, CacheBackend
from .persistence import LLMPersistence


class LedgerClient:
    """
    LLM client with deterministic caching, provenance tracking, and persistence.
    
    Thin wrapper around LiteLLM that adds:
    - SHA256-based deterministic caching
    - Full request/response persistence
    - Provenance metadata tracking
    - Token usage accounting
    - Automatic retry and rate limiting (via LiteLLM)
    """
    
    def __init__(
        self,
        cache_backend: Optional[CacheBackend] = None,
        database_url: Optional[str] = None,
        cache_ttl: Optional[int] = None,
        enable_persistence: bool = True,
        enable_cache: bool = True,
        default_model: str = "openai/gpt-4o",
        litellm_kwargs: Optional[Dict[str, Any]] = None
    ):
        """
        Initialize LLM client.
        
        Args:
            cache_backend: Cache backend (defaults to InMemoryCache)
            database_url: Database connection string (defaults to SQLite)
            cache_ttl: Cache TTL in seconds (None = no expiration)
            enable_persistence: Whether to persist requests/responses
            enable_cache: Whether to use caching
            default_model: Default model for completions
            litellm_kwargs: Additional kwargs for LiteLLM configuration
        """
        # Set up caching
        self.enable_cache = enable_cache
        if enable_cache:
            backend = cache_backend or InMemoryCache()
            self.cache = DeterministicCache(backend=backend, ttl=cache_ttl)
        else:
            self.cache = None
        
        # Set up persistence
        self.enable_persistence = enable_persistence
        if enable_persistence:
            db_url = database_url or "sqlite:///llm_gateway.db"
            self.persistence = LLMPersistence(db_url)
        else:
            self.persistence = None
        
        self.default_model = default_model
        
        # Configure LiteLLM
        if litellm_kwargs:
            for key, value in litellm_kwargs.items():
                setattr(litellm, key, value)
        
        # Statistics
        self._total_requests = 0
        self._cache_hits = 0
        self._total_tokens = 0
        self._total_cost = 0.0
    
    def completion(self) -> CompletionBuilder:
        """Create a fluent completion builder."""
        return CompletionBuilder(self)
    
    def complete(
        self,
        request: LLMRequest,
        use_cache: bool = True,
        num_retries: int = 3,
        timeout: int = 300
    ) -> LLMResponse:
        """
        Execute completion request (synchronous).
        
        Args:
            request: LLM request with provenance
            use_cache: Whether to use cache
            num_retries: Number of retries on failure
            timeout: Request timeout in seconds
            
        Returns:
            LLM response with provenance
        """
        self._total_requests += 1
        
        # Check cache first
        if use_cache and self.enable_cache:
            cached_response = self.cache.get(request)
            if cached_response:
                self._cache_hits += 1
                return cached_response
        
        # Execute via LiteLLM
        start_time = time.time()
        
        try:
            litellm_response = completion(
                **request.to_litellm_params(),
                num_retries=num_retries,
                timeout=timeout
            )
        except Exception as e:
            # Log error and re-raise
            if self.enable_persistence:
                # Could log failed requests here
                pass
            raise
        
        latency_ms = (time.time() - start_time) * 1000
        
        # Extract response data
        content = litellm_response.choices[0].message.content
        usage_data = litellm_response.usage
        
        usage = TokenUsage(
            prompt_tokens=usage_data.prompt_tokens,
            completion_tokens=usage_data.completion_tokens,
            total_tokens=usage_data.total_tokens
        )
        
        # Create response object
        response = LLMResponse(
            request_id=request.id,
            content=content,
            usage=usage,
            from_cache=False,
            provider=litellm_response.model.split("/")[0] if "/" in litellm_response.model else "unknown",
            model=request.model,
            provider_request_id=getattr(litellm_response, 'id', None),
            latency_ms=latency_ms,
            metadata=request.metadata
        )
        
        # Update statistics
        self._total_tokens += usage.total_tokens
        self._total_cost += usage.cost_estimate
        
        # Cache response
        if use_cache and self.enable_cache:
            cache_key = self.cache.set(request, response)
            response.cache_key = cache_key
        
        # Persist to database
        if self.enable_persistence:
            cache_key = self.cache.generate_cache_key(request) if self.enable_cache else ""
            self.persistence.save_request(request, cache_key)
            self.persistence.save_response(response)
        
        return response
    
    async def complete_async(
        self,
        request: LLMRequest,
        use_cache: bool = True,
        num_retries: int = 3,
        timeout: int = 300
    ) -> LLMResponse:
        """
        Execute completion request (asynchronous).
        
        Args:
            request: LLM request with provenance
            use_cache: Whether to use cache
            num_retries: Number of retries on failure
            timeout: Request timeout in seconds
            
        Returns:
            LLM response with provenance
        """
        self._total_requests += 1
        
        # Check cache first
        if use_cache and self.enable_cache:
            cached_response = self.cache.get(request)
            if cached_response:
                self._cache_hits += 1
                return cached_response
        
        # Execute via LiteLLM
        start_time = time.time()
        
        try:
            litellm_response = await acompletion(
                **request.to_litellm_params(),
                num_retries=num_retries,
                timeout=timeout
            )
        except Exception as e:
            if self.enable_persistence:
                pass
            raise
        
        latency_ms = (time.time() - start_time) * 1000
        
        # Extract response data
        content = litellm_response.choices[0].message.content
        usage_data = litellm_response.usage
        
        usage = TokenUsage(
            prompt_tokens=usage_data.prompt_tokens,
            completion_tokens=usage_data.completion_tokens,
            total_tokens=usage_data.total_tokens
        )
        
        # Create response object
        response = LLMResponse(
            request_id=request.id,
            content=content,
            usage=usage,
            from_cache=False,
            provider=litellm_response.model.split("/")[0] if "/" in litellm_response.model else "unknown",
            model=request.model,
            provider_request_id=getattr(litellm_response, 'id', None),
            latency_ms=latency_ms,
            metadata=request.metadata
        )
        
        # Update statistics
        self._total_tokens += usage.total_tokens
        self._total_cost += usage.cost_estimate
        
        # Cache response
        if use_cache and self.enable_cache:
            cache_key = self.cache.set(request, response)
            response.cache_key = cache_key
        
        # Persist to database
        if self.enable_persistence:
            cache_key = self.cache.generate_cache_key(request) if self.enable_cache else ""
            self.persistence.save_request(request, cache_key)
            self.persistence.save_response(response)
        
        return response
    
    def quick_complete(
        self,
        prompt: str,
        system: Optional[str] = None,
        model: Optional[str] = None,
        **kwargs
    ) -> LLMResponse:
        """
        Quick completion without builder (convenience method).
        
        Args:
            prompt: User prompt
            system: System prompt (optional)
            model: Model to use (defaults to default_model)
            **kwargs: Additional parameters (temperature, max_tokens, etc.)
            
        Returns:
            LLM response with provenance
        """
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        
        request = LLMRequest(
            model=model or self.default_model,
            messages=messages,
            **kwargs
        )
        
        return self.complete(request)
    
    def get_stats(self) -> dict:
        """Get client statistics."""
        stats = {
            "total_requests": self._total_requests,
            "total_tokens": self._total_tokens,
            "total_cost_estimate": self._total_cost,
            "cache_hit_rate": self._cache_hits / self._total_requests if self._total_requests > 0 else 0.0
        }
        
        if self.enable_cache and self.cache:
            stats["cache_stats"] = self.cache.stats
        
        return stats
    
    def clear_cache(self) -> None:
        """Clear all cached responses."""
        if self.enable_cache and self.cache:
            self.cache.clear()
    
    def get_request(self, request_id: str) -> Optional[LLMRequest]:
        """Retrieve persisted request by ID."""
        if self.enable_persistence:
            from uuid import UUID
            return self.persistence.get_request(UUID(request_id))
        return None
    
    def get_response(self, response_id: str) -> Optional[LLMResponse]:
        """Retrieve persisted response by ID."""
        if self.enable_persistence:
            from uuid import UUID
            return self.persistence.get_response(UUID(response_id))
        return None
    
    def query_requests(
        self,
        workflow_id: Optional[str] = None,
        chunk_id: Optional[str] = None,
        limit: int = 100
    ) -> List[LLMRequest]:
        """Query requests by provenance metadata."""
        if self.enable_persistence:
            return self.persistence.get_requests_by_metadata(
                workflow_id=workflow_id,
                chunk_id=chunk_id,
                limit=limit
            )
        return []
