#!/usr/bin/env python3
"""
rag_hyde.py - HyDE (Hypothetical Document Embeddings) para RAG

Genera un documento hipotético con Ollama a partir de una consulta,
lo embeddea con bge-m3, y busca por similitud coseno en la BD.

Uso:
  python rag_hyde.py "¿qué es una variable de entorno?"
  python rag_hyde.py "explica la configuración de red" --top 10
"""

import argparse
import sqlite3
import sys
import numpy as np
from pathlib import Path

DB_PATH = Path.home() / ".cerebro" / "rag_conocimiento.db"
OLLAMA_MODEL = "llama3.2"
EMBED_MODEL = "bge-m3"

HYDE_PROMPT = """Eres un asistente experto en Linux y administración de sistemas.
Dada una pregunta, genera un documento hipotético de 3-5 párrafos que RESPONDA
a la pregunta como si fuera un artículo del blog o la documentación oficial.

Sé técnico pero claro. Usa markdown si es necesario.
No digas "basado en" ni "hipotético". Responde DIRECTAMENTE.

Pregunta: {consulta}

Documento:"""


def _ollama():
    try:
        import ollama
    except ImportError:
        print("Necesitas instalar ollama: pip install ollama")
        sys.exit(1)
    return ollama


def generar_doc_hipotetico(consulta: str) -> str:
    ollama = _ollama()
    prompt = HYDE_PROMPT.format(consulta=consulta)
    respuesta = ollama.generate(model=OLLAMA_MODEL, prompt=prompt)
    return respuesta["response"]


def embed_text(texto: str) -> list[float]:
    ollama = _ollama()
    respuesta = ollama.embed(model=EMBED_MODEL, input=texto)
    return respuesta["embeddings"][0]


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


def cargar_embeddings_db() -> list[tuple[int, str, np.ndarray]]:
    if not DB_PATH.exists():
        print(f"BD no encontrada: {DB_PATH}")
        sys.exit(1)
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    cursor.execute("""SELECT c.id, c.content, e.vector
                       FROM chunks c JOIN embeddings e ON e.chunk_id = c.id""")
    filas = cursor.fetchall()
    conn.close()
    if not filas:
        print("No hay chunks con embeddings en la BD.")
        sys.exit(1)
    return [(chunk_id, contenido, np.frombuffer(blob, dtype=np.float32).copy())
            for chunk_id, contenido, blob in filas]


def buscar_semantica(query_emb: np.ndarray, chunks: list, top_k: int = 5) -> list:
    query_norm = query_emb / (np.linalg.norm(query_emb) + 1e-10)
    scores = []
    for chunk_id, contenido, vector in chunks:
        vec_norm = vector / (np.linalg.norm(vector) + 1e-10)
        sim = cosine_sim(query_norm, vec_norm)
        scores.append((chunk_id, contenido, sim))
    scores.sort(key=lambda x: x[2], reverse=True)
    return scores[:top_k]


def main():
    parser = argparse.ArgumentParser(description="HyDE: genera un doc hipotético y busca en la BD")
    parser.add_argument("consulta", help="Consulta en lenguaje natural")
    parser.add_argument("--top", type=int, default=5, help="Resultados a mostrar")
    args = parser.parse_args()

    doc_hipotetico = generar_doc_hipotetico(args.consulta)
    print(f"Documento generado:\n{─ * 50}\n{doc_hipotetico[:500]}...\n{─ * 50}\n")

    query_emb = np.array(embed_text(doc_hipotetico), dtype=np.float32)
    chunks = cargar_embeddings_db()
    resultados = buscar_semantica(query_emb, chunks, args.top)

    for i, (chunk_id, contenido, score) in enumerate(resultados, 1):
        preview = contenido[:150].replace("\n", " ")
        print(f"  {i}. [{score:.4f}] (chunk #{chunk_id})")
        print(f"     {preview}...")


if __name__ == "__main__":
    main()