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

-- requires: rag-core

BEGIN;

CREATE FUNCTION rag.export_collection_json(
  p_collection_id uuid
) RETURNS jsonb
LANGUAGE plpgsql
STABLE
AS $$
DECLARE
  v_result jsonb;
  v_collection record;
  v_models jsonb;
  v_documents jsonb;
BEGIN
  -- Get collection
  SELECT * INTO v_collection
  FROM rag.collection
  WHERE id = p_collection_id;

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

  -- Get models used by this collection
  SELECT jsonb_agg(
    jsonb_build_object(
      'name', em.name,
      'provider', em.provider,
      'dimensions', em.dimensions,
      'metadata', em.metadata
    )
  ) INTO v_models
  FROM rag.collection_model cm
  JOIN rag.embedding_model em ON em.id = cm.model_id
  WHERE cm.collection_id = p_collection_id;

  -- Get documents with chunks and embeddings
  SELECT jsonb_agg(
    jsonb_build_object(
      'id', d.id,
      'source_uri', d.source_uri,
      'source_type', d.source_type,
      'content', d.content,
      'content_hash', d.content_hash,
      'metadata', d.metadata,
      'chunks', (
        SELECT jsonb_agg(
          jsonb_build_object(
            'content', c.content,
            'index', c.chunk_index,
            'start_offset', c.start_offset,
            'end_offset', c.end_offset,
            'token_count', c.token_count,
            'metadata', c.metadata,
            'embeddings', (
              SELECT jsonb_object_agg(
                em.name,
                e.vector::text
              )
              FROM rag.embedding e
              JOIN rag.embedding_model em ON em.id = e.model_id
              WHERE e.chunk_id = c.id
            )
          )
          ORDER BY c.chunk_index
        )
        FROM rag.chunk c
        WHERE c.document_id = d.id
      )
    )
  ) INTO v_documents
  FROM rag.document d
  WHERE d.collection_id = p_collection_id;

  -- Build final result
  v_result := jsonb_build_object(
    'version', '1.0',
    'exported_at', now(),
    'collection', jsonb_build_object(
      'id', v_collection.id,
      'name', v_collection.name,
      'description', v_collection.description,
      'chunk_config', v_collection.chunk_config,
      'metadata', v_collection.metadata
    ),
    'models', COALESCE(v_models, '[]'::jsonb),
    'documents', COALESCE(v_documents, '[]'::jsonb)
  );

  RETURN v_result;
END;
$$;

COMMENT ON FUNCTION rag.export_collection_json(uuid) IS 'Export a collection with all documents, chunks, and embeddings as JSON';

COMMIT;
