import numpy as np
import time
import random
import logging
from typing import List, Dict, Any

logger = logging.getLogger(__name__)

class SemanticRetriever:
    def __init__(self, ali_client=None, openai_client=None):
        from sentence_transformers import SentenceTransformer
        # Gunakan model multilingual E5 untuk kapabilitas lintas bahasa (termasuk Indonesia)
        self.model = SentenceTransformer("intfloat/multilingual-e5-small", device="cpu")
        self.vector_dim = 384

    def _cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
        """Calculate cosine similarity between two vectors."""
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)
        if norm1 == 0 or norm2 == 0:
            return 0.0
        return float(np.dot(vec1, vec2) / (norm1 * norm2))

    def _embed_batch(self, texts: List[str], is_query: bool = False) -> List[List[float]]:
        """Embed a list of texts using the local SentenceTransformer model with E5 prefix rules."""
        prefix = "query: " if is_query else "passage: "
        
        # E5 models have a max sequence length of 512 tokens. 
        # We truncate by character (approx 1500 chars to be safe) before encoding.
        MAX_CHARS = 1500
        
        prefixed_texts = []
        for t in texts:
            cleaned = (t or "").strip()
            if len(cleaned) > MAX_CHARS:
                logger.warning(f"Text exceeded {MAX_CHARS} chars, truncating before embedding.")
                cleaned = cleaned[:MAX_CHARS]
            prefixed_texts.append(f"{prefix}{cleaned}")
        
        # model.encode returns a numpy array or list of numpy arrays, we convert to list of floats
        embeddings = self.model.encode(prefixed_texts, batch_size=32, show_progress_bar=False)
        return [emb.tolist() for emb in embeddings]

    def find_top_k_parents(self, queries: List[str], child_chunks: List[Dict[str, Any]], top_k: int = 3) -> Dict[str, List[str]]:
        """
        Embed queries and child chunks, then calculate similarities.
        Returns a dictionary mapping each query to a list of top_k parent_ids.
        """
        if not child_chunks or not queries:
            return {q: [] for q in queries}

        logger.info(f"Embedding {len(child_chunks)} child chunks with local model...")
        child_texts = [c["content"] for c in child_chunks]
        child_embeddings = self._embed_batch(child_texts, is_query=False)
        child_vecs = np.array(child_embeddings)

        logger.info(f"Embedding {len(queries)} queries with local model...")
        query_embeddings = self._embed_batch(queries, is_query=True)
        query_vecs = np.array(query_embeddings)

        results = {}
        for q_idx, query in enumerate(queries):
            q_vec = query_vecs[q_idx]
            
            # Skip if query embedding failed (zero vector)
            if np.linalg.norm(q_vec) == 0:
                 results[query] = []
                 continue
                 
            scores = []
            for c_idx, c_vec in enumerate(child_vecs):
                score = self._cosine_similarity(q_vec, c_vec)
                scores.append((score, child_chunks[c_idx]["parent_id"]))
            
            # Sort by score descending
            scores.sort(key=lambda x: x[0], reverse=True)
            
            # Extract unique top parent IDs (a parent might have multiple top scoring children)
            top_parents = []
            seen_parents = set()
            for score, p_id in scores:
                if p_id not in seen_parents:
                    top_parents.append(p_id)
                    seen_parents.add(p_id)
                if len(top_parents) >= top_k:
                    break
                    
            results[query] = top_parents
            logger.info(f"Query: '{query}' -> Top parents: {top_parents}")

        return results

if __name__ == "__main__":
    # Smoke test with actual local model
    retriever = SemanticRetriever()
    dummy_children = [
        {"parent_id": "p1", "content": "Information about machine learning and neural networks."},
        {"parent_id": "p1", "content": "Deep learning models require large datasets."},
        {"parent_id": "p2", "content": "Photosynthesis is the process by which plants make food."},
    ]
    queries = ["What is machine learning?", "How do plants eat?"]
    
    print("Running local retrieval with E5 prefix and truncation guards...")
    res = retriever.find_top_k_parents(queries, dummy_children, top_k=1)
    for q, p_ids in res.items():
        print(f"'{q}': {p_ids}")
