import marimo

__generated_with = "0.15.5"
app = marimo.App(width="columns")


@app.cell(column=0)
def _():
    import marimo as mo
    return (mo,)


@app.cell
def _():
    from huggingface_hub import snapshot_download

    snapshot_download(
        repo_id="nomic-ai/nomic-embed-text-v1.5",
        local_dir=".model/nomic-embed-text-v1.5"
    )
    return


@app.cell
def _():
    from transformers import AutoTokenizer, AutoModel
    import torch

    # Ruta al modelo descargado
    model_path = ".model/nomic-embed-text-v1.5"

    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)

    return model, tokenizer, torch


@app.cell
def _(model, tokenizer, torch):
    texts = [
        "La inteligencia artificial está transformando el mundo.",
        "Los embeddings convierten texto en vectores numéricos."
    ]

    # Tokenizar y obtener embeddings
    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=8192)

    with torch.no_grad():
        embeddings = model(**inputs).last_hidden_state.mean(dim=1)

    for text, _vector in zip(texts, embeddings):
        print(f"Texto: {text}\nDimensión: {_vector.shape[0]}\nPrimeros valores: {_vector[:5].tolist()}\n")
    return


@app.cell
def _():
    import os
    import textwrap
    import yaml
    import subprocess
    from pathlib import Path

    # ==== Parámetros ====
    PUERTO_POSTGRES = 55455
    SERVICIO = "postgres_ext"
    PASSWORD = "mipassword"
    USER = "postgres"
    DBNAME = "basededatos"
    return (
        DBNAME,
        PASSWORD,
        PUERTO_POSTGRES,
        Path,
        SERVICIO,
        USER,
        subprocess,
        textwrap,
        yaml,
    )


@app.cell
def _(
    DBNAME,
    HOST,
    PASSWORD,
    PUERTO_POSTGRES,
    USER,
    model,
    psycopg2,
    tokenizer,
    torch,
):

    conn = psycopg2.connect(
        dbname=DBNAME,
        user=USER,
        password=PASSWORD,
        host=HOST,
        port=PUERTO_POSTGRES
    )

    conn.autocommit = True   # 👈 importante
    cur = conn.cursor()


    base_texts = [
        "La inteligencia artificial está transformando el mundo.",
        "Los embeddings convierten texto en vectores numéricos.",
        "PostgreSQL con pgvector permite búsquedas semánticas.",
        "El aprendizaje profundo impulsa avances en visión computacional.",
        "Transformers cambiaron el campo del procesamiento del lenguaje natural.",
        "Los modelos de lenguaje grande permiten nuevas aplicaciones.",
        "La ciencia de datos combina estadística y programación.",
        "El big data requiere arquitecturas distribuidas.",
        "El machine learning mejora con más datos y cómputo.",
        "El deep learning usa redes neuronales profundas.",
    ]
    # Duplicamos con variaciones para llegar a 20
    _texts = base_texts + [t + f" Ejemplo {i}" for i, t in enumerate(base_texts, start=1)]

    # Tokenizar y obtener embeddings
    _inputs = tokenizer(_texts, return_tensors="pt", padding=True, truncation=True, max_length=512)
    with torch.no_grad():
        _embeddings = model(**_inputs).last_hidden_state.mean(dim=1)

    # Insertar en la base de datos
    for _text, vector in zip(_texts, _embeddings):
        vector_list = vector.tolist()
        # Convertimos a formato adecuado para pgvector: [v1, v2, v3, ...]
        vector_str = "[" + ",".join(f"{v:.6f}" for v in vector_list) + "]"
        cur.execute(
            "INSERT INTO nota (titulo, embedding) VALUES (%s, %s::vector)",
            (_text, vector_str)
        )

    print("✅ Se insertaron 20 textos con sus embeddings en la tabla 'nota'.")

    cur.close()
    conn.close()
    return


@app.cell
def _():
    return


@app.cell
def _(Engine, PUERTO_POSTGRES, create_engine, quote_plus):

    def conectar_postgres(
        host: str = "localhost",
        port: int = PUERTO_POSTGRES,
        dbname: str = "basededatos",
        user: str = "postgres",
        password: str = "mipassword"
    ) -> Engine:
        """
        Devuelve un objeto SQLAlchemy Engine conectado a PostgreSQL.
        Compatible con pandas.to_sql y read_sql.
        """
        pwd = quote_plus(password)  # Escapar caracteres especiales del password
        url = f"postgresql+psycopg2://{user}:{pwd}@{host}:{port}/{dbname}"

        engine = create_engine(
            url,
            pool_size=5,
            max_overflow=10,
            future=True
        )
        return engine
    return


@app.cell
def _(Engine, pd):
    def consultar_df(conn: Engine, query: str, params: dict | None = None) -> pd.DataFrame:
        """
        Ejecuta una consulta SQL usando SQLAlchemy y devuelve los resultados como un DataFrame.
        - conn: Engine devuelto por conectar_postgres()
        - query: str con la consulta SQL
        - params: dict opcional con parámetros de la consulta
        """
        return pd.read_sql(query, con=conn, params=params)
    return


@app.cell
def _(DATABASE_URL, create_engine, sessionmaker):
    # Create engine
    engine = create_engine(DATABASE_URL, echo=True)

    # Create session
    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
    session = SessionLocal()

    print("✅ Conexión establecida con la BBDD")


    return


@app.cell
def _(conn2, mo):
    _df = mo.sql(
        f"""
        SELECT * FROM public.nota
        """,
        engine=conn2
    )
    return


@app.cell
def _():
    from sqlalchemy import create_engine
    from sqlalchemy.engine import Engine
    from urllib.parse import quote_plus
    import pandas as pd
    # import psycopg2
    import sqlglot
    # import os
    # import yaml
    # import subprocess
    return Engine, create_engine, pd, quote_plus


@app.cell(disabled=True)
def _(Path):
    # import os
    # import textwrap
    # import yaml
    # import subprocess
    # from pathlib import Path

    # # ==== Parámetros ====
    # PUERTO_POSTGRES = 55455
    # SERVICIO = "postgres_ext"
    # PASSWORD = "mipassword"
    # USER = "postgres"
    # DBNAME = "basededatos"

    # Lista de extensiones que quieres habilitar
    EXTENSIONES = [
        # builtin
        "hstore", "citext", "uuid-ossp", "pg_trgm",
        # requieren paquetes
        "postgis", "pgvector",
        # "timescaledb"  # <- si quieres usar base de timescaledb, activa la bandera abajo
    ]

    # Usa imagen base de timescaledb cuando la extensión 'timescaledb' esté en la lista
    timescaledb_base_image = False  # pon True si quieres usar la imagen base de TimescaleDB

    RUTA_PROYECTO = Path(".").resolve()
    return EXTENSIONES, RUTA_PROYECTO, timescaledb_base_image


@app.function(disabled=True)
def pkgs_para_extensiones(exts, pg_major=15, use_timescale_base=False):
    """
    Devuelve (pkgs_apt, builtins) para las extensiones solicitadas.
    builtins = extensiones que no requieren apt
    pkgs_apt = paquetes apt necesarios para otras extensiones
    """
    builtins = []
    pkgs_apt = []

    for e in exts:
        e_low = e.lower()
        if e_low in {"hstore", "citext", "uuid-ossp", "pg_trgm"}:
            builtins.append(e_low)
        elif e_low == "postgis":
            pkgs_apt += [
                f"postgresql-{pg_major}-postgis-3",
                f"postgresql-{pg_major}-postgis-3-scripts",
                "gdal-bin",
                "proj-bin",
            ]
        elif e_low == "pgvector":
            pkgs_apt += [f"postgresql-{pg_major}-pgvector"]
        elif e_low == "timescaledb":
            if not use_timescale_base:
                # Para instalar timescaledb necesitarías repos adicionales
                pass
        else:
            raise ValueError(f"Extensión no soportada en este helper: {e}")

    pkgs_apt = sorted(set(pkgs_apt))
    return pkgs_apt, builtins


@app.cell(disabled=True)
def _(DBNAME, Path, USER, textwrap):
    def generar_dockerfile(ruta: Path, exts, use_timescale_base=False, pg_major=15):
        ruta.mkdir(parents=True, exist_ok=True)
        pkgs_apt, builtins = pkgs_para_extensiones(exts, pg_major, use_timescale_base)

        base_image = (
            f"timescale/timescaledb:2.16-pg{pg_major}" if use_timescale_base else f"postgres:{pg_major}"
        )

        apt_block = ""
        if pkgs_apt:
            apt_lines = [
                "RUN apt-get update && \\",
                "    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \\",
            ]
            # 👇 Backslash en **todas** las líneas de paquetes
            for pkg in pkgs_apt:
                apt_lines.append(f"        {pkg} \\")
            # y ya colgamos el rm del último paquete
            apt_lines.append("    && rm -rf /var/lib/apt/lists/*")
            apt_block = "\n".join(apt_lines)

        dockerfile = f"""
        FROM {base_image}

        # Variables de entorno útiles
        ENV POSTGRES_USER={USER} \\
            POSTGRES_DB={DBNAME}

        {apt_block}

        # Copiamos scripts de inicialización
        COPY docker-entrypoint-initdb.d/ /docker-entrypoint-initdb.d/
        """
        (ruta / "Dockerfile").write_text(textwrap.dedent(dockerfile).strip() + "\n", encoding="utf-8")
    return (generar_dockerfile,)


@app.cell(disabled=True)
def _(Path):
    def generar_init_sql(ruta: Path, exts):
        init_dir = ruta / "docker-entrypoint-initdb.d"
        init_dir.mkdir(parents=True, exist_ok=True)

        lines = ["-- Habilitar extensiones solicitadas"]
        for e in exts:
            e_low = e.lower()
            ext_name = {
                "pgvector": "vector",
                "postgis": "postgis",
                "hstore": "hstore",
                "citext": "citext",
                "uuid-ossp": "\"uuid-ossp\"",
                "pg_trgm": "pg_trgm",
                "timescaledb": "timescaledb",
            }.get(e_low, e_low)
            lines.append(f"CREATE EXTENSION IF NOT EXISTS {ext_name};")

        sql = "\n".join(lines) + "\n"
        (init_dir / "10-extensions.sql").write_text(sql, encoding="utf-8")
    return (generar_init_sql,)


@app.cell(disabled=True)
def _(DBNAME, Path, USER, yaml):
    def crear_docker_compose(ruta: Path, servicio: str, puerto_host: int, password: str):
        compose = {
            "version": "3.8",
            "services": {
                servicio: {
                    "build": {"context": "."},
                    "restart": "always",
                    "ports": [f"{puerto_host}:5432"],
                    "environment": {
                        "POSTGRES_PASSWORD": password,
                        "POSTGRES_USER": USER,
                        "POSTGRES_DB": DBNAME,
                    },
                    "healthcheck": {
                        "test": ["CMD-SHELL", "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}"],
                        "interval": "10s",
                        "timeout": "5s",
                        "retries": 5,
                    },
                    "volumes": [
                        "postgres_data:/var/lib/postgresql/data"
                    ],
                }
            },
            "volumes": {"postgres_data": {}}
        }
        (ruta / "docker-compose.yml").write_text(yaml.dump(compose, sort_keys=False), encoding="utf-8")
    return (crear_docker_compose,)


@app.cell(disabled=True)
def _(Path, subprocess):
    def construir_y_levantar(ruta: Path):
        def _run(cmd):
            subprocess.run(cmd, cwd=ruta, check=True)

        try:
            _run(["docker", "compose", "build"])
            _run(["docker", "compose", "up", "-d"])
        except Exception:
            _run(["docker-compose", "build"])
            _run(["docker-compose", "up", "-d"])
    return (construir_y_levantar,)


@app.cell
def _(
    EXTENSIONES,
    PASSWORD,
    PUERTO_POSTGRES,
    RUTA_PROYECTO,
    SERVICIO,
    construir_y_levantar,
    crear_docker_compose,
    generar_dockerfile,
    generar_init_sql,
    timescaledb_base_image,
):
    if __name__ == "__main__":
        RUTA_PROYECTO.mkdir(parents=True, exist_ok=True)
        generar_dockerfile(RUTA_PROYECTO, EXTENSIONES, timescaledb_base_image)
        generar_init_sql(RUTA_PROYECTO, EXTENSIONES)
        crear_docker_compose(RUTA_PROYECTO, SERVICIO, PUERTO_POSTGRES, PASSWORD)
        construir_y_levantar(RUTA_PROYECTO)
    return


@app.cell(column=1)
def _(DBNAME, PASSWORD, PUERTO_POSTGRES, USER):
    import psycopg2

    # Sentencias SQL
    sql = """
    CREATE EXTENSION IF NOT EXISTS vector;

    CREATE TABLE IF NOT EXISTS nota (
        id SERIAL PRIMARY KEY,
        titulo TEXT NOT NULL,
        embedding VECTOR(768)
    );
    """

    HOST = "127.0.0.1"

    def init_db():
        conn = psycopg2.connect(
            dbname=DBNAME,
            user=USER,
            password=PASSWORD,
            host=HOST,
            port=PUERTO_POSTGRES
        )
        conn.autocommit = True
        cur = conn.cursor()
        cur.execute(sql)
        cur.close()
        conn.close()
        print("✅ Tabla 'nota' creada con pgvector.")

    if __name__ == "__main__":
        init_db()
    return HOST, psycopg2


@app.cell
def _():
    return


@app.cell(column=2)
def _(mo):
    # Cell 2
    # Caja de texto para la consulta
    query_input = mo.ui.text(label="Texto de búsqueda", full_width=True)
    query_input
    return (query_input,)


@app.cell
def _(model, query_input, tokenizer, torch):
    # Cell 3
    # Generar embedding del texto introducido
    if query_input.value.strip():
        _inputs = tokenizer(
            [query_input.value],
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=512
        )
        with torch.no_grad():
            embedding = model(**_inputs).last_hidden_state.mean(dim=1)[0].tolist()
        embedding_str = "[" + ",".join(f"{v:.6f}" for v in embedding) + "]"
    else:
        embedding_str = None

    embedding_str
    return (embedding_str,)


@app.cell
def _(
    DBNAME,
    HOST,
    PASSWORD,
    PUERTO_POSTGRES,
    USER,
    embedding_str,
    pd,
    psycopg2,
):
    if embedding_str:

        conn3 = psycopg2.connect(
            dbname=DBNAME,
            user=USER,
            password=PASSWORD,
            host=HOST,
            port=PUERTO_POSTGRES
        )

        conn3.autocommit = True
        _cur = conn3.cursor()

        _cur.execute(
            """
            SELECT id, titulo, embedding <#> %s::vector AS distancia
            FROM nota
            ORDER BY embedding <#> %s::vector
            LIMIT 3;
            """,
            (embedding_str, embedding_str)
        )
        resultados = _cur.fetchall()
        _cur.close()
        conn3.close()

        df = pd.DataFrame(resultados, columns=["ID", "Título", "Distancia"])
    else:
        df = pd.DataFrame(columns=["ID", "Título", "Distancia"])

    df
    return


@app.cell
def _():
    return


if __name__ == "__main__":
    app.run()