from abc import ABC, abstractmethod
from typing import Dict, List, Any, AsyncGenerator, Optional, Union
from enum import Enum
import asyncio
import json
from dataclasses import dataclass
import os

# Standard message format for all providers
@dataclass
class ChatMessage:
    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 ChatMessage")
    
    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 ChatMessage")
    
    @classmethod
    def from_dict(cls, data: Dict[str, str]) -> 'ChatMessage':
        """Create ChatMessage from dictionary"""
        return cls(role=data['role'], content=data['content'])
    
    def to_dict(self) -> Dict[str, str]:
        """Convert ChatMessage to dictionary"""
        return {'role': self.role, 'content': self.content}

@dataclass
class ChatResponse:
    content: str
    model: str
    finish_reason: Optional[str] = None
    usage: Optional[Dict[str, int]] = None

class LLMProvider(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[ChatMessage], 
        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 list_models(self) -> List[str]:
        """List available models"""
        pass

class OllamaAdapter(BaseLLMAdapter):
    """Adapter for Ollama"""
    
    def __init__(self, **config):
        super().__init__(**config)
        import ollama
        self.client = ollama.AsyncClient( # type: ignore
            host=config.get('host', 'http://localhost:11434')
        )
    
    async def chat(
        self, 
        model: str, 
        messages: List[ChatMessage], 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        # Convert ChatMessage 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
            )
            return ChatResponse(
                content=response['message']['content'],
                model=model,
                finish_reason=response.get('done_reason')
            )
    
    async def _stream_chat(self, model: str, messages: List[Dict], **kwargs):
        async for chunk in self.client.chat(
            model=model, 
            messages=messages, 
            stream=True,
            **kwargs
        ):
            if chunk.get('message', {}).get('content'):
                yield ChatResponse(
                    content=chunk['message']['content'],
                    model=model,
                    finish_reason=chunk.get('done_reason')
                )
    
    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
            )
            return ChatResponse(
                content=response['response'],
                model=model,
                finish_reason=response.get('done_reason')
            )
    
    async def _stream_generate(self, model: str, prompt: str, **kwargs):
        async for chunk in self.client.generate(
            model=model, 
            prompt=prompt, 
            stream=True,
            **kwargs
        ):
            if chunk.get('response'):
                yield ChatResponse(
                    content=chunk['response'],
                    model=model,
                    finish_reason=chunk.get('done_reason')
                )
    
    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"""
    
    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'))
        )
    
    async def chat(
        self, 
        model: str, 
        messages: List[ChatMessage], 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        # Convert ChatMessage 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
            )
            return ChatResponse(
                content=response.choices[0].message.content,
                model=model,
                finish_reason=response.choices[0].finish_reason,
                usage=response.usage.model_dump() if response.usage else None
            )
    
    async def _stream_chat(self, model: str, messages: List[Dict], **kwargs):
        async for chunk in await self.client.chat.completions.create(
            model=model,
            messages=messages,
            stream=True,
            **kwargs
        ):
            if chunk.choices[0].delta.content:
                yield ChatResponse(
                    content=chunk.choices[0].delta.content,
                    model=model,
                    finish_reason=chunk.choices[0].finish_reason
                )
    
    async def generate(
        self, 
        model: str, 
        prompt: str, 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        # Convert to chat format for OpenAI
        messages = [ChatMessage(role="user", content=prompt)]
        return await self.chat(model, messages, stream, **kwargs)
    
    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"""
    
    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'))
        )
    
    async def chat(
        self, 
        model: str, 
        messages: List[ChatMessage], 
        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
            )
            return ChatResponse(
                content=response.content[0].text,
                model=model,
                finish_reason=response.stop_reason,
                usage={
                    "input_tokens": response.usage.input_tokens,
                    "output_tokens": response.usage.output_tokens
                }
            )
    
    async def _stream_chat(self, **kwargs):
        async with self.client.messages.stream(**kwargs) as stream:
            async for text in stream.text_stream:
                yield ChatResponse(
                    content=text,
                    model=kwargs['model']
                )
    
    async def generate(
        self, 
        model: str, 
        prompt: str, 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        messages = [ChatMessage(role="user", content=prompt)]
        return await self.chat(model, messages, stream, **kwargs)
    
    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"""
    
    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
    
    async def chat(
        self, 
        model: str, 
        messages: List[ChatMessage], 
        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)
            return ChatResponse(
                content=response.text,
                model=model,
                finish_reason=response.candidates[0].finish_reason.name if response.candidates else None
            )
    
    async def _stream_chat(self, chat, message: str, **kwargs):
        async for chunk in chat.send_message_async(message, stream=True, **kwargs):
            if chunk.text:
                yield ChatResponse(
                    content=chunk.text,
                    model=chat.model.model_name
                )
    
    async def generate(
        self, 
        model: str, 
        prompt: str, 
        stream: bool = False,
        **kwargs
    ) -> Union[ChatResponse, AsyncGenerator[ChatResponse, None]]:
        
        messages = [ChatMessage(role="user", content=prompt)]
        return await self.chat(model, messages, stream, **kwargs)
    
    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[ChatMessage], 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 ChatMessage objects
        normalized_messages = []
        for msg in messages:
            if isinstance(msg, ChatMessage):
                normalized_messages.append(msg)
            elif isinstance(msg, dict):
                normalized_messages.append(ChatMessage.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[ChatMessage], List[Dict[str, str]]], 
        provider: Optional[LLMProvider] = None,
        **kwargs
    ) -> AsyncGenerator[ChatResponse, None]:
        """Stream chat messages using specified or default provider"""
        
        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[ChatMessage], 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 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()
    
    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("LLM proxy not initialized")

        # Configure Ollama if available
        ollama_host = os.getenv('OLLAMA_HOST', 'http://localhost:11434')
        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("LLM proxy 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()