"""
Ollama Local AI Service

Provides local AI capabilities for lightweight tasks like:
- Quick sentiment analysis
- Simple text summarization
- Fast pattern classification
- Embeddings generation

Falls back to OpenRouter for complex tasks.
"""

import httpx
import logging
from typing import Optional, List, Dict, Any
from app.config import settings

logger = logging.getLogger(__name__)


class OllamaService:
    """Service for local AI using Ollama."""
    
    def __init__(self):
        self.base_url = settings.OLLAMA_BASE_URL
        self.model = settings.OLLAMA_MODEL
        self.embed_model = settings.OLLAMA_MODEL_EMBED
        self.timeout = settings.OLLAMA_TIMEOUT
        self._available = None  # Cached availability status
    
    async def is_available(self) -> bool:
        """Check if Ollama is running and has the required model."""
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                response = await client.get(f"{self.base_url}/api/tags")
                if response.status_code == 200:
                    data = response.json()
                    models = [m["name"] for m in data.get("models", [])]
                    self._available = self.model in models or any(self.model.split(":")[0] in m for m in models)
                    return self._available
        except Exception as e:
            logger.debug(f"Ollama not available: {e}")
            self._available = False
        return False
    
    async def generate(
        self,
        prompt: str,
        system: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: int = 500,
        model: Optional[str] = None
    ) -> Optional[str]:
        """
        Generate text using local Ollama model.
        
        Args:
            prompt: The user prompt
            system: Optional system prompt
            temperature: Sampling temperature (0-1)
            max_tokens: Maximum tokens to generate
            model: Override default model
            
        Returns:
            Generated text or None if failed
        """
        if not await self.is_available():
            logger.warning("Ollama not available, skipping local generation")
            return None
        
        use_model = model or self.model
        
        payload = {
            "model": use_model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens,
            }
        }
        
        if system:
            payload["system"] = system
        
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    f"{self.base_url}/api/generate",
                    json=payload
                )
                
                if response.status_code == 200:
                    data = response.json()
                    return data.get("response", "").strip()
                else:
                    logger.error(f"Ollama generate failed: {response.status_code}")
                    return None
                    
        except Exception as e:
            logger.error(f"Ollama generate error: {e}")
            return None
    
    async def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 500,
        model: Optional[str] = None
    ) -> Optional[str]:
        """
        Chat completion using local Ollama model.
        
        Args:
            messages: List of {"role": "user/assistant/system", "content": "..."}
            temperature: Sampling temperature
            max_tokens: Maximum tokens to generate
            model: Override default model
            
        Returns:
            Assistant response or None if failed
        """
        if not await self.is_available():
            return None
        
        use_model = model or self.model
        
        payload = {
            "model": use_model,
            "messages": messages,
            "stream": False,
            "options": {
                "temperature": temperature,
                "num_predict": max_tokens,
            }
        }
        
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    f"{self.base_url}/api/chat",
                    json=payload
                )
                
                if response.status_code == 200:
                    data = response.json()
                    return data.get("message", {}).get("content", "").strip()
                else:
                    logger.error(f"Ollama chat failed: {response.status_code}")
                    return None
                    
        except Exception as e:
            logger.error(f"Ollama chat error: {e}")
            return None
    
    async def embed(
        self,
        text: str,
        model: Optional[str] = None
    ) -> Optional[List[float]]:
        """
        Generate embeddings using local model.
        
        Args:
            text: Text to embed
            model: Override default embedding model
            
        Returns:
            Embedding vector or None if failed
        """
        if not await self.is_available():
            return None
        
        use_model = model or self.embed_model
        
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    f"{self.base_url}/api/embeddings",
                    json={"model": use_model, "prompt": text}
                )
                
                if response.status_code == 200:
                    data = response.json()
                    return data.get("embedding")
                else:
                    logger.error(f"Ollama embed failed: {response.status_code}")
                    return None
                    
        except Exception as e:
            logger.error(f"Ollama embed error: {e}")
            return None
    
    async def quick_sentiment(self, text: str) -> Optional[Dict[str, Any]]:
        """
        Quick sentiment analysis using local model.
        Optimized for speed over accuracy.
        
        Args:
            text: Text to analyze
            
        Returns:
            {"sentiment": "positive/negative/neutral", "confidence": 0.0-1.0}
        """
        system = """You are a sentiment analyzer. Respond ONLY with JSON in this exact format:
{"sentiment": "positive" or "negative" or "neutral", "confidence": 0.0 to 1.0}
No other text."""
        
        prompt = f"Analyze the sentiment of this text:\n\n{text[:500]}"  # Limit input
        
        result = await self.generate(
            prompt=prompt,
            system=system,
            temperature=0.1,
            max_tokens=50
        )
        
        if result:
            try:
                import json
                # Try to extract JSON from response
                if "{" in result:
                    json_str = result[result.find("{"):result.rfind("}")+1]
                    return json.loads(json_str)
            except:
                pass
        
        return None
    
    async def quick_classify(
        self,
        text: str,
        categories: List[str]
    ) -> Optional[str]:
        """
        Quick text classification into predefined categories.
        
        Args:
            text: Text to classify
            categories: List of possible categories
            
        Returns:
            Selected category or None
        """
        categories_str = ", ".join(categories)
        system = f"You are a classifier. Respond with ONLY one of these categories: {categories_str}. No other text."
        
        prompt = f"Classify this text into one category:\n\n{text[:500]}"
        
        result = await self.generate(
            prompt=prompt,
            system=system,
            temperature=0.1,
            max_tokens=20
        )
        
        if result:
            # Find matching category
            result_lower = result.lower().strip()
            for cat in categories:
                if cat.lower() in result_lower:
                    return cat
        
        return None
    
    async def quick_summarize(self, text: str, max_sentences: int = 2) -> Optional[str]:
        """
        Quick text summarization.
        
        Args:
            text: Text to summarize
            max_sentences: Maximum sentences in summary
            
        Returns:
            Summary or None
        """
        system = f"Summarize in {max_sentences} sentence(s) or less. Be concise and direct."
        
        result = await self.generate(
            prompt=text[:2000],  # Limit input
            system=system,
            temperature=0.3,
            max_tokens=150
        )
        
        return result


# Global instance
ollama_service = OllamaService()


async def get_ai_response(
    prompt: str,
    system: Optional[str] = None,
    use_local: bool = True,
    fallback_to_cloud: bool = True
) -> Optional[str]:
    """
    Unified AI response function that tries local first, then cloud.
    
    Args:
        prompt: User prompt
        system: System prompt
        use_local: Whether to try Ollama first
        fallback_to_cloud: Whether to fallback to OpenRouter if local fails
        
    Returns:
        AI response or None
    """
    # Try local first if enabled
    if use_local and settings.USE_LOCAL_AI:
        result = await ollama_service.generate(prompt, system)
        if result:
            logger.info("Used local Ollama for AI response")
            return result
    
    # Fallback to cloud
    if fallback_to_cloud and settings.OPENROUTER_API_KEY:
        from app.services.openrouter import openrouter_service
        # This would need a simple generate method in openrouter
        logger.info("Falling back to OpenRouter for AI response")
        # For now, return None - full integration would go here
        pass
    
    return None
