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

-- requires: rag-core

BEGIN;

CREATE FUNCTION rag.import_collection_json(
  p_data jsonb,
  p_overwrite boolean DEFAULT false
) RETURNS uuid
LANGUAGE plpgsql
AS $$
DECLARE
  v_collection_id uuid;
  v_collection_name text;
  v_model record;
  v_model_id uuid;
  v_doc record;
  v_document_id uuid;
  v_chunk record;
  v_chunk_id uuid;
  v_embedding_model_name text;
  v_embedding_vector text;
BEGIN
  -- Validate version
  IF (p_data->>'version') IS NULL OR (p_data->>'version') != '1.0' THEN
    RAISE EXCEPTION 'Unsupported or missing version. Expected 1.0, got: %', p_data->>'version';
  END IF;

  v_collection_name := p_data->'collection'->>'name';
  
  IF v_collection_name IS NULL THEN
    RAISE EXCEPTION 'Collection name is required';
  END IF;

  -- Check if collection exists
  SELECT id INTO v_collection_id
  FROM rag.collection
  WHERE name = v_collection_name;

  IF FOUND AND NOT p_overwrite THEN
    RAISE EXCEPTION 'Collection already exists: %. Use p_overwrite=true to replace.', v_collection_name;
  END IF;

  -- Delete existing collection if overwriting
  IF FOUND AND p_overwrite THEN
    DELETE FROM rag.collection WHERE id = v_collection_id;
  END IF;

  -- Create collection
  INSERT INTO rag.collection (name, description, chunk_config, metadata)
  VALUES (
    v_collection_name,
    p_data->'collection'->>'description',
    COALESCE(p_data->'collection'->'chunk_config', '{"strategy": "semantic", "max_tokens": 512, "overlap_tokens": 50}'::jsonb),
    COALESCE(p_data->'collection'->'metadata', '{}'::jsonb)
  )
  RETURNING id INTO v_collection_id;

  -- Import models
  FOR v_model IN SELECT * FROM jsonb_array_elements(COALESCE(p_data->'models', '[]'::jsonb))
  LOOP
    -- Insert or get existing model
    INSERT INTO rag.embedding_model (name, provider, dimensions, metadata)
    VALUES (
      v_model.value->>'name',
      v_model.value->>'provider',
      (v_model.value->>'dimensions')::int,
      COALESCE(v_model.value->'metadata', '{}'::jsonb)
    )
    ON CONFLICT (name) DO UPDATE SET
      provider = EXCLUDED.provider,
      dimensions = EXCLUDED.dimensions,
      metadata = EXCLUDED.metadata,
      updated_at = now()
    RETURNING id INTO v_model_id;

    -- Link model to collection
    INSERT INTO rag.collection_model (collection_id, model_id)
    VALUES (v_collection_id, v_model_id)
    ON CONFLICT (collection_id, model_id) DO NOTHING;
  END LOOP;

  -- Import documents
  FOR v_doc IN SELECT * FROM jsonb_array_elements(COALESCE(p_data->'documents', '[]'::jsonb))
  LOOP
    INSERT INTO rag.document (collection_id, source_uri, source_type, content, content_hash, metadata, status)
    VALUES (
      v_collection_id,
      v_doc.value->>'source_uri',
      v_doc.value->>'source_type',
      v_doc.value->>'content',
      v_doc.value->>'content_hash',
      COALESCE(v_doc.value->'metadata', '{}'::jsonb),
      'embedded'
    )
    RETURNING id INTO v_document_id;

    -- Import chunks
    FOR v_chunk IN SELECT * FROM jsonb_array_elements(COALESCE(v_doc.value->'chunks', '[]'::jsonb))
    LOOP
      INSERT INTO rag.chunk (document_id, content, chunk_index, start_offset, end_offset, token_count, metadata)
      VALUES (
        v_document_id,
        v_chunk.value->>'content',
        (v_chunk.value->>'index')::int,
        (v_chunk.value->>'start_offset')::int,
        (v_chunk.value->>'end_offset')::int,
        (v_chunk.value->>'token_count')::int,
        COALESCE(v_chunk.value->'metadata', '{}'::jsonb)
      )
      RETURNING id INTO v_chunk_id;

      -- Import embeddings
      IF v_chunk.value->'embeddings' IS NOT NULL THEN
        FOR v_embedding_model_name, v_embedding_vector IN 
          SELECT key, value FROM jsonb_each_text(v_chunk.value->'embeddings')
        LOOP
          SELECT id INTO v_model_id
          FROM rag.embedding_model
          WHERE name = v_embedding_model_name;

          IF FOUND THEN
            INSERT INTO rag.embedding (chunk_id, model_id, vector)
            VALUES (v_chunk_id, v_model_id, v_embedding_vector::vector);
          END IF;
        END LOOP;
      END IF;
    END LOOP;
  END LOOP;

  RETURN v_collection_id;
END;
$$;

COMMENT ON FUNCTION rag.import_collection_json(jsonb, boolean) IS 'Import a collection from JSON export format';

COMMIT;
