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

-- requires: rag-core

BEGIN;

CREATE FUNCTION rag.search_similar(
  p_collection_id uuid,
  p_query_vector vector,
  p_model_name text DEFAULT NULL,
  p_limit int DEFAULT 10,
  p_distance_threshold float DEFAULT NULL
) RETURNS TABLE (
  chunk_id uuid,
  document_id uuid,
  content text,
  distance float,
  metadata jsonb
)
LANGUAGE plpgsql
STABLE
AS $$
DECLARE
  v_model_id uuid;
BEGIN
  -- Get model ID if specified
  IF p_model_name IS NOT NULL THEN
    SELECT id INTO v_model_id
    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 INTO v_model_id
    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 no default, use first model
    IF NOT FOUND THEN
      SELECT em.id INTO v_model_id
      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;

  RETURN QUERY
  SELECT 
    e.chunk_id,
    c.document_id,
    c.content,
    (e.vector <=> p_query_vector)::float AS distance,
    c.metadata
  FROM rag.embedding e
  JOIN rag.chunk c ON c.id = e.chunk_id
  JOIN rag.document d ON d.id = c.document_id
  WHERE d.collection_id = p_collection_id
    AND e.model_id = v_model_id
    AND (p_distance_threshold IS NULL OR (e.vector <=> p_query_vector) <= p_distance_threshold)
  ORDER BY e.vector <=> p_query_vector
  LIMIT p_limit;
END;
$$;

COMMENT ON FUNCTION rag.search_similar(uuid, vector, text, int, float) IS 'Search for similar chunks using cosine distance';

-- Overload for searching by collection name
CREATE FUNCTION rag.search_similar(
  p_collection_name text,
  p_query_vector vector,
  p_model_name text DEFAULT NULL,
  p_limit int DEFAULT 10,
  p_distance_threshold float DEFAULT NULL
) RETURNS TABLE (
  chunk_id uuid,
  document_id uuid,
  content text,
  distance float,
  metadata jsonb
)
LANGUAGE plpgsql
STABLE
AS $$
DECLARE
  v_collection_id uuid;
BEGIN
  SELECT id INTO v_collection_id
  FROM rag.collection
  WHERE name = p_collection_name;

  IF NOT FOUND THEN
    RAISE EXCEPTION 'Collection not found: %', p_collection_name;
  END IF;

  RETURN QUERY
  SELECT * FROM rag.search_similar(v_collection_id, p_query_vector, p_model_name, p_limit, p_distance_threshold);
END;
$$;

COMMENT ON FUNCTION rag.search_similar(text, vector, text, int, float) IS 'Search for similar chunks by collection name using cosine distance';

COMMIT;
