"""
VSS LLM Client
Client for text-based LLM inference (routing and SQL generation).
Supports both Ollama and vLLM backends.
"""

import logging
import httpx
from typing import Optional, List
from abc import ABC, abstractmethod

from config.settings import LLMConfig, get_settings

logger = logging.getLogger(__name__)


class BaseLLMClient(ABC):
    """Abstract base class for LLM clients."""
    
    @abstractmethod
    async def generate(
        self,
        prompt: str,
        max_tokens: int = 512,
        temperature: float = 0.1,
        stop: Optional[List[str]] = None,
    ) -> str:
        """Generate text completion."""
        pass
    
    @abstractmethod
    async def health_check(self) -> bool:
        """Check if LLM service is available."""
        pass


class OllamaClient(BaseLLMClient):
    """Ollama-based LLM client."""
    
    def __init__(self, config: Optional[LLMConfig] = None):
        self.config = config or get_settings().llm
        self.base_url = self.config.endpoint
        self.model = self.config.model_name
        self._client = httpx.AsyncClient(timeout=60.0)
    
    async def generate(
        self,
        prompt: str,
        max_tokens: int = 512,
        temperature: float = 0.1,
        stop: Optional[List[str]] = None,
    ) -> str:
        """Generate text using Ollama API."""
        url = f"{self.base_url}/api/generate"
        
        payload = {
            "model": self.model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "num_predict": max_tokens,
                "temperature": temperature,
            }
        }
        
        if stop:
            payload["options"]["stop"] = stop
        
        logger.debug(f"Ollama request to {url}, prompt length: {len(prompt)}")
        
        try:
            response = await self._client.post(url, json=payload)
            response.raise_for_status()
            result = response.json()
            generated_text = result.get("response", "")
            logger.debug(f"Ollama response length: {len(generated_text)}")
            return generated_text.strip()
        except httpx.HTTPError as e:
            logger.error(f"Ollama HTTP error: {e}")
            raise
        except Exception as e:
            logger.error(f"Ollama generation error: {e}")
            raise
    
    async def health_check(self) -> bool:
        """Check Ollama service health."""
        try:
            response = await self._client.get(f"{self.base_url}/api/tags")
            return response.status_code == 200
        except Exception as e:
            logger.warning(f"Ollama health check failed: {e}")
            return False
    
    async def close(self):
        """Close HTTP client."""
        await self._client.aclose()


class VLLMTextClient(BaseLLMClient):
    """vLLM-based LLM client using OpenAI-compatible API."""
    
    def __init__(self, config: Optional[LLMConfig] = None):
        self.config = config or get_settings().llm
        self.base_url = self.config.vllm_endpoint
        self.model = self.config.model_name
        self._client = httpx.AsyncClient(timeout=60.0)
    
    async def generate(
        self,
        prompt: str,
        max_tokens: int = 512,
        temperature: float = 0.1,
        stop: Optional[List[str]] = None,
    ) -> str:
        """Generate text using vLLM OpenAI-compatible API."""
        url = f"{self.base_url}/chat/completions"
        
        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens,
            "temperature": temperature,
        }
        
        if stop:
            payload["stop"] = stop
        
        try:
            response = await self._client.post(url, json=payload)
            response.raise_for_status()
            result = response.json()
            return result["choices"][0]["message"]["content"].strip()
        except Exception as e:
            logger.error(f"vLLM generation error: {e}")
            raise
    
    async def health_check(self) -> bool:
        """Check vLLM service health."""
        try:
            response = await self._client.get(f"{self.base_url}/models")
            return response.status_code == 200
        except Exception as e:
            logger.warning(f"vLLM health check failed: {e}")
            return False
    
    async def close(self):
        await self._client.aclose()


# === Factory Function ===

def get_llm_client(config: Optional[LLMConfig] = None) -> BaseLLMClient:
    """Get appropriate LLM client based on configuration."""
    config = config or get_settings().llm
    if config.use_vllm:
        logger.info("Using vLLM backend for LLM")
        return VLLMTextClient(config)
    else:
        logger.info("Using Ollama backend for LLM")
        return OllamaClient(config)


# === Singleton Instance ===
# IMPORTANT: This is a SYNC function that returns the client directly
# The client's methods (generate, health_check) are async

_llm_client: Optional[BaseLLMClient] = None


def get_llm() -> BaseLLMClient:
    """
    Get or create LLM client singleton.
    
    NOTE: This is a SYNC function. Do NOT await it.
    The returned client has async methods (generate, health_check).
    
    Usage:
        llm = get_llm()  # No await
        result = await llm.generate(...)  # Await the method
        healthy = await llm.health_check()  # Await the method
    """
    global _llm_client
    if _llm_client is None:
        _llm_client = get_llm_client()
    return _llm_client