"""
Embeddings utilities using Ollama and Nomic.
"""
import os
import requests
import numpy as np
from typing import List, Optional, Union
import ollama
from app.config import Config

class EmbeddingService:
    """Service for generating embeddings using Ollama and Nomic."""
    
    @staticmethod
    def get_ollama_embeddings(texts: List[str], model: Optional[str] = None) -> List[List[float]]:
        """
        Generate embeddings using Ollama.
        
        Args:
            texts: List of texts to generate embeddings for
            model: Ollama model to use for embeddings (default from config)
            
        Returns:
            List of embeddings as float arrays
        """
        if model is None:
            # Use model from config
            model = Config.OLLAMA_MODEL
        
        # Set Ollama host from config
        ollama.host = Config.OLLAMA_HOST
        
        embeddings = []
        
        for text in texts:
            try:
                # Call Ollama API for embeddings
                response = ollama.embeddings(model=model, prompt=text)
                embedding = response.get("embedding", [])
                embeddings.append(embedding)
            except Exception as e:
                print(f"Error generating Ollama embedding: {e}")
                # Return a zero embedding as fallback
                embeddings.append([0.0] * 768)  # typical dimension for text embeddings
        
        return embeddings
    
    @staticmethod
    def get_nomic_embeddings(texts: List[str]) -> List[List[float]]:
        """
        Generate embeddings using Nomic.
        
        Args:
            texts: List of texts to generate embeddings for
            
        Returns:
            List of embeddings as float arrays
        """
        try:
            # The new version of Nomic requires a Cohere API key, so we'll fall back to Ollama
            # if we don't have one configured
            cohere_api_key = Config.COHERE_API_KEY
            
            if not cohere_api_key:
                print("No Cohere API key found for Nomic embeddings, falling back to Ollama")
                return EmbeddingService.get_ollama_embeddings(texts)
            
            # Dynamically import nomic embedders to avoid startup errors if not available
            from nomic.embedders import CohereEmbedder
            
            # Create a Nomic embedding model using CohereEmbedder with API key
            embedding_model = CohereEmbedder(cohere_api_key=cohere_api_key)
            
            # Generate embeddings for the texts
            embeddings = []
            for text in texts:
                embedding = embedding_model.embed(text)
                embeddings.append(embedding)
            
            return embeddings
        except Exception as e:
            print(f"Error generating Nomic embeddings: {e}")
            # Fall back to Ollama embeddings
            print("Falling back to Ollama embeddings")
            return EmbeddingService.get_ollama_embeddings(texts)
    
    @staticmethod
    def get_embeddings(texts: Union[str, List[str]], use_nomic: Optional[bool] = None) -> List[List[float]]:
        """
        Generate embeddings using either Nomic or Ollama.
        
        Args:
            texts: Text or list of texts to generate embeddings for
            use_nomic: Whether to use Nomic (True) or Ollama (False), defaults to config setting
            
        Returns:
            List of embeddings as float arrays
        """
        # Convert single text to list
        if isinstance(texts, str):
            texts = [texts]
        
        # If use_nomic is not specified, use the config setting
        if use_nomic is None:
            use_nomic = Config.USE_NOMIC_EMBEDDINGS
        
        # Generate embeddings using chosen method
        if use_nomic:
            return EmbeddingService.get_nomic_embeddings(texts)
        else:
            return EmbeddingService.get_ollama_embeddings(texts)