"""
LLM client abstraction for model orchestration.

Supports multiple backends:
- HuggingFace Transformers
- vLLM
- Ollama
- OpenAI-compatible APIs
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


@dataclass
class GenerationConfig:
    """Configuration for text generation."""
    max_tokens: int = 2048
    temperature: float = 0.7
    top_p: float = 0.9
    top_k: int = 50
    repetition_penalty: float = 1.1
    stop_sequences: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
            "top_p": self.top_p,
            "top_k": self.top_k,
            "repetition_penalty": self.repetition_penalty,
            "stop_sequences": self.stop_sequences
        }


@dataclass
class GenerationResult:
    """Result from LLM generation."""
    text: str
    tokens_used: int = 0
    finish_reason: str = "stop"
    model: str = ""
    metadata: Dict[str, Any] = field(default_factory=dict)


class LLMClient(ABC):
    """Abstract base class for LLM clients."""
    
    @abstractmethod
    def generate(
        self,
        prompt: str,
        config: Optional[GenerationConfig] = None
    ) -> GenerationResult:
        """
        Generate text from prompt.
        
        Args:
            prompt: Input prompt
            config: Generation configuration
            
        Returns:
            GenerationResult
        """
        pass
    
    @abstractmethod
    def generate_batch(
        self,
        prompts: List[str],
        config: Optional[GenerationConfig] = None
    ) -> List[GenerationResult]:
        """
        Generate text for multiple prompts.
        
        Args:
            prompts: List of prompts
            config: Generation configuration
            
        Returns:
            List of GenerationResults
        """
        pass
    
    def chat(
        self,
        messages: List[Dict[str, str]],
        config: Optional[GenerationConfig] = None
    ) -> GenerationResult:
        """
        Chat-style generation with message history.
        
        Args:
            messages: List of message dicts with 'role' and 'content'
            config: Generation configuration
            
        Returns:
            GenerationResult
        """
        # Default implementation: format messages as prompt
        prompt = self._format_chat_messages(messages)
        return self.generate(prompt, config)
    
    def _format_chat_messages(self, messages: List[Dict[str, str]]) -> str:
        """Format chat messages as a prompt."""
        formatted = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            
            if role == "system":
                formatted.append(f"System: {content}")
            elif role == "assistant":
                formatted.append(f"Assistant: {content}")
            else:
                formatted.append(f"User: {content}")
        
        formatted.append("Assistant:")
        return "\n\n".join(formatted)


class HuggingFaceLLM(LLMClient):
    """HuggingFace Transformers LLM client."""
    
    def __init__(
        self,
        model_name: str = "mistralai/Mistral-7B-Instruct-v0.2",
        device: str = "auto",
        quantization: Optional[str] = "4bit",
        trust_remote_code: bool = True,
        torch_dtype: str = "auto"
    ):
        """
        Initialize HuggingFace LLM.
        
        Args:
            model_name: Model name or path
            device: Device to use
            quantization: Quantization mode (None, "4bit", "8bit")
            trust_remote_code: Trust remote code
            torch_dtype: Torch dtype
        """
        self.model_name = model_name
        self.device = device
        self.quantization = quantization
        self.trust_remote_code = trust_remote_code
        self.torch_dtype = torch_dtype
        
        self._model = None
        self._tokenizer = None
    
    def _load_model(self):
        """Lazy load the model."""
        if self._model is not None:
            return
        
        logger.info(f"Loading HuggingFace model: {self.model_name}")
        
        try:
            import torch
            from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
        except ImportError:
            raise ImportError(
                "transformers and torch are required. "
                "Install with: pip install transformers torch"
            )
        
        # Resolve device
        if self.device == "auto":
            if torch.cuda.is_available():
                device = "cuda"
            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
                device = "mps"
            else:
                device = "cpu"
        else:
            device = self.device
        
        # Resolve dtype
        if self.torch_dtype == "auto":
            dtype = torch.float16 if device != "cpu" else torch.float32
        else:
            dtype = getattr(torch, self.torch_dtype, torch.float32)
        
        # Load tokenizer
        self._tokenizer = AutoTokenizer.from_pretrained(
            self.model_name,
            trust_remote_code=self.trust_remote_code
        )
        
        # Configure quantization
        model_kwargs = {
            "trust_remote_code": self.trust_remote_code,
            "torch_dtype": dtype,
        }
        
        if self.quantization and device == "cuda":
            try:
                if self.quantization == "4bit":
                    model_kwargs["quantization_config"] = BitsAndBytesConfig(
                        load_in_4bit=True,
                        bnb_4bit_compute_dtype=dtype
                    )
                elif self.quantization == "8bit":
                    model_kwargs["quantization_config"] = BitsAndBytesConfig(
                        load_in_8bit=True
                    )
            except Exception as e:
                logger.warning(f"Quantization not available: {e}")
        else:
            model_kwargs["device_map"] = device if device != "cpu" else None
        
        # Load model
        self._model = AutoModelForCausalLM.from_pretrained(
            self.model_name,
            **model_kwargs
        )
        
        if device == "cpu":
            self._model = self._model.to(device)
        
        logger.info(f"Loaded model on {device}")
    
    @trace_function(name="hf_generate")
    def generate(
        self,
        prompt: str,
        config: Optional[GenerationConfig] = None
    ) -> GenerationResult:
        """Generate text using HuggingFace model."""
        self._load_model()
        
        config = config or GenerationConfig()
        
        # Tokenize
        inputs = self._tokenizer(prompt, return_tensors="pt")
        inputs = {k: v.to(self._model.device) for k, v in inputs.items()}
        input_length = inputs["input_ids"].shape[1]
        
        # Generate
        import torch
        
        with torch.no_grad():
            outputs = self._model.generate(
                **inputs,
                max_new_tokens=config.max_tokens,
                temperature=config.temperature,
                top_p=config.top_p,
                top_k=config.top_k,
                repetition_penalty=config.repetition_penalty,
                do_sample=config.temperature > 0,
                pad_token_id=self._tokenizer.eos_token_id
            )
        
        # Decode
        generated_tokens = outputs[0][input_length:]
        text = self._tokenizer.decode(generated_tokens, skip_special_tokens=True)
        
        return GenerationResult(
            text=text.strip(),
            tokens_used=len(generated_tokens),
            model=self.model_name
        )
    
    def generate_batch(
        self,
        prompts: List[str],
        config: Optional[GenerationConfig] = None
    ) -> List[GenerationResult]:
        """Generate text for multiple prompts."""
        return [self.generate(p, config) for p in prompts]


class OllamaLLM(LLMClient):
    """Ollama LLM client."""
    
    def __init__(
        self,
        model_name: str = "mistral",
        host: str = "http://localhost:11434"
    ):
        """
        Initialize Ollama client.
        
        Args:
            model_name: Ollama model name
            host: Ollama server URL
        """
        self.model_name = model_name
        self.host = host.rstrip("/")
    
    @trace_function(name="ollama_generate")
    def generate(
        self,
        prompt: str,
        config: Optional[GenerationConfig] = None
    ) -> GenerationResult:
        """Generate text using Ollama."""
        import requests
        
        config = config or GenerationConfig()
        
        response = requests.post(
            f"{self.host}/api/generate",
            json={
                "model": self.model_name,
                "prompt": prompt,
                "options": {
                    "num_predict": config.max_tokens,
                    "temperature": config.temperature,
                    "top_p": config.top_p,
                    "top_k": config.top_k,
                    "repeat_penalty": config.repetition_penalty
                },
                "stream": False
            }
        )
        response.raise_for_status()
        
        data = response.json()
        
        return GenerationResult(
            text=data.get("response", "").strip(),
            tokens_used=data.get("eval_count", 0),
            model=self.model_name,
            metadata={"total_duration": data.get("total_duration")}
        )
    
    def generate_batch(
        self,
        prompts: List[str],
        config: Optional[GenerationConfig] = None
    ) -> List[GenerationResult]:
        """Generate text for multiple prompts."""
        return [self.generate(p, config) for p in prompts]
    
    def chat(
        self,
        messages: List[Dict[str, str]],
        config: Optional[GenerationConfig] = None
    ) -> GenerationResult:
        """Chat using Ollama's chat API."""
        import requests
        
        config = config or GenerationConfig()
        
        response = requests.post(
            f"{self.host}/api/chat",
            json={
                "model": self.model_name,
                "messages": messages,
                "options": {
                    "num_predict": config.max_tokens,
                    "temperature": config.temperature,
                    "top_p": config.top_p
                },
                "stream": False
            }
        )
        response.raise_for_status()
        
        data = response.json()
        
        return GenerationResult(
            text=data.get("message", {}).get("content", "").strip(),
            tokens_used=data.get("eval_count", 0),
            model=self.model_name
        )


