import hashlib
import json
import logging

from openai import AsyncOpenAI, OpenAI

from config import OPENAI_API_KEY

logger = logging.getLogger(__name__)

__all__ = [
    "create_embedding",
    "create_embedding_async",
    "create_embeddings_async",
    "get_async_embedding_client",
    "get_embedding_client",
]


class EmbeddingClientManager:
    """
    Singleton Class quản lý OpenAI Embedding Client (Sync & Async).
    """

    def __init__(self):
        self._client: OpenAI | None = None
        self._async_client: AsyncOpenAI | None = None

    def get_client(self) -> OpenAI:
        """Sync Client lazy loading"""
        if self._client is None:
            if not OPENAI_API_KEY:
                raise RuntimeError("CRITICAL: OPENAI_API_KEY chưa được thiết lập")
            self._client = OpenAI(api_key=OPENAI_API_KEY)
        return self._client

    def get_async_client(self) -> AsyncOpenAI:
        """Async Client lazy loading"""
        if self._async_client is None:
            if not OPENAI_API_KEY:
                raise RuntimeError("CRITICAL: OPENAI_API_KEY chưa được thiết lập")
            self._async_client = AsyncOpenAI(api_key=OPENAI_API_KEY)
        return self._async_client


logger = logging.getLogger(__name__)

from common.cache import redis_cache

# --- Singleton ---
_manager = EmbeddingClientManager()
get_embedding_client = _manager.get_client
get_async_embedding_client = _manager.get_async_client


def create_embedding(text: str) -> list[float]:
    """Sync embedding generation with Layer 2 Cache"""
    try:
        # 1. Try Cache (Sync wrapper for get_embedding if needed, but here we just use what we have)
        # Note: common.cache is async, so sync create_embedding will still call OpenAI
        # unless we add a sync cache method. For now, focus on async.
        client = get_embedding_client()
        response = client.embeddings.create(model="text-embedding-3-small", input=text)
        return response.data[0].embedding
    except Exception as e:
        logger.error(f"Error creating embedding (sync): {e}")
        return []


async def create_embedding_async(text: str) -> list[float]:
    """
    Async embedding generation with Layer 2 Cache.
    Saves OpenAI costs by reusing embeddings for identical queries.
    """
    import time
    total_start = time.time()
    
    try:
        # 1. Try Layer 2 Cache
        cache_start = time.time()
        cached = await redis_cache.get_embedding(text)
        cache_time = (time.time() - cache_start) * 1000
        
        # Validate dimension (must be 1536 for text-embedding-3-small)
        if cached and len(cached) == 1536:
            logger.info(f"⚡ [EMBEDDING CACHE HIT] Total: {(time.time() - total_start) * 1000:.2f}ms (cache_check: {cache_time:.2f}ms)")
            return cached
        elif cached:
            logger.warning(f"⚠️ Cached embedding has wrong dimension ({len(cached)}). Regenerating...")

        # 2. Call OpenAI (this is the slow part!)
        openai_start = time.time()
        client = get_async_embedding_client()
        response = await client.embeddings.create(model="text-embedding-3-small", input=text)
        embedding = response.data[0].embedding
        openai_time = (time.time() - openai_start) * 1000

        # 3. Store in Cache
        write_start = time.time()
        if embedding:
            await redis_cache.set_embedding(text, embedding)
        write_time = (time.time() - write_start) * 1000
        
        total_time = (time.time() - total_start) * 1000
        logger.info(f"🔄 [EMBEDDING CACHE MISS] Total: {total_time:.2f}ms (cache_check: {cache_time:.2f}ms, openai: {openai_time:.2f}ms, write: {write_time:.2f}ms)")
            
        return embedding
    except Exception as e:
        logger.error(f"Error creating embedding (async): {e}")
        return []


async def create_embeddings_async(texts: list[str]) -> list[list[float]]:
    """
    Batch async embedding generation with per-item Layer 2 Cache.
    """
    try:
        if not texts:
            return []

        results = [[] for _ in texts]
        missed_indices = []
        missed_texts = []

        client = redis_cache.get_client()
        if client:
            keys = []
            for text in texts:
                text_hash = hashlib.md5(text.strip().lower().encode()).hexdigest()
                keys.append(f"emb_cache:{text_hash}")

            cached_values = await client.mget(keys)
            for i, cached in enumerate(cached_values):
                use_cache = False
                if cached:
                    loaded_emb = json.loads(cached)
                    # Validate dimension
                    if len(loaded_emb) == 1536:
                        results[i] = loaded_emb
                        use_cache = True
                    else:
                        logger.warning(f"⚠️ Cached embedding for item {i} has wrong dimension ({len(loaded_emb)}). Ignoring.")
                
                if not use_cache:
                    missed_indices.append(i)
                    missed_texts.append(texts[i])
        else:
            # Fallback: no redis client, treat all as miss
            missed_indices = list(range(len(texts)))
            missed_texts = texts

        # 2. Call OpenAI for missed texts
        if missed_texts:
            client = get_async_embedding_client()
            response = await client.embeddings.create(model="text-embedding-3-small", input=missed_texts)
            
            # OpenAI returns embeddings in the same order as missed_texts
            for i, data_item in enumerate(response.data):
                idx = missed_indices[i]
                embedding = data_item.embedding
                results[idx] = embedding
                
                # 3. Cache the new embedding
                await redis_cache.set_embedding(missed_texts[i], embedding)

        return results
    except Exception as e:
        logger.error(f"Error creating batch embeddings (async): {e}")
        return [[] for _ in texts]
