framework/vectorstore/pinecone_test.go

package vectorstore

import (
	"context"
	"testing"
	"time"

	bifrost "github.com/maximhq/bifrost/core"
	"github.com/maximhq/bifrost/core/schemas"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

const (
	PineconeTestTimeout          = 30 * time.Second
	PineconeTestNamespace        = "bifrost-test-namespace"
	PineconeTestDimension        = 1536 // Matches text-embedding-3-small dimension
	PineconeTestDefaultAPIKey    = "pclocal"        // Pinecone Local doesn't validate API keys
	PineconeTestDefaultIndexHost = "localhost:5081" // Pinecone Local default port
)

type PineconeTestSetup struct {
	Store  *PineconeStore
	Logger schemas.Logger
	Config PineconeConfig
	ctx    context.Context
	cancel context.CancelFunc
}

func NewPineconeTestSetup(t *testing.T) *PineconeTestSetup {
	apiKey := schemas.NewEnvVar(getEnvWithDefault("PINECONE_API_KEY", PineconeTestDefaultAPIKey))
	indexHost := schemas.NewEnvVar(getEnvWithDefault("PINECONE_INDEX_HOST", PineconeTestDefaultIndexHost))

	config := PineconeConfig{
		APIKey:    *apiKey,
		IndexHost: *indexHost,
	}

	logger := bifrost.NewDefaultLogger(schemas.LogLevelInfo)
	ctx, cancel := context.WithTimeout(context.Background(), PineconeTestTimeout)

	store, err := newPineconeStore(ctx, &config, logger)
	if err != nil {
		cancel()
		t.Fatalf("Failed to create Pinecone store: %v", err)
	}

	setup := &PineconeTestSetup{
		Store:  store,
		Logger: logger,
		Config: config,
		ctx:    ctx,
		cancel: cancel,
	}

	return setup
}

func (ts *PineconeTestSetup) Cleanup(t *testing.T) {
	defer ts.cancel()

	if !testing.Short() {
		ts.cleanupTestData(t)
	}

	if err := ts.Store.Close(ts.ctx, PineconeTestNamespace); err != nil {
		t.Logf("Warning: Failed to close store: %v", err)
	}
}

func (ts *PineconeTestSetup) cleanupTestData(t *testing.T) {
	// Delete all vectors in the test namespace
	err := ts.Store.DeleteNamespace(ts.ctx, PineconeTestNamespace)
	if err != nil {
		t.Logf("Warning: Failed to cleanup test namespace: %v", err)
	}
	t.Logf("Cleaned up test namespace: %s", PineconeTestNamespace)
}

// ============================================================================
// UNIT TESTS
// ============================================================================

func TestPineconeConfig_Validation(t *testing.T) {
	logger := bifrost.NewDefaultLogger(schemas.LogLevelInfo)
	ctx := context.Background()

	tests := []struct {
		name        string
		config      PineconeConfig
		expectError bool
		errorMsg    string
	}{
		{
			name: "missing api key",
			config: PineconeConfig{
				IndexHost: *schemas.NewEnvVar("https://my-index.svc.environment.pinecone.io"),
			},
			expectError: true,
			errorMsg:    "pinecone api_key is required",
		},
		{
			name: "missing index host",
			config: PineconeConfig{
				APIKey: *schemas.NewEnvVar("test-api-key"),
			},
			expectError: true,
			errorMsg:    "pinecone index_host is required",
		},
		{
			name: "empty api key",
			config: PineconeConfig{
				APIKey:    *schemas.NewEnvVar(""),
				IndexHost: *schemas.NewEnvVar("https://my-index.svc.environment.pinecone.io"),
			},
			expectError: true,
			errorMsg:    "pinecone api_key is required",
		},
		{
			name: "empty index host",
			config: PineconeConfig{
				APIKey:    *schemas.NewEnvVar("test-api-key"),
				IndexHost: *schemas.NewEnvVar(""),
			},
			expectError: true,
			errorMsg:    "pinecone index_host is required",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			store, err := newPineconeStore(ctx, &tt.config, logger)

			if tt.expectError {
				assert.Error(t, err)
				assert.Nil(t, store)
				if tt.errorMsg != "" {
					assert.Contains(t, err.Error(), tt.errorMsg)
				}
			} else {
				// Note: This will fail with connection error in unit tests
				// but should pass config validation
				if err != nil {
					assert.Contains(t, err.Error(), "failed to connect")
				}
			}
		})
	}
}

func TestBuildPineconeFilter(t *testing.T) {
	tests := []struct {
		name     string
		queries  []Query
		expected bool
	}{
		{
			name:     "empty queries",
			queries:  []Query{},
			expected: false,
		},
		{
			name: "single string query",
			queries: []Query{
				{Field: "category", Operator: QueryOperatorEqual, Value: "tech"},
			},
			expected: true,
		},
		{
			name: "single numeric query",
			queries: []Query{
				{Field: "size", Operator: QueryOperatorGreaterThan, Value: 1000},
			},
			expected: true,
		},
		{
			name: "multiple queries",
			queries: []Query{
				{Field: "category", Operator: QueryOperatorEqual, Value: "tech"},
				{Field: "public", Operator: QueryOperatorEqual, Value: true},
			},
			expected: true,
		},
		{
			name: "not equal query",
			queries: []Query{
				{Field: "status", Operator: QueryOperatorNotEqual, Value: "deleted"},
			},
			expected: true,
		},
		{
			name: "range queries",
			queries: []Query{
				{Field: "count", Operator: QueryOperatorGreaterThanOrEqual, Value: 10},
				{Field: "score", Operator: QueryOperatorLessThanOrEqual, Value: 100},
			},
			expected: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result, err := buildPineconeFilter(tt.queries)
			assert.NoError(t, err)

			if tt.expected {
				assert.NotNil(t, result)
			} else {
				assert.Nil(t, result)
			}
		})
	}
}

func TestBuildPineconeCondition(t *testing.T) {
	tests := []struct {
		name     string
		query    Query
		expected map[string]interface{}
	}{
		{
			name:     "equal operator",
			query:    Query{Field: "category", Operator: QueryOperatorEqual, Value: "tech"},
			expected: map[string]interface{}{"$eq": "tech"},
		},
		{
			name:     "not equal operator",
			query:    Query{Field: "status", Operator: QueryOperatorNotEqual, Value: "deleted"},
			expected: map[string]interface{}{"$ne": "deleted"},
		},
		{
			name:     "greater than operator",
			query:    Query{Field: "count", Operator: QueryOperatorGreaterThan, Value: 10},
			expected: map[string]interface{}{"$gt": 10},
		},
		{
			name:     "greater than or equal operator",
			query:    Query{Field: "count", Operator: QueryOperatorGreaterThanOrEqual, Value: 10},
			expected: map[string]interface{}{"$gte": 10},
		},
		{
			name:     "less than operator",
			query:    Query{Field: "score", Operator: QueryOperatorLessThan, Value: 100},
			expected: map[string]interface{}{"$lt": 100},
		},
		{
			name:     "less than or equal operator",
			query:    Query{Field: "score", Operator: QueryOperatorLessThanOrEqual, Value: 100},
			expected: map[string]interface{}{"$lte": 100},
		},
		{
			name:     "contains any operator",
			query:    Query{Field: "tags", Operator: QueryOperatorContainsAny, Value: []string{"a", "b"}},
			expected: map[string]interface{}{"$in": []string{"a", "b"}},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := buildPineconeCondition(tt.query)
			assert.Equal(t, tt.expected, result)
		})
	}
}

func TestMatchesQueries(t *testing.T) {
	tests := []struct {
		name     string
		props    map[string]interface{}
		queries  []Query
		expected bool
	}{
		{
			name:     "empty queries matches all",
			props:    map[string]interface{}{"type": "document"},
			queries:  []Query{},
			expected: true,
		},
		{
			name:     "equal match",
			props:    map[string]interface{}{"type": "document"},
			queries:  []Query{{Field: "type", Operator: QueryOperatorEqual, Value: "document"}},
			expected: true,
		},
		{
			name:     "equal no match",
			props:    map[string]interface{}{"type": "document"},
			queries:  []Query{{Field: "type", Operator: QueryOperatorEqual, Value: "image"}},
			expected: false,
		},
		{
			name:     "not equal match",
			props:    map[string]interface{}{"type": "document"},
			queries:  []Query{{Field: "type", Operator: QueryOperatorNotEqual, Value: "image"}},
			expected: true,
		},
		{
			name:     "is null match",
			props:    map[string]interface{}{"type": "document"},
			queries:  []Query{{Field: "author", Operator: QueryOperatorIsNull, Value: nil}},
			expected: true,
		},
		{
			name:     "is not null match",
			props:    map[string]interface{}{"type": "document", "author": "alice"},
			queries:  []Query{{Field: "author", Operator: QueryOperatorIsNotNull, Value: nil}},
			expected: true,
		},
		{
			name:  "multiple queries all match",
			props: map[string]interface{}{"type": "document", "public": true},
			queries: []Query{
				{Field: "type", Operator: QueryOperatorEqual, Value: "document"},
				{Field: "public", Operator: QueryOperatorEqual, Value: true},
			},
			expected: true,
		},
		{
			name:  "multiple queries one fails",
			props: map[string]interface{}{"type": "document", "public": false},
			queries: []Query{
				{Field: "type", Operator: QueryOperatorEqual, Value: "document"},
				{Field: "public", Operator: QueryOperatorEqual, Value: true},
			},
			expected: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := matchesQueries(tt.props, tt.queries)
			assert.Equal(t, tt.expected, result)
		})
	}
}

func TestFilterPropertiesPinecone(t *testing.T) {
	props := map[string]interface{}{
		"type":   "document",
		"author": "alice",
		"size":   1024,
		"public": true,
	}

	tests := []struct {
		name         string
		selectFields []string
		expected     map[string]interface{}
	}{
		{
			name:         "empty select returns all",
			selectFields: []string{},
			expected:     props,
		},
		{
			name:         "select single field",
			selectFields: []string{"type"},
			expected:     map[string]interface{}{"type": "document"},
		},
		{
			name:         "select multiple fields",
			selectFields: []string{"type", "author"},
			expected:     map[string]interface{}{"type": "document", "author": "alice"},
		},
		{
			name:         "select non-existent field",
			selectFields: []string{"missing"},
			expected:     map[string]interface{}{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := filterPropertiesPinecone(props, tt.selectFields)
			assert.Equal(t, tt.expected, result)
		})
	}
}

// ============================================================================
// INTEGRATION TESTS (require real Pinecone instance)
// ============================================================================

func TestPineconeStore_Integration(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	setup := NewPineconeTestSetup(t)
	defer setup.Cleanup(t)

	// Test Ping
	err := setup.Store.Ping(setup.ctx)
	require.NoError(t, err)

	// Test Add and GetChunk
	key := generateUUID()
	embedding := generateTestEmbedding(PineconeTestDimension)
	metadata := map[string]interface{}{
		"type":   "document",
		"author": "test",
	}

	err = setup.Store.Add(setup.ctx, PineconeTestNamespace, key, embedding, metadata)
	require.NoError(t, err)

	// Wait for eventual consistency
	time.Sleep(2 * time.Second)

	result, err := setup.Store.GetChunk(setup.ctx, PineconeTestNamespace, key)
	require.NoError(t, err)
	assert.Equal(t, key, result.ID)
	assert.Equal(t, "document", result.Properties["type"])
	assert.Equal(t, "test", result.Properties["author"])

	// Test GetChunks
	key2 := generateUUID()
	err = setup.Store.Add(setup.ctx, PineconeTestNamespace, key2, generateTestEmbedding(PineconeTestDimension), map[string]interface{}{"type": "image"})
	require.NoError(t, err)

	time.Sleep(2 * time.Second)

	results, err := setup.Store.GetChunks(setup.ctx, PineconeTestNamespace, []string{key, key2})
	require.NoError(t, err)
	assert.Len(t, results, 2)
}

func TestPineconeStore_VectorSearch(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	setup := NewPineconeTestSetup(t)
	defer setup.Cleanup(t)

	// Add test vectors
	emb := generateTestEmbedding(PineconeTestDimension)
	err := setup.Store.Add(setup.ctx, PineconeTestNamespace, generateUUID(), emb, map[string]interface{}{"type": "tech"})
	require.NoError(t, err)

	err = setup.Store.Add(setup.ctx, PineconeTestNamespace, generateUUID(), generateTestEmbedding(PineconeTestDimension), map[string]interface{}{"type": "sports"})
	require.NoError(t, err)

	// Wait for eventual consistency
	time.Sleep(3 * time.Second)

	// Test vector similarity search
	results, err := setup.Store.GetNearest(setup.ctx, PineconeTestNamespace, emb, nil, []string{"type"}, 0.1, 10)
	require.NoError(t, err)
	assert.GreaterOrEqual(t, len(results), 1)

	if len(results) > 0 {
		require.NotNil(t, results[0].Score)
	}

	// Test with filter
	queries := []Query{{Field: "type", Operator: QueryOperatorEqual, Value: "tech"}}
	results, err = setup.Store.GetNearest(setup.ctx, PineconeTestNamespace, emb, queries, []string{"type"}, 0.1, 10)
	require.NoError(t, err)
	for _, result := range results {
		assert.Equal(t, "tech", result.Properties["type"])
	}
}

func TestPineconeStore_Delete(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	setup := NewPineconeTestSetup(t)
	defer setup.Cleanup(t)

	// Add a vector
	key := generateUUID()
	err := setup.Store.Add(setup.ctx, PineconeTestNamespace, key, generateTestEmbedding(PineconeTestDimension), map[string]interface{}{"type": "to-delete"})
	require.NoError(t, err)

	time.Sleep(2 * time.Second)

	// Verify it exists
	_, err = setup.Store.GetChunk(setup.ctx, PineconeTestNamespace, key)
	require.NoError(t, err)

	// Delete it
	err = setup.Store.Delete(setup.ctx, PineconeTestNamespace, key)
	require.NoError(t, err)

	time.Sleep(2 * time.Second)

	// Verify it's gone
	_, err = setup.Store.GetChunk(setup.ctx, PineconeTestNamespace, key)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "not found")
}

func TestPineconeStore_ErrorHandling(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	setup := NewPineconeTestSetup(t)
	defer setup.Cleanup(t)

	// Test GetChunk with non-existent ID
	_, err := setup.Store.GetChunk(setup.ctx, PineconeTestNamespace, generateUUID())
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "not found")

	// Test Add with empty ID
	err = setup.Store.Add(setup.ctx, PineconeTestNamespace, "", generateTestEmbedding(PineconeTestDimension), map[string]interface{}{"type": "test"})
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "id is required")

	// Test Delete with empty ID
	err = setup.Store.Delete(setup.ctx, PineconeTestNamespace, "")
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "id is required")
}

func TestPineconeStore_SemanticCacheWorkflow(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	setup := NewPineconeTestSetup(t)
	defer setup.Cleanup(t)

	// Simulate a semantic cache workflow
	cacheEntries := []struct {
		key       string
		embedding []float32
		metadata  map[string]interface{}
	}{
		{
			generateUUID(),
			generateTestEmbedding(PineconeTestDimension),
			map[string]interface{}{
				"request_hash": "abc123",
				"user":         "u1",
				"lang":         "en",
				"response":     "answer1",
			},
		},
		{
			generateUUID(),
			generateTestEmbedding(PineconeTestDimension),
			map[string]interface{}{
				"request_hash": "def456",
				"user":         "u1",
				"lang":         "es",
				"response":     "answer2",
			},
		},
	}

	// Add cache entries
	for _, entry := range cacheEntries {
		err := setup.Store.Add(setup.ctx, PineconeTestNamespace, entry.key, entry.embedding, entry.metadata)
		require.NoError(t, err)
	}

	time.Sleep(3 * time.Second)

	// Test semantic search with user filter
	userFilter := []Query{{Field: "user", Operator: QueryOperatorEqual, Value: "u1"}}
	results, err := setup.Store.GetNearest(setup.ctx, PineconeTestNamespace, cacheEntries[0].embedding, userFilter, []string{"request_hash", "user", "lang", "response"}, 0.1, 10)
	require.NoError(t, err)
	assert.GreaterOrEqual(t, len(results), 1)

	// Verify user filter worked
	for _, result := range results {
		assert.Equal(t, "u1", result.Properties["user"])
	}
}

// ============================================================================
// INTERFACE COMPLIANCE TESTS
// ============================================================================

func TestPineconeStore_InterfaceCompliance(t *testing.T) {
	var _ VectorStore = (*PineconeStore)(nil)
}

func TestVectorStoreFactory_Pinecone(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration tests in short mode")
	}

	apiKey := schemas.NewEnvVar(getEnvWithDefault("PINECONE_API_KEY", PineconeTestDefaultAPIKey))
	indexHost := schemas.NewEnvVar(getEnvWithDefault("PINECONE_INDEX_HOST", PineconeTestDefaultIndexHost))

	logger := bifrost.NewDefaultLogger(schemas.LogLevelInfo)
	config := &Config{
		Enabled: true,
		Type:    VectorStoreTypePinecone,
		Config: PineconeConfig{
			APIKey:    *apiKey,
			IndexHost: *indexHost,
		},
	}

	store, err := NewVectorStore(context.Background(), config, logger)
	if err != nil {
		t.Skipf("Could not create Pinecone store: %v", err)
	}
	defer store.Close(context.Background(), PineconeTestNamespace)

	pineconeStore, ok := store.(*PineconeStore)
	assert.True(t, ok)
	assert.NotNil(t, pineconeStore)
}