from abc import ABC, abstractmethod
from typing import Dict, List, Any, AsyncGenerator, Optional, Union
from pydantic import BaseModel, Field # type: ignore
from enum import Enum
import asyncio
import json
from dataclasses import dataclass
import os
import defines
from logger import logger

# Standard message format for all providers
@dataclass
class LLMMessage:
    role: str  # "user", "assistant", "system"
    content: str
    
    def __getitem__(self, key: str):
        """Allow dictionary-style access for backward compatibility"""
        if key == 'role':
            return self.role
        elif key == 'content':
            return self.content
        else:
            raise KeyError(f"'{key}' not found in LLMMessage")
    
    def __setitem__(self, key: str, value: str):
        """Allow dictionary-style assignment"""
        if key == 'role':
            self.role = value
        elif key == 'content':
            self.content = value
        else:
            raise KeyError(f"'{key}' not found in LLMMessage")
    
    @classmethod
    def from_dict(cls, data: Dict[str, str]) -> 'LLMMessage':
        """Create LLMMessage from dictionary"""
        return cls(role=data['role'], content=data['content'])
    
    def to_dict(self) -> Dict[str, str]:
        """Convert LLMMessage to dictionary"""
        return {'role': self.role, 'content': self.content}
    
# Enhanced usage statistics model
class UsageStats(BaseModel):
    """Comprehensive usage statistics across all providers"""
    # Token counts (standardized across providers)
    prompt_tokens: Optional[int] = Field(default=None, description="Number of tokens in the prompt")
    completion_tokens: Optional[int] = Field(default=None, description="Number of tokens in the completion")
    total_tokens: Optional[int] = Field(default=None, description="Total number of tokens used")
    
    # Ollama-specific detailed stats
    prompt_eval_count: Optional[int] = Field(default=None, description="Number of tokens evaluated in prompt")
    prompt_eval_duration: Optional[int] = Field(default=None, description="Time spent evaluating prompt (nanoseconds)")
    eval_count: Optional[int] = Field(default=None, description="Number of tokens generated")
    eval_duration: Optional[int] = Field(default=None, description="Time spent generating tokens (nanoseconds)")
    total_duration: Optional[int] = Field(default=None, description="Total request duration (nanoseconds)")
    
    # Performance metrics
    tokens_per_second: Optional[float] = Field(default=None, description="Generation speed in tokens/second")
    prompt_tokens_per_second: Optional[float] = Field(default=None, description="Prompt processing speed")
    
    # Additional provider-specific stats
    extra_stats: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Provider-specific additional statistics")
    
    def calculate_derived_stats(self) -> None:
        """Calculate derived statistics where possible"""
        # Calculate tokens per second for Ollama
        if self.eval_count and self.eval_duration and self.eval_duration > 0:
            # Convert nanoseconds to seconds and calculate tokens/sec
            duration_seconds = self.eval_duration / 1_000_000_000
            self.tokens_per_second = self.eval_count / duration_seconds
        
        # Calculate prompt processing speed for Ollama
        if self.prompt_eval_count and self.prompt_eval_duration and self.prompt_eval_duration > 0:
            duration_seconds = self.prompt_eval_duration / 1_000_000_000
            self.prompt_tokens_per_second = self.prompt_eval_count / duration_seconds
        
        # Standardize token counts across providers
        if not self.total_tokens and self.prompt_tokens and self.completion_tokens:
            self.total_tokens = self.prompt_tokens + self.completion_tokens
        
        # Map Ollama counts to standard format if not already set
        if not self.prompt_tokens and self.prompt_eval_count:
            self.prompt_tokens = self.prompt_eval_count
        if not self.completion_tokens and self.eval_count:
            self.completion_tokens = self.eval_count

# Embedding response models
class EmbeddingData(BaseModel):
    """Single embedding result"""
    embedding: List[float] = Field(description="The embedding vector")
    index: int = Field(description="Index in the input list")

class EmbeddingResponse(BaseModel):
    """Response from embeddings API"""
    data: List[EmbeddingData] = Field(description="List of embedding results")
    model: str = Field(description="Model used for embeddings")
    usage: Optional[UsageStats] = Field(default=None, description="Usage statistics")
    
    def get_single_embedding(self) -> List[float]:
        """Get the first embedding for single-text requests"""
        if not self.data:
            raise ValueError("No embeddings in response")
        return self.data[0].embedding
    
    def get_embeddings(self) -> List[List[float]]:
        """Get all embeddings as a list of vectors"""
        return [item.embedding for item in self.data]

class ChatResponse(BaseModel):
    content: str
    model: str
    finish_reason: Optional[str] = Field(default="")
    usage: Optional[UsageStats] = Field(default=None)
    # Keep legacy usage field for backward compatibility
    usage_legacy: Optional[Dict[str, int]] = Field(default=None, alias="usage_dict")
    
    def get_usage_dict(self) -> Dict[str, Any]:
        """Get usage statistics as dictionary for backward compatibility"""
        if self.usage:
            return self.usage.model_dump(exclude_none=True)
        return self.usage_legacy or {}

class LLMProvider(str, Enum):
    OLLAMA = "ollama"
    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    GEMINI = "gemini"
    GROK = "grok"

class BaseLLMAdapter(ABC):
    """Abstract base class for all LLM adapters"""
    
    def __init__(self, **config):
        self.config = config
    
    @abstractmethod
    async def chat(
        self, 
        model: str, 
        messages: List[LLMMessage], 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        """Send chat messages and get response"""
        pass
    
    @abstractmethod
    async def generate(
        self, 
        model: str, 
        prompt: str, 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        """Generate text from prompt"""
        pass
    
    @abstractmethod
    async def embeddings(
        self,
        model: str,
        input_texts: Union[str, List[str]],
        **kwargs
    ) -> EmbeddingResponse:
        """Generate embeddings for input text(s)"""
        pass
    
    @abstractmethod
    async def list_models(self) -> List[str]:
        """List available models"""
        pass

class OllamaAdapter(BaseLLMAdapter):
    """Adapter for Ollama with enhanced statistics"""
    
    def __init__(self, **config):
        super().__init__(**config)
        import ollama
        self.client = ollama.AsyncClient( # type: ignore
            host=config.get('host', defines.ollama_api_url)
        )
    
    def _create_usage_stats(self, response_data: Dict[str, Any]) -> UsageStats:
        """Create comprehensive usage statistics from Ollama response"""
        usage_stats = UsageStats()
        
        # Extract Ollama-specific stats
        if 'prompt_eval_count' in response_data:
            usage_stats.prompt_eval_count = response_data['prompt_eval_count']
        if 'prompt_eval_duration' in response_data:
            usage_stats.prompt_eval_duration = response_data['prompt_eval_duration']
        if 'eval_count' in response_data:
            usage_stats.eval_count = response_data['eval_count']
        if 'eval_duration' in response_data:
            usage_stats.eval_duration = response_data['eval_duration']
        if 'total_duration' in response_data:
            usage_stats.total_duration = response_data['total_duration']
        
        # Store any additional stats
        extra_stats = {}
        if type(response_data) is dict:
            for key, value in response_data.items():
                if key not in ['prompt_eval_count', 'prompt_eval_duration', 'eval_count', 
                            'eval_duration', 'total_duration', 'message', 'done_reason', 'response']:
                    extra_stats[key] = value
        
        if extra_stats:
            usage_stats.extra_stats = extra_stats
        
        # Calculate derived statistics
        usage_stats.calculate_derived_stats()
        
        return usage_stats
    
    async def chat(
        self, 
        model: str, 
        messages: List[LLMMessage], 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        # Convert LLMMessage objects to Ollama format
        ollama_messages = []
        for msg in messages:
            ollama_messages.append({
                "role": msg.role, 
                "content": msg.content
            })
        
        if stream:
            return self._stream_chat(model, ollama_messages, **kwargs)
        else:
            response = await self.client.chat(
                model=model, 
                messages=ollama_messages, 
                stream=False,
                **kwargs
            )
            
            usage_stats = self._create_usage_stats(response)
            
            return ChatResponse(
                content=response['message']['content'],
                model=model,
                finish_reason=response.get('done_reason'),
                usage=usage_stats
            )
    
    async def _stream_chat(self, model: str, messages: List[Dict], **kwargs):
        # Await the chat call first, then iterate over the result
        stream = await self.client.chat(
            model=model, 
            messages=messages, 
            stream=True,
            **kwargs
        )
        
        # Accumulate stats for final chunk
        accumulated_stats = {}
        
        async for chunk in stream:
            # Update accumulated stats
            accumulated_stats.update(chunk)
            content = chunk.get('message', {}).get('content', '')                
            usage_stats = None
            if chunk.done:
                usage_stats = self._create_usage_stats(accumulated_stats)
                yield ChatResponse(
                    content=content,
                    model=model,
                    finish_reason=chunk.get('done_reason'),
                    usage=usage_stats
                )
            else:
                yield ChatResponse(
                    content=content,
                    model=model,
                    finish_reason=None,
                    usage=None
                )
    
    async def generate(
        self, 
        model: str, 
        prompt: str, 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        if stream:
            return self._stream_generate(model, prompt, **kwargs)
        else:
            response = await self.client.generate(
                model=model, 
                prompt=prompt, 
                stream=False,
                **kwargs
            )
            
            usage_stats = self._create_usage_stats(response)
            
            return ChatResponse(
                content=response['response'],
                model=model,
                finish_reason=response.get('done_reason'),
                usage=usage_stats
            )
    
    async def _stream_generate(self, model: str, prompt: str, **kwargs):
        # Await the generate call first, then iterate over the result
        stream = await self.client.generate(
            model=model, 
            prompt=prompt, 
            stream=True,
            **kwargs
        )
        
        accumulated_stats = {}
        
        async for chunk in stream:
            if chunk.get('response'):
                accumulated_stats.update(chunk)
                
                usage_stats = None
                if chunk.get('done', False):  # Only create stats for final chunk
                    usage_stats = self._create_usage_stats(accumulated_stats)
                
                yield ChatResponse(
                    content=chunk['response'],
                    model=model,
                    finish_reason=chunk.get('done_reason'),
                    usage=usage_stats
                )

    async def embeddings(
        self,
        model: str,
        input_texts: Union[str, List[str]],
        **kwargs
    ) -> EmbeddingResponse:
        """Generate embeddings using Ollama"""
        # Normalize input to list
        if isinstance(input_texts, str):
            texts = [input_texts]
        else:
            texts = input_texts
        
        # Ollama embeddings API typically handles one text at a time
        results = []
        final_response = None
        
        for i, text in enumerate(texts):
            response = await self.client.embeddings(
                model=model,
                prompt=text,
                **kwargs
            )
            
            results.append(EmbeddingData(
                embedding=response['embedding'],
                index=i
            ))
            final_response = response
        
        # Create usage stats if available from the last response
        usage_stats = None
        if final_response and len(results) == 1:
            usage_stats = self._create_usage_stats(final_response)
        
        return EmbeddingResponse(
            data=results,
            model=model,
            usage=usage_stats
        )
    
    async def list_models(self) -> List[str]:
        models = await self.client.list()
        return [model['name'] for model in models['models']]

class OpenAIAdapter(BaseLLMAdapter):
    """Adapter for OpenAI with enhanced statistics"""
    
    def __init__(self, **config):
        super().__init__(**config)
        import openai # type: ignore
        self.client = openai.AsyncOpenAI(
            api_key=config.get('api_key', os.getenv('OPENAI_API_KEY'))
        )
    
    def _create_usage_stats(self, usage_data: Any) -> UsageStats:
        """Create usage statistics from OpenAI response"""
        if not usage_data:
            return UsageStats()
        
        usage_dict = usage_data.model_dump() if hasattr(usage_data, 'model_dump') else usage_data
        
        return UsageStats(
            prompt_tokens=usage_dict.get('prompt_tokens'),
            completion_tokens=usage_dict.get('completion_tokens'),
            total_tokens=usage_dict.get('total_tokens'),
            extra_stats={k: v for k, v in usage_dict.items() 
                        if k not in ['prompt_tokens', 'completion_tokens', 'total_tokens']}
        )
    
    async def chat(
        self, 
        model: str, 
        messages: List[LLMMessage], 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        # Convert LLMMessage objects to OpenAI format
        openai_messages = []
        for msg in messages:
            openai_messages.append({
                "role": msg.role, 
                "content": msg.content
            })
        
        if stream:
            return self._stream_chat(model, openai_messages, **kwargs)
        else:
            response = await self.client.chat.completions.create(
                model=model,
                messages=openai_messages,
                stream=False,
                **kwargs
            )
            
            usage_stats = self._create_usage_stats(response.usage)
            
            return ChatResponse(
                content=response.choices[0].message.content,
                model=model,
                finish_reason=response.choices[0].finish_reason,
                usage=usage_stats
            )
    
    async def _stream_chat(self, model: str, messages: List[Dict], **kwargs):
        # Await the stream creation first, then iterate
        stream = await self.client.chat.completions.create(
            model=model,
            messages=messages,
            stream=True,
            **kwargs
        )
        
        async for chunk in stream:
            if chunk.choices[0].delta.content:
                # Usage stats are only available in the final chunk for OpenAI streaming
                usage_stats = None
                if (chunk.choices[0].finish_reason is not None and 
                    hasattr(chunk, 'usage') and chunk.usage):
                    usage_stats = self._create_usage_stats(chunk.usage)
                
                yield ChatResponse(
                    content=chunk.choices[0].delta.content,
                    model=model,
                    finish_reason=chunk.choices[0].finish_reason,
                    usage=usage_stats
                )
    
    async def generate(
        self, 
        model: str, 
        prompt: str, 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        # Convert to chat format for OpenAI
        messages = [LLMMessage(role="user", content=prompt)]
        return await self.chat(model, messages, stream, **kwargs)

    async def embeddings(
        self,
        model: str,
        input_texts: Union[str, List[str]],
        **kwargs
    ) -> EmbeddingResponse:
        """Generate embeddings using OpenAI"""
        # Normalize input to list
        if isinstance(input_texts, str):
            texts = [input_texts]
        else:
            texts = input_texts
        
        response = await self.client.embeddings.create(
            model=model,
            input=texts,
            **kwargs
        )
        
        # Convert OpenAI response to our format
        results = []
        for item in response.data:
            results.append(EmbeddingData(
                embedding=item.embedding,
                index=item.index
            ))
        
        # Create usage stats
        usage_stats = self._create_usage_stats(response.usage)
        
        return EmbeddingResponse(
            data=results,
            model=model,
            usage=usage_stats
        )
    
    async def list_models(self) -> List[str]:
        models = await self.client.models.list()
        return [model.id for model in models.data]

class AnthropicAdapter(BaseLLMAdapter):
    """Adapter for Anthropic Claude with enhanced statistics"""
    
    def __init__(self, **config):
        super().__init__(**config)
        import anthropic # type: ignore
        self.client = anthropic.AsyncAnthropic(
            api_key=config.get('api_key', os.getenv('ANTHROPIC_API_KEY'))
        )
    
    def _create_usage_stats(self, usage_data: Any) -> UsageStats:
        """Create usage statistics from Anthropic response"""
        if not usage_data:
            return UsageStats()
        
        usage_dict = usage_data.model_dump() if hasattr(usage_data, 'model_dump') else usage_data
        
        return UsageStats(
            prompt_tokens=usage_dict.get('input_tokens'),
            completion_tokens=usage_dict.get('output_tokens'),
            total_tokens=(usage_dict.get('input_tokens', 0) + usage_dict.get('output_tokens', 0)) or None,
            extra_stats={k: v for k, v in usage_dict.items() 
                        if k not in ['input_tokens', 'output_tokens']}
        )
    
    async def chat(
        self, 
        model: str, 
        messages: List[LLMMessage], 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        # Anthropic requires system message to be separate
        system_message = None
        anthropic_messages = []
        
        for msg in messages:
            if msg.role == "system":
                system_message = msg.content
            else:
                anthropic_messages.append({
                    "role": msg.role, 
                    "content": msg.content
                })
        
        request_kwargs = {
            "model": model,
            "messages": anthropic_messages,
            "max_tokens": kwargs.pop('max_tokens', 1000),
            **kwargs
        }
        
        if system_message:
            request_kwargs["system"] = system_message
        
        if stream:
            return self._stream_chat(**request_kwargs)
        else:
            response = await self.client.messages.create(
                stream=False,
                **request_kwargs
            )
            
            usage_stats = self._create_usage_stats(response.usage)
            
            return ChatResponse(
                content=response.content[0].text,
                model=model,
                finish_reason=response.stop_reason,
                usage=usage_stats
            )
    
    async def _stream_chat(self, **kwargs):
        model = kwargs['model']
        
        async with self.client.messages.stream(**kwargs) as stream:
            final_usage_stats = None
            
            async for text in stream.text_stream:
                yield ChatResponse(
                    content=text,
                    model=model,
                    usage=None  # Usage stats not available until stream completion
                )
            
            # Collect usage stats after stream completion
            try:
                if hasattr(stream, 'get_final_message'):
                    final_message = await stream.get_final_message()
                    if final_message and hasattr(final_message, 'usage'):
                        final_usage_stats = self._create_usage_stats(final_message.usage)
                        
                        # Yield a final empty response with usage stats
                        yield ChatResponse(
                            content="",
                            model=model,
                            usage=final_usage_stats,
                            finish_reason="stop"
                        )
            except Exception as e:
                logger.debug(f"Could not retrieve final usage stats: {e}")
    
    async def generate(
        self, 
        model: str, 
        prompt: str, 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        messages = [LLMMessage(role="user", content=prompt)]
        return await self.chat(model, messages, stream, **kwargs)

    async def embeddings(
        self,
        model: str,
        input_texts: Union[str, List[str]],
        **kwargs
    ) -> EmbeddingResponse:
        """Anthropic doesn't provide embeddings API"""
        raise NotImplementedError(
            "Anthropic does not provide embeddings API. "
            "Consider using OpenAI, Ollama, or Gemini for embeddings."
        )    
    
    async def list_models(self) -> List[str]:
        # Anthropic doesn't have a list models endpoint, return known models
        return [
            "claude-3-5-sonnet-20241022",
            "claude-3-5-haiku-20241022", 
            "claude-3-opus-20240229"
        ]

class GeminiAdapter(BaseLLMAdapter):
    """Adapter for Google Gemini with enhanced statistics"""
    
    def __init__(self, **config):
        super().__init__(**config)
        import google.generativeai as genai # type: ignore
        genai.configure(api_key=config.get('api_key', os.getenv('GEMINI_API_KEY')))
        self.genai = genai
    
    def _create_usage_stats(self, response: Any) -> UsageStats:
        """Create usage statistics from Gemini response"""
        usage_stats = UsageStats()
        
        # Gemini usage metadata extraction
        if hasattr(response, 'usage_metadata'):
            usage = response.usage_metadata
            usage_stats.prompt_tokens = getattr(usage, 'prompt_token_count', None)
            usage_stats.completion_tokens = getattr(usage, 'candidates_token_count', None)
            usage_stats.total_tokens = getattr(usage, 'total_token_count', None)
            
            # Store additional Gemini-specific stats
            extra_stats = {}
            for attr in dir(usage):
                if not attr.startswith('_') and attr not in ['prompt_token_count', 'candidates_token_count', 'total_token_count']:
                    extra_stats[attr] = getattr(usage, attr)
            
            if extra_stats:
                usage_stats.extra_stats = extra_stats
        
        return usage_stats
    
    async def chat(
        self, 
        model: str, 
        messages: List[LLMMessage], 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        model_instance = self.genai.GenerativeModel(model)
        
        # Convert messages to Gemini format
        chat_history = []
        current_message = None
        
        for i, msg in enumerate(messages[:-1]):  # All but last message go to history
            if msg.role == "user":
                chat_history.append({"role": "user", "parts": [msg.content]})
            elif msg.role == "assistant":
                chat_history.append({"role": "model", "parts": [msg.content]})
        
        # Last message is the current prompt
        if messages:
            current_message = messages[-1].content
        
        chat = model_instance.start_chat(history=chat_history)
        if not current_message:
            raise ValueError("No current message provided for chat")
        
        if stream:
            return self._stream_chat(chat, current_message, **kwargs)
        else:
            response = await chat.send_message_async(current_message, **kwargs)
            usage_stats = self._create_usage_stats(response)
            
            return ChatResponse(
                content=response.text,
                model=model,
                finish_reason=response.candidates[0].finish_reason.name if response.candidates else None,
                usage=usage_stats
            )
    
    async def _stream_chat(self, chat, message: str, **kwargs):
        # Await the stream creation first, then iterate
        stream = await chat.send_message_async(message, stream=True, **kwargs)
        
        final_chunk = None
        
        async for chunk in stream:
            if chunk.text:
                final_chunk = chunk  # Keep reference to last chunk
                
                yield ChatResponse(
                    content=chunk.text,
                    model=chat.model.model_name,
                    usage=None  # Don't include usage stats until final chunk
                )
        
        # After streaming is complete, yield final response with usage stats
        if final_chunk:
            usage_stats = self._create_usage_stats(final_chunk)
            if usage_stats and any([usage_stats.prompt_tokens, usage_stats.completion_tokens, usage_stats.total_tokens]):
                yield ChatResponse(
                    content="",
                    model=chat.model.model_name,
                    usage=usage_stats,
                    finish_reason=final_chunk.candidates[0].finish_reason.name if final_chunk.candidates else None
                )
    
    async def generate(
        self, 
        model: str, 
        prompt: str, 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        messages = [LLMMessage(role="user", content=prompt)]
        return await self.chat(model, messages, stream, **kwargs)

    async def embeddings(
        self,
        model: str,
        input_texts: Union[str, List[str]],
        **kwargs
    ) -> EmbeddingResponse:
        """Generate embeddings using Google Gemini"""
        # Normalize input to list
        if isinstance(input_texts, str):
            texts = [input_texts]
        else:
            texts = input_texts
        
        results = []
        
        # Gemini embeddings - use embedding model
        for i, text in enumerate(texts):
            response = self.genai.embed_content(
                model=model,
                content=text,
                **kwargs
            )
            
            results.append(EmbeddingData(
                embedding=response['embedding'],
                index=i
            ))
        
        return EmbeddingResponse(
            data=results,
            model=model,
            usage=None  # Gemini embeddings don't typically return usage stats
        )

    async def list_models(self) -> List[str]:
        models = self.genai.list_models()
        return [model.name for model in models if 'generateContent' in model.supported_generation_methods]

class UnifiedLLMProxy:
    """Main proxy class that provides unified interface to all LLM providers"""
    
    def __init__(self, default_provider: LLMProvider = LLMProvider.OLLAMA):
        self.adapters: Dict[LLMProvider, BaseLLMAdapter] = {}
        self.default_provider = default_provider
        self._initialized_providers = set()
    
    def configure_provider(self, provider: LLMProvider, **config):
        """Configure a specific provider with its settings"""
        adapter_classes = {
            LLMProvider.OLLAMA: OllamaAdapter,
            LLMProvider.OPENAI: OpenAIAdapter,
            LLMProvider.ANTHROPIC: AnthropicAdapter,
            LLMProvider.GEMINI: GeminiAdapter,
            # Add other providers as needed
        }
        
        if provider in adapter_classes:
            self.adapters[provider] = adapter_classes[provider](**config)
            self._initialized_providers.add(provider)
        else:
            raise ValueError(f"Unsupported provider: {provider}")
    
    def set_default_provider(self, provider: LLMProvider):
        """Set the default provider for requests"""
        if provider not in self._initialized_providers:
            raise ValueError(f"Provider {provider} not configured")
        self.default_provider = provider
    
    async def chat(
        self, 
        model: str, 
        messages: Union[List[LLMMessage], List[Dict[str, str]]], 
        provider: Optional[LLMProvider] = None,
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        """Send chat messages using specified or default provider"""
        
        provider = provider or self.default_provider
        adapter = self._get_adapter(provider)
        
        # Normalize messages to LLMMessage objects
        normalized_messages = []
        for msg in messages:
            if isinstance(msg, LLMMessage):
                normalized_messages.append(msg)
            elif isinstance(msg, dict):
                normalized_messages.append(LLMMessage.from_dict(msg))
            else:
                raise ValueError(f"Invalid message type: {type(msg)}")
        
        return await adapter.chat(model, normalized_messages, stream, **kwargs)
    
    async def chat_stream(
        self, 
        model: str, 
        messages: Union[List[LLMMessage], List[Dict[str, str]]], 
        provider: Optional[LLMProvider] = None,
        stream: bool = True,
        **kwargs
    ) -> AsyncGenerator[ChatResponse, None]:
        """Stream chat messages using specified or default provider"""
        if stream is False:
            raise ValueError("stream must be True for chat_stream")
        result = await self.chat(model, messages, provider, stream=True, **kwargs)
        # Type checker now knows this is an AsyncGenerator due to stream=True
        async for chunk in result:  # type: ignore
            yield chunk
    
    async def chat_single(
        self, 
        model: str, 
        messages: Union[List[LLMMessage], List[Dict[str, str]]], 
        provider: Optional[LLMProvider] = None,
        **kwargs
    ) -> ChatResponse:
        """Get single chat response using specified or default provider"""
        
        result = await self.chat(model, messages, provider, stream=False, **kwargs)
        # Type checker now knows this is a ChatResponse due to stream=False
        return result  # type: ignore
    
    async def generate(
        self, 
        model: str, 
        prompt: str,
        provider: Optional[LLMProvider] = None,
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        """Generate text using specified or default provider"""
        
        provider = provider or self.default_provider
        adapter = self._get_adapter(provider)
        
        return await adapter.generate(model, prompt, stream, **kwargs)
    
    async def generate_stream(
        self, 
        model: str, 
        prompt: str,
        provider: Optional[LLMProvider] = None,
        **kwargs
    ) -> AsyncGenerator[ChatResponse, None]:
        """Stream text generation using specified or default provider"""
        
        result = await self.generate(model, prompt, provider, stream=True, **kwargs)
        async for chunk in result:  # type: ignore
            yield chunk
    
    async def generate_single(
        self, 
        model: str, 
        prompt: str,
        provider: Optional[LLMProvider] = None,
        **kwargs
    ) -> ChatResponse:
        """Get single generation response using specified or default provider"""
        
        result = await self.generate(model, prompt, provider, stream=False, **kwargs)
        return result  # type: ignore
    
    async def embeddings(
        self,
        model: str,
        input_texts: Union[str, List[str]],
        provider: Optional[LLMProvider] = None,
        **kwargs
    ) -> EmbeddingResponse:
        """Generate embeddings using specified or default provider"""
        provider = provider or self.default_provider
        adapter = self._get_adapter(provider)
        
        return await adapter.embeddings(model, input_texts, **kwargs)
    
    async def list_models(self, provider: Optional[LLMProvider] = None) -> List[str]:
        """List available models for specified or default provider"""
        
        provider = provider or self.default_provider
        adapter = self._get_adapter(provider)
        
        return await adapter.list_models()
    
    async def list_embedding_models(self, provider: Optional[LLMProvider] = None) -> List[str]:
        """List available embedding models for specified or default provider"""
        provider = provider or self.default_provider
        
        # Provider-specific embedding models
        embedding_models = {
            LLMProvider.OLLAMA: [
                "nomic-embed-text",
                "mxbai-embed-large", 
                "all-minilm",
                "snowflake-arctic-embed"
            ],
            LLMProvider.OPENAI: [
                "text-embedding-3-small",
                "text-embedding-3-large", 
                "text-embedding-ada-002"
            ],
            LLMProvider.ANTHROPIC: [],  # No embeddings API
            LLMProvider.GEMINI: [
                "models/embedding-001",
                "models/text-embedding-004"
            ]
        }
        
        if provider == LLMProvider.ANTHROPIC:
            raise NotImplementedError("Anthropic does not provide embeddings API")
        
        # For Ollama, check which embedding models are actually available
        if provider == LLMProvider.OLLAMA:
            try:
                all_models = await self.list_models(provider)
                available_embedding_models = []
                suggested_models = embedding_models[provider]
                
                for model in all_models:
                    # Check if it's a known embedding model
                    if any(emb_model in model for emb_model in suggested_models):
                        available_embedding_models.append(model)
                
                return available_embedding_models if available_embedding_models else suggested_models
            except:
                return embedding_models[provider]
        
        return embedding_models.get(provider, [])
    
    def _get_adapter(self, provider: LLMProvider) -> BaseLLMAdapter:
        """Get adapter for specified provider"""
        if provider not in self.adapters:
            raise ValueError(f"Provider {provider} not configured")
        return self.adapters[provider]

# Example usage and configuration
class LLMManager:
    """Singleton manager for the unified LLM proxy"""
    
    _instance = None
    _proxy = None
    
    @classmethod
    def get_instance(cls):
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance
    
    def __init__(self):
        if LLMManager._proxy is None:
            LLMManager._proxy = UnifiedLLMProxy()
            self._configure_from_environment()
    
    def _configure_from_environment(self):
        """Configure providers based on environment variables"""
        if not self._proxy:
            raise RuntimeError("UnifiedLLMProxy instance not initialized")
        # Configure Ollama if available
        ollama_host = os.getenv('OLLAMA_HOST', defines.ollama_api_url)
        self._proxy.configure_provider(LLMProvider.OLLAMA, host=ollama_host)
        
        # Configure OpenAI if API key is available
        if os.getenv('OPENAI_API_KEY'):
            self._proxy.configure_provider(LLMProvider.OPENAI)
        
        # Configure Anthropic if API key is available
        if os.getenv('ANTHROPIC_API_KEY'):
            self._proxy.configure_provider(LLMProvider.ANTHROPIC)
        
        # Configure Gemini if API key is available
        if os.getenv('GEMINI_API_KEY'):
            self._proxy.configure_provider(LLMProvider.GEMINI)
        
        # Set default provider from environment or use Ollama
        default_provider = os.getenv('DEFAULT_LLM_PROVIDER', 'ollama')
        try:
            self._proxy.set_default_provider(LLMProvider(default_provider))
        except ValueError:
            # Fallback to Ollama if specified provider not available
            self._proxy.set_default_provider(LLMProvider.OLLAMA)
    
    def get_proxy(self) -> UnifiedLLMProxy:
        """Get the unified LLM proxy instance"""
        if not self._proxy:
            raise RuntimeError("UnifiedLLMProxy instance not initialized")
        return self._proxy

# Convenience function for easy access
def get_llm() -> UnifiedLLMProxy:
    """Get the configured LLM proxy"""
    return LLMManager.get_instance().get_proxy()

# Example usage with detailed statistics
async def example_usage():
    """Example showing how to access detailed statistics"""
    llm = get_llm()
    
    # Configure providers
    llm.configure_provider(LLMProvider.OLLAMA, host="http://localhost:11434")
    
    # Simple chat
    messages = [
        LLMMessage(role="user", content="Explain quantum computing in one paragraph")
    ]
    
    response = await llm.chat_single("llama2", messages)
    
    print(f"Content: {response.content}")
    print(f"Model: {response.model}")
    
    if response.usage:
        print(f"Usage Statistics:")
        print(f"  Prompt tokens: {response.usage.prompt_tokens}")
        print(f"  Completion tokens: {response.usage.completion_tokens}")
        print(f"  Total tokens: {response.usage.total_tokens}")
        
        # Ollama-specific stats
        if response.usage.prompt_eval_count:
            print(f"  Prompt eval count: {response.usage.prompt_eval_count}")
        if response.usage.eval_count:
            print(f"  Eval count: {response.usage.eval_count}")
        if response.usage.tokens_per_second:
            print(f"  Generation speed: {response.usage.tokens_per_second:.2f} tokens/sec")
        if response.usage.prompt_tokens_per_second:
            print(f"  Prompt processing speed: {response.usage.prompt_tokens_per_second:.2f} tokens/sec")
        
        # Access as dictionary for backward compatibility
        usage_dict = response.get_usage_dict()
        print(f"Usage as dict: {usage_dict}")

async def example_embeddings_usage():
    """Example showing how to use the embeddings API"""
    llm = get_llm()
    
    # Configure providers
    llm.configure_provider(LLMProvider.OLLAMA, host="http://localhost:11434")
    if os.getenv('OPENAI_API_KEY'):
        llm.configure_provider(LLMProvider.OPENAI)
    
    # List available embedding models
    print("=== Available Embedding Models ===")
    try:
        ollama_embedding_models = await llm.list_embedding_models(LLMProvider.OLLAMA)
        print(f"Ollama embedding models: {ollama_embedding_models}")
    except Exception as e:
        print(f"Could not list Ollama embedding models: {e}")
    
    if os.getenv('OPENAI_API_KEY'):
        try:
            openai_embedding_models = await llm.list_embedding_models(LLMProvider.OPENAI)
            print(f"OpenAI embedding models: {openai_embedding_models}")
        except Exception as e:
            print(f"Could not list OpenAI embedding models: {e}")
    
    # Single text embedding
    print("\n=== Single Text Embedding ===")
    single_text = "The quick brown fox jumps over the lazy dog"
    
    try:
        response = await llm.embeddings("nomic-embed-text", single_text, provider=LLMProvider.OLLAMA)
        print(f"Model: {response.model}")
        print(f"Embedding dimension: {len(response.get_single_embedding())}")
        print(f"First 5 values: {response.get_single_embedding()[:5]}")
        
        if response.usage:
            print(f"Usage: {response.usage.model_dump(exclude_none=True)}")
    except Exception as e:
        print(f"Ollama embedding failed: {e}")
    
    # Batch embeddings
    print("\n=== Batch Text Embeddings ===")
    texts = [
        "Machine learning is fascinating",
        "Natural language processing enables AI communication", 
        "Deep learning uses neural networks",
        "Transformers revolutionized AI"
    ]
    
    try:
        # Try OpenAI if available
        if os.getenv('OPENAI_API_KEY'):
            response = await llm.embeddings("text-embedding-3-small", texts, provider=LLMProvider.OPENAI)
            print(f"OpenAI Model: {response.model}")
            print(f"Number of embeddings: {len(response.data)}")
            print(f"Embedding dimension: {len(response.data[0].embedding)}")
            
            # Calculate similarity between first two texts (requires numpy)
            try:
                import numpy as np # type: ignore
                emb1 = np.array(response.data[0].embedding)
                emb2 = np.array(response.data[1].embedding)
                similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
                print(f"Cosine similarity between first two texts: {similarity:.4f}")
            except ImportError:
                print("Install numpy to calculate similarity: pip install numpy")
            
            if response.usage:
                print(f"Usage: {response.usage.model_dump(exclude_none=True)}")
        else:
            # Fallback to Ollama for batch
            response = await llm.embeddings("nomic-embed-text", texts, provider=LLMProvider.OLLAMA)
            print(f"Ollama Model: {response.model}")
            print(f"Number of embeddings: {len(response.data)}")
            print(f"Embedding dimension: {len(response.data[0].embedding)}")
            
    except Exception as e:
        print(f"Batch embedding failed: {e}")

async def example_streaming_with_stats():
    """Example showing how to collect usage stats from streaming responses"""
    llm = get_llm()
    
    messages = [
        LLMMessage(role="user", content="Write a short story about AI")
    ]
    
    print("Streaming response:")
    final_stats = None
    
    async for chunk in llm.chat_stream("llama2", messages):
        # Print content as it streams
        if chunk.content:
            print(chunk.content, end="", flush=True)
        
        # Collect usage stats from final chunk
        if chunk.usage:
            final_stats = chunk.usage
    
    print("\n\nFinal Usage Statistics:")
    if final_stats:
        print(f"  Total tokens: {final_stats.total_tokens}")
        print(f"  Generation speed: {final_stats.tokens_per_second:.2f} tokens/sec" if final_stats.tokens_per_second else "  Generation speed: N/A")
    else:
        print("  No usage statistics available")

if __name__ == "__main__":
    asyncio.run(example_usage())
    print("\n" + "="*50 + "\n")
    asyncio.run(example_streaming_with_stats())
    print("\n" + "="*50 + "\n")
    asyncio.run(example_embeddings_usage())