#!/usr/bin/env python3
"""rag_hyde.py — HyDE: Hypothetical Document Embeddings"""
import requests, numpy as np, sqlite3, os, math
from numpy.linalg import norm

DB_PATH = os.path.expanduser("~/.cerebro/rag_conocimiento.db")
OLLAMA_URL = "http://localhost:11434/api"

def generate_hypothetical_document(query: str, model: str ="llama3.2") -> str:
    prompt = f"""Genera un texto breve y factual que responda a la siguiente pregunta.
El texto debe ser informativo, objetivo y escrito en español.

Pregunta: {query}

Texto informativo:"""
    resp = requests.post(f"{OLLAMA_URL}/generate",
        json={"model": model, "prompt": prompt, "stream": False})
    return resp.json()["response"]

def get_embedding(text: str, model: str ="bge-m3") -> np.ndarray:
    resp = requests.post(f"{OLLAMA_URL}/embed",
        json={"model": model, "input": text})
    return np.array(resp.json()["embeddings"][0], dtype=np.float32)

def blob_to_vector(blob: bytes) -> np.ndarray:
    return np.frombuffer(blob, dtype=np.float32).copy()

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    return float(np.dot(a, b))

def hyde_search(query: str, top_k: int = 5):
    hypothetic_doc = generate_hypothetical_document(query)
    print(f"[HyDE] Documento generado: {hypothetic_doc[:100]}...")
    hyde_emb = get_embedding(hypothetic_doc)
    hyde_emb = hyde_emb / (norm(hyde_emb) + 1e-10)
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.execute("SELECT id, contenido, embedding FROM chunks")
    results = []
    for row in cursor:
        contenido, emb_blob = row[1], row[2]
        emb_vec = blob_to_vector(emb_blob)
        emb_vec = emb_vec / (norm(emb_vec) + 1e-10)
        cos_sim = cosine_similarity(hyde_emb, emb_vec)
        results.append((cos_sim, contenido[:200]))
    
    results.sort(key=lambda x: x[0], reverse=True)
    conn.close()
    return results[:top_k]

def hyde_hybrid_search(query: str, alpha: float = 0.3, top_k: int = 5):
    """HyDE + Hybrid search combinado. Alpha bajo porque HyDE ya captura semántica."""
    hyde_doc = generate_hypothetical_document(query)
    hyde_emb = get_embedding(hyde_doc)
    hyde_emb = hyde_emb / (norm(hyde_emb) + 1e-10)
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.execute("""
        SELECT c.contenido, c.embedding,
               bm25(chunks_fts, 0, 0, 5, 5) as fts_score
        FROM chunks_fts JOIN chunks c ON chunks_fts.rowid = c.id
        WHERE chunks_fts MATCH ?
    """, (query,))
    
    results = []
    for row in cursor:
        contenido, emb_blob, fts_score = row
        emb_vec = blob_to_vector(emb_blob)
        emb_vec = emb_vec / (norm(emb_vec) + 1e-10)
        cos_sim = cosine_similarity(hyde_emb, emb_vec)
        fts_norm = 1.0 / (1.0 + math.exp(-fts_score / 10.0))
        score_total = alpha * fts_norm + (1.0 - alpha) * cos_sim
        results.append((score_total, contenido[:200]))
    
    results.sort(reverse=True)
    conn.close()
    return results[:top_k]