-- Deploy schemas/rag/procedures/create_hnsw_index to pg
-- made with <3 @ constructive.io

-- requires: rag-core

BEGIN;

CREATE FUNCTION rag.create_hnsw_index(
  p_collection_id uuid,
  p_model_name text DEFAULT NULL,
  p_m int DEFAULT 16,
  p_ef_construction int DEFAULT 64,
  p_distance_op text DEFAULT 'cosine'
) RETURNS text
LANGUAGE plpgsql
AS $$
DECLARE
  v_model_id uuid;
  v_model_name text;
  v_model_dimensions int;
  v_index_name text;
  v_op_class text;
  v_sql text;
BEGIN
  -- Get model
  IF p_model_name IS NOT NULL THEN
    SELECT id, name, dimensions INTO v_model_id, v_model_name, v_model_dimensions
    FROM rag.embedding_model
    WHERE name = p_model_name;

    IF NOT FOUND THEN
      RAISE EXCEPTION 'Model not found: %', p_model_name;
    END IF;
  ELSE
    -- Use default model for collection
    SELECT em.id, em.name, em.dimensions INTO v_model_id, v_model_name, v_model_dimensions
    FROM rag.collection_model cm
    JOIN rag.embedding_model em ON em.id = cm.model_id
    WHERE cm.collection_id = p_collection_id
      AND cm.is_default = true
    LIMIT 1;

    IF NOT FOUND THEN
      SELECT em.id, em.name, em.dimensions INTO v_model_id, v_model_name, v_model_dimensions
      FROM rag.collection_model cm
      JOIN rag.embedding_model em ON em.id = cm.model_id
      WHERE cm.collection_id = p_collection_id
      LIMIT 1;
    END IF;

    IF NOT FOUND THEN
      RAISE EXCEPTION 'No embedding model found for collection: %', p_collection_id;
    END IF;
  END IF;

  -- Determine operator class based on distance type
  CASE p_distance_op
    WHEN 'cosine' THEN v_op_class := 'vector_cosine_ops';
    WHEN 'l2', 'euclidean' THEN v_op_class := 'vector_l2_ops';
    WHEN 'ip', 'inner_product' THEN v_op_class := 'vector_ip_ops';
    ELSE RAISE EXCEPTION 'Unsupported distance operator: %. Use cosine, l2, or ip', p_distance_op;
  END CASE;

  -- Generate index name
  v_index_name := 'idx_rag_embedding_hnsw_' || 
    replace(p_collection_id::text, '-', '_') || '_' ||
    replace(v_model_name, '-', '_');

  -- Build and execute CREATE INDEX
  -- Cast vector to specific dimensions to support dimensionless vector columns
  v_sql := format(
    'CREATE INDEX IF NOT EXISTS %I ON rag.embedding USING hnsw ((vector::vector(%s)) %s) WITH (m = %s, ef_construction = %s) WHERE model_id = %L',
    v_index_name,
    v_model_dimensions,
    v_op_class,
    p_m,
    p_ef_construction,
    v_model_id
  );

  EXECUTE v_sql;

  RETURN v_index_name;
END;
$$;

COMMENT ON FUNCTION rag.create_hnsw_index(uuid, text, int, int, text) IS 'Create an HNSW index for a collection and model. Parameters: m (connections per layer, default 16), ef_construction (index build quality, default 64), distance_op (cosine, l2, ip)';

COMMIT;