class OpenAICompatibleLLM(LLMClient):
    """OpenAI-compatible API client."""
    
    def __init__(
        self,
        model_name: str = "gpt-3.5-turbo",
        api_base: Optional[str] = None,
        api_key: Optional[str] = None
    ):
        """
        Initialize OpenAI-compatible client.
        
        Args:
            model_name: Model name
            api_base: API base URL (None for OpenAI)
            api_key: API key
        """
        self.model_name = model_name
        self.api_base = api_base
        self.api_key = api_key
        self._client = None
    
    def _get_client(self):
        """Get or create OpenAI client."""
        if self._client is not None:
            return self._client
        
        try:
            from openai import OpenAI
        except ImportError:
            raise ImportError("openai is required. Install with: pip install openai")
        
        kwargs = {}
        if self.api_base:
            kwargs["base_url"] = self.api_base
        if self.api_key:
            kwargs["api_key"] = self.api_key
        
        self._client = OpenAI(**kwargs)
        return self._client
    
    @trace_function(name="openai_generate")
    def generate(
        self,
        prompt: str,
        config: Optional[GenerationConfig] = None
    ) -> GenerationResult:
        """Generate text using OpenAI-compatible API."""
        messages = [{"role": "user", "content": prompt}]
        return self.chat(messages, config)
    
    def generate_batch(
        self,
        prompts: List[str],
        config: Optional[GenerationConfig] = None
    ) -> List[GenerationResult]:
        """Generate text for multiple prompts."""
        return [self.generate(p, config) for p in prompts]
    
    def chat(
        self,
        messages: List[Dict[str, str]],
        config: Optional[GenerationConfig] = None
    ) -> GenerationResult:
        """Chat using OpenAI-compatible API."""
        client = self._get_client()
        config = config or GenerationConfig()
        
        response = client.chat.completions.create(
            model=self.model_name,
            messages=messages,
            max_tokens=config.max_tokens,
            temperature=config.temperature,
            top_p=config.top_p
        )
        
        choice = response.choices[0]
        
        return GenerationResult(
            text=choice.message.content.strip(),
            tokens_used=response.usage.total_tokens if response.usage else 0,
            finish_reason=choice.finish_reason,
            model=self.model_name
        )


class VLLMClient(LLMClient):
    """vLLM client for high-throughput inference."""
    
    def __init__(
        self,
        model_name: str = "mistralai/Mistral-7B-Instruct-v0.2",
        tensor_parallel_size: int = 1
    ):
        """
        Initialize vLLM client.
        
        Args:
            model_name: Model name
            tensor_parallel_size: Number of GPUs for tensor parallelism
        """
        self.model_name = model_name
        self.tensor_parallel_size = tensor_parallel_size
        self._llm = None
    
    def _load_model(self):
        """Lazy load vLLM model."""
        if self._llm is not None:
            return
        
        try:
            from vllm import LLM
        except ImportError:
            raise ImportError("vllm is required. Install with: pip install vllm")
        
        logger.info(f"Loading vLLM model: {self.model_name}")
        
        self._llm = LLM(
            model=self.model_name,
            tensor_parallel_size=self.tensor_parallel_size
        )
    
    @trace_function(name="vllm_generate")
    def generate(
        self,
        prompt: str,
        config: Optional[GenerationConfig] = None
    ) -> GenerationResult:
        """Generate text using vLLM."""
        results = self.generate_batch([prompt], config)
        return results[0]
    
    def generate_batch(
        self,
        prompts: List[str],
        config: Optional[GenerationConfig] = None
    ) -> List[GenerationResult]:
        """Generate text for multiple prompts using vLLM."""
        self._load_model()
        
        try:
            from vllm import SamplingParams
        except ImportError:
            raise ImportError("vllm is required")
        
        config = config or GenerationConfig()
        
        sampling_params = SamplingParams(
            max_tokens=config.max_tokens,
            temperature=config.temperature,
            top_p=config.top_p,
            top_k=config.top_k,
            repetition_penalty=config.repetition_penalty
        )
        
        outputs = self._llm.generate(prompts, sampling_params)
        
        results = []
        for output in outputs:
            text = output.outputs[0].text
            results.append(GenerationResult(
                text=text.strip(),
                tokens_used=len(output.outputs[0].token_ids),
                model=self.model_name
            ))
        
        return results


def create_llm_client(settings) -> LLMClient:
    """
    Factory function to create LLM client from settings.
    
    Args:
        settings: LLMSettings or Settings object
        
    Returns:
        LLMClient instance
    """
    if hasattr(settings, "llm"):
        llm_settings = settings.llm
    else:
        llm_settings = settings
    
    backend = llm_settings.backend
    
    if backend == "huggingface":
        return HuggingFaceLLM(
            model_name=llm_settings.model_name,
            device=llm_settings.device,
            quantization=llm_settings.quantization,
            trust_remote_code=llm_settings.huggingface.trust_remote_code,
            torch_dtype=llm_settings.huggingface.torch_dtype
        )
    elif backend == "ollama":
        return OllamaLLM(
            model_name=llm_settings.model_name,
            host=llm_settings.ollama.host
        )
    elif backend == "openai":
        return OpenAICompatibleLLM(
            model_name=llm_settings.model_name,
            api_base=llm_settings.openai.api_base,
            api_key=llm_settings.openai.api_key
        )
    elif backend == "vllm":
        return VLLMClient(
            model_name=llm_settings.model_name,
            tensor_parallel_size=llm_settings.vllm.tensor_parallel_size
        )
    else:
        raise ValueError(f"Unknown LLM backend: {backend}")
