package server

import (
	"context"
	"errors"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/mem9-ai/drive9/pkg/backend"
	"github.com/mem9-ai/drive9/pkg/embedding"
	"github.com/mem9-ai/drive9/pkg/meta"
)

func TestNewWithConfigWiresEmbeddingResolver(t *testing.T) {
	db := newTenantDeleteDBInfo(t)
	runtime := &backend.AppEmbeddingRuntime{Client: embedding.NopClient{}, Generation: 7, ConfigRevision: 11}
	provider := &EmbeddingDefaultProvider{APIBase: "https://api.example.com", APIKey: "secret", Model: "model"}
	srv := NewWithConfig(Config{
		Meta: db.Meta, Pool: db.Pool,
		EmbeddingDefaultRuntime: runtime, EmbeddingDefaultProvider: provider,
	})
	t.Cleanup(srv.Close)

	if srv.embeddingResolver == nil {
		t.Fatal("NewWithConfig did not wire the embedding resolver")
	}
	resolved, err := srv.embeddingResolver.Resolve(context.Background(), "absent-tenant")
	if err != nil {
		t.Fatal(err)
	}
	if resolved.state != embeddingResolveReady || resolved.runtime.Generation != runtime.Generation {
		t.Fatalf("resolved default = state %d generation %d, want ready generation %d",
			resolved.state, resolved.runtime.Generation, runtime.Generation)
	}
	if srv.embeddingDefaultView == nil || srv.embeddingDefaultView.Model != provider.Model {
		t.Fatalf("default provider view = %#v, want model %q", srv.embeddingDefaultView, provider.Model)
	}
}

type fakeEmbeddingConfigStore struct {
	mu    sync.Mutex
	row   *meta.TenantEmbeddingConfig
	err   error
	calls int
	start chan struct{}
	block chan struct{}
}

func (s *fakeEmbeddingConfigStore) GetTenantEmbeddingConfig(context.Context, string) (*meta.TenantEmbeddingConfig, error) {
	if s.start != nil {
		select {
		case s.start <- struct{}{}:
		default:
		}
	}
	if s.block != nil {
		<-s.block
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	s.calls++
	if s.err != nil {
		return nil, s.err
	}
	if s.row == nil {
		return nil, meta.ErrNotFound
	}
	copy := *s.row
	copy.APIKeyCipher = append([]byte(nil), s.row.APIKeyCipher...)
	return &copy, nil
}

type fakeEmbeddingDecryptor struct{ err error }

func (d fakeEmbeddingDecryptor) Decrypt(_ context.Context, cipher []byte) ([]byte, error) {
	if d.err != nil {
		return nil, d.err
	}
	return append([]byte(nil), cipher...), nil
}

func TestEmbeddingConfigResolverStatesAndCache(t *testing.T) {
	defaultRuntime := &backend.AppEmbeddingRuntime{Client: embedding.NopClient{}, Generation: 9, ConfigRevision: 11}

	t.Run("confirmed absent falls back and caches", func(t *testing.T) {
		store := &fakeEmbeddingConfigStore{}
		r := newEmbeddingConfigResolver(store, fakeEmbeddingDecryptor{}, defaultRuntime, time.Minute)
		t.Cleanup(r.Close)
		for range 2 {
			got, err := r.Resolve(context.Background(), "tenant")
			if err != nil || got.state != embeddingResolveReady || got.runtime.Generation != 9 || got.runtime.TenantManaged {
				t.Fatalf("Resolve = %+v, %v", got, err)
			}
		}
		if store.calls != 1 {
			t.Fatalf("store calls = %d, want 1", store.calls)
		}
	})

	t.Run("off marker blocks fallback", func(t *testing.T) {
		store := &fakeEmbeddingConfigStore{row: &meta.TenantEmbeddingConfig{TenantID: "tenant", Generation: 2, ConfigRevision: 3}}
		r := newEmbeddingConfigResolver(store, fakeEmbeddingDecryptor{}, defaultRuntime, time.Minute)
		t.Cleanup(r.Close)
		got, err := r.Resolve(context.Background(), "tenant")
		if err != nil || got.state != embeddingResolveDisabled {
			t.Fatalf("Resolve = %+v, %v", got, err)
		}
	})

	t.Run("meta and decrypt errors fail closed", func(t *testing.T) {
		for _, tc := range []struct {
			name  string
			store *fakeEmbeddingConfigStore
			dec   fakeEmbeddingDecryptor
		}{
			{name: "meta", store: &fakeEmbeddingConfigStore{err: errors.New("meta unavailable")}},
			{name: "decrypt", store: &fakeEmbeddingConfigStore{row: &meta.TenantEmbeddingConfig{
				TenantID: "tenant", Enabled: true, APIBase: "https://api.example.com/v1",
				APIKeyCipher: []byte("cipher"), Model: "model", Generation: 1, ConfigRevision: 1,
			}}, dec: fakeEmbeddingDecryptor{err: errors.New("decrypt failed")}},
		} {
			t.Run(tc.name, func(t *testing.T) {
				r := newEmbeddingConfigResolver(tc.store, tc.dec, defaultRuntime, time.Minute)
				t.Cleanup(r.Close)
				if _, err := r.Resolve(context.Background(), "tenant"); err == nil {
					t.Fatal("Resolve succeeded")
				}
			})
		}
	})
}

func TestEmbeddingConfigResolverEvictDropsCachedResolution(t *testing.T) {
	store := &fakeEmbeddingConfigStore{}
	r := newEmbeddingConfigResolver(store, fakeEmbeddingDecryptor{}, nil, time.Minute)
	t.Cleanup(r.Close)
	if _, err := r.Resolve(context.Background(), "tenant"); err != nil {
		t.Fatal(err)
	}
	r.Evict("tenant")
	if _, err := r.Resolve(context.Background(), "tenant"); err != nil {
		t.Fatal(err)
	}
	if store.calls != 2 {
		t.Fatalf("store calls = %d, want 2", store.calls)
	}
}

func TestEmbeddingConfigResolverEvictFencesInFlightResolution(t *testing.T) {
	store := &fakeEmbeddingConfigStore{start: make(chan struct{}, 1), block: make(chan struct{})}
	r := newEmbeddingConfigResolver(store, fakeEmbeddingDecryptor{}, nil, time.Minute)
	t.Cleanup(r.Close)
	errCh := make(chan error, 1)
	go func() {
		_, err := r.Resolve(context.Background(), "tenant")
		errCh <- err
	}()
	select {
	case <-store.start:
	case <-time.After(time.Second):
		t.Fatal("resolve did not reach store")
	}
	r.Evict("tenant")
	close(store.block)
	if err := <-errCh; err == nil || !strings.Contains(err.Error(), "changed during resolution") {
		t.Fatalf("in-flight Resolve error = %v, want generation fence", err)
	}
	if len(r.cache) != 0 {
		t.Fatalf("stale resolution populated cache: %+v", r.cache)
	}
}

func TestEmbeddingConfigResolverCachesFailureForShortBackoff(t *testing.T) {
	store := &fakeEmbeddingConfigStore{err: errors.New("meta unavailable")}
	r := newEmbeddingConfigResolver(store, fakeEmbeddingDecryptor{}, nil, time.Minute)
	t.Cleanup(r.Close)
	for range 2 {
		if _, err := r.Resolve(context.Background(), "tenant"); err == nil {
			t.Fatal("Resolve succeeded")
		}
	}
	if store.calls != 1 {
		t.Fatalf("store calls during failure backoff = %d, want 1", store.calls)
	}
	r.mu.Lock()
	r.cache["tenant"] = embeddingConfigCacheEntry{err: store.err, expiresAt: time.Now().Add(-time.Second)}
	r.mu.Unlock()
	if _, err := r.Resolve(context.Background(), "tenant"); err == nil {
		t.Fatal("Resolve after failure expiry succeeded")
	}
	if store.calls != 2 {
		t.Fatalf("store calls after failure expiry = %d, want 2", store.calls)
	}
}

func TestEmbeddingConfigResolverJanitorDropsExpiredRuntime(t *testing.T) {
	store := &fakeEmbeddingConfigStore{row: &meta.TenantEmbeddingConfig{
		TenantID: "tenant", Enabled: true, APIBase: "https://api.example.com/v1",
		APIKeyCipher: []byte("plaintext-test-key"), Model: "model", Generation: 1, ConfigRevision: 1,
	}}
	r := newEmbeddingConfigResolver(store, fakeEmbeddingDecryptor{}, nil, 20*time.Millisecond)
	t.Cleanup(r.Close)
	resolved, err := r.Resolve(context.Background(), "tenant")
	if err != nil {
		t.Fatal(err)
	}
	if !resolved.runtime.TenantManaged {
		t.Fatal("tenant config resolver returned a process-default runtime identity")
	}
	deadline := time.Now().Add(time.Second)
	for {
		r.mu.Lock()
		_, cached := r.cache["tenant"]
		r.mu.Unlock()
		if !cached {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("janitor retained expired plaintext-bearing runtime")
		}
		time.Sleep(5 * time.Millisecond)
	}
}