"""
VSS VLM Client
Client for Cosmos-Reason2-8B via vLLM OpenAI-compatible API.
"""

import logging
import base64
from typing import Optional, List, Dict, Any, Union

import httpx

from config.settings import VLMConfig, get_settings

logger = logging.getLogger(__name__)


class VLMClient:
    """
    Vision-Language Model client for Cosmos-Reason2-8B.
    Uses vLLM's OpenAI-compatible API for inference.
    """
    
    def __init__(self, config: Optional[VLMConfig] = None):
        self.config = config or get_settings().vlm
        self.base_url = self.config.endpoint
        self.model = self.config.model_name
        self._client = httpx.AsyncClient(timeout=120.0)
    
    async def generate(
        self,
        prompt: str,
        images: Optional[List[Union[str, bytes]]] = None,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
    ) -> str:
        """
        Generate response from VLM.
        
        Args:
            prompt: Text prompt/question
            images: List of images (base64 strings or raw bytes)
            max_tokens: Max response tokens
            temperature: Sampling temperature
        
        Returns:
            Generated text response
        """
        url = f"{self.base_url}/chat/completions"
        max_tokens = max_tokens or self.config.max_tokens
        temperature = temperature or self.config.temperature
        
        # Build message content
        content = []
        
        # Add images first
        if images:
            for img in images:
                img_b64 = self._to_base64(img)
                content.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}
                })
        
        # Add text prompt
        content.append({"type": "text", "text": prompt})
        
        payload = {
            "model": self.model,
            "messages": [{"role": "user", "content": content}],
            "max_tokens": max_tokens,
            "temperature": temperature,
        }
        
        logger.debug(f"VLM request: {len(images or [])} images, prompt: {len(prompt)} chars")
        
        try:
            response = await self._client.post(url, json=payload)
            response.raise_for_status()
            result = response.json()
            generated = result["choices"][0]["message"]["content"]
            logger.debug(f"VLM response: {len(generated)} chars")
            return generated.strip()
        except httpx.HTTPStatusError as e:
            logger.error(f"VLM HTTP error: {e.response.status_code}")
            raise
        except Exception as e:
            logger.error(f"VLM generation error: {e}")
            raise
    
    async def describe_scene(
        self,
        images: List[Union[str, bytes]],
        context: str = ""
    ) -> str:
        """Generate scene description for surveillance frame."""
        prompt = f"""Analyze this surveillance frame.
{f'Context: {context}' if context else ''}

Describe:
- People visible and their activities
- Vehicles (types, colors, positions)
- Notable objects or events
- Overall scene assessment

Be specific and factual."""
        return await self.generate(prompt=prompt, images=images)
    
    async def health_check(self) -> bool:
        """Check if VLM server is available."""
        try:
            # vLLM models endpoint
            response = await self._client.get(f"{self.base_url}/models")
            return response.status_code == 200
        except Exception as e:
            logger.warning(f"VLM health check failed: {e}")
            return False
    
    def _to_base64(self, data: Union[str, bytes]) -> str:
        """Convert data to base64 string."""
        if isinstance(data, bytes):
            return base64.b64encode(data).decode("utf-8")
        elif isinstance(data, str):
            # Check if already base64
            try:
                base64.b64decode(data)
                return data
            except Exception:
                # Assume file path
                with open(data, "rb") as f:
                    return base64.b64encode(f.read()).decode("utf-8")
        else:
            raise ValueError(f"Unsupported data type: {type(data)}")
    
    async def close(self):
        """Close HTTP client."""
        await self._client.aclose()


# === Singleton Instance ===
# IMPORTANT: This is a SYNC function that returns the client directly

_vlm_client: Optional[VLMClient] = None


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