#!/usr/bin/env python3
"""
demo-sqlite-vec.py — Demo de sqlite-vec + Ollama embeddings

Uso:
    pip install sqlite-vec ollama
    python demo-sqlite-vec.py

Descripción:
    Demuestra cómo usar sqlite-vec para almacenar y buscar vectores de
    embeddings generados por Ollama. Este es el fundamento del RAG local
    que exploraremos en episodios posteriores.

    Flujo:
    1. Carga la extensión sqlite-vec
    2. Crea una tabla virtual para vectores de 768 dimensiones
    3. Genera embeddings con Ollama (nomic-embed-text)
    4. Inserta documentos con sus vectores
    5. Busca por similitud semántica

Requiere:
    - Ollama instalado y corriendo
    - Modelo nomic-embed-text descargado (ollama pull nomic-embed-text)
    - Python 3.10+
    - pip install sqlite-vec ollama

Referencia:
    https://github.com/asg017/sqlite-vec
"""

import sqlite3
import sys


def main():
    # --- Comprobar que tenemos los modelos necesarios --------------------
    print("=" * 50)
    print("Demo: sqlite-vec + Ollama embeddings")
    print("=" * 50)
    print()

    # --- 1. Conectar a SQLite en memoria --------------------------------
    print("[1/5] Conectando a SQLite en memoria...")
    conn = sqlite3.connect(":memory:")
    conn.enable_load_extension(True)

    try:
        conn.load_extension("vec0")
        print("  ✓ Extensión sqlite-vec cargada")
    except sqlite3.OperationalError as e:
        print(f"  ✗ No se pudo cargar vec0: {e}")
        print("  Asegúrate de tener sqlite-vec instalado:")
        print("    pip install sqlite-vec")
        sys.exit(1)

    # --- 2. Crear tabla vectorial ---------------------------------------
    print()
    print("[2/5] Creando tabla vectorial vec_docs (768 dimensiones)...")
    conn.execute("""
        CREATE VIRTUAL TABLE vec_docs USING vec0(
            embedding float[768]
        )
    """)
    print("  ✓ Tabla vec_docs creada")

    # --- 3. Preparar documentos -----------------------------------------
    print()
    print("[3/5] Preparando documentos de ejemplo...")

    documentos = [
        {
            "id": 1,
            "titulo": "Linux",
            "texto": "Linux es un sistema operativo libre y de código abierto "
            "basado en el kernel de Linux. Fue creado por Linus Torvalds "
            "en 1991 y se distribuye bajo la GPL.",
        },
        {
            "id": 2,
            "titulo": "Python",
            "texto": "Python es un lenguaje de programación interpretado, "
            "multiparadigma y de alto nivel. Creado por Guido van Rossum "
            "en 1991, es conocido por su sintaxis clara y legible.",
        },
        {
            "id": 3,
            "titulo": "IA Local",
            "texto": "La inteligencia artificial local permite ejecutar modelos "
            "de lenguaje en tu propio hardware sin conexión a internet. "
            "Herramientas como Ollama facilitan este proceso.",
        },
        {
            "id": 4,
            "titulo": "Docker",
            "texto": "Docker es una plataforma de contenedores que permite "
            "empaquetar aplicaciones y sus dependencias en unidades "
            "portables llamadas contenedores.",
        },
    ]

    for doc in documentos:
        print(f"  - [{doc['id']}] {doc['titulo']}")

    # --- 4. Generar embeddings e insertar --------------------------------
    print()
    print("[4/5] Generando embeddings con Ollama (nomic-embed-text)...")
    print("  (Esto puede tardar unos segundos por documento)")
    print()

    try:
        import ollama
    except ImportError:
        print("  ✗ No se pudo importar ollama. Instálalo con:")
        print("    pip install ollama")
        sys.exit(1)

    for doc in documentos:
        # Generar embedding
        response = ollama.embeddings(model="nomic-embed-text", prompt=doc["texto"])
        vector = response["embedding"]

        # Insertar en sqlite-vec
        conn.execute(
            "INSERT INTO vec_docs(rowid, embedding) VALUES (?, ?)",
            [doc["id"], str(vector)],
        )
        print(f"  ✓ Embedding generado para: {doc['titulo']}")

    # --- 5. Búsqueda semántica -------------------------------------------
    print()
    print("[5/5] Búsqueda semántica...")
    print()

    # Consulta de ejemplo
    consulta = "¿Qué sistema operativo libre puedo usar?"
    print(f'  Consulta: "{consulta}"')
    print()

    # Generar embedding de la consulta
    response = ollama.embeddings(model="nomic-embed-text", prompt=consulta)
    query_vector = response["embedding"]

    # Buscar en sqlite-vec
    cursor = conn.execute(
        """
        SELECT rowid, distance
        FROM vec_docs
        WHERE embedding match ?
        ORDER BY distance
        LIMIT 3
    """,
        [str(query_vector)],
    )

    resultados = cursor.fetchall()

    print("  Resultados (más cercano primero):")
    print()
    for rowid, distancia in resultados:
        doc = documentos[rowid - 1]  # rowid empieza en 1
        print(f"  [{rowid}] {doc['titulo']} (distancia: {distancia:.4f})")
        print(f"      {doc['texto'][:80]}...")
        print()

    # --- Cerrar conexión ------------------------------------------------
    conn.close()
    print("=" * 50)
    print("Demo completada")
    print()
    print("  Este es el fundamento del RAG local.")
    print("  En episodios posteriores veremos cómo:")
    print("    - Indexar documentos completos")
    print("    - Combinar búsqueda semántica + keywords")
    print("    - Integrar con Ollama para respuestas contextuales")
    print("=" * 50)


if __name__ == "__main__":
    main()