package server

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"golang.org/x/sync/singleflight"

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

const (
	defaultEmbeddingConfigCacheTTL = 30 * time.Second
	embeddingConfigFailureBackoff  = time.Second
	defaultEmbeddingTimeout        = 20 * time.Second
)

type embeddingResolveState uint8

const (
	embeddingResolveNotConfigured embeddingResolveState = iota
	embeddingResolveReady
	embeddingResolveDisabled
)

type resolvedEmbeddingConfig struct {
	state   embeddingResolveState
	runtime backend.AppEmbeddingRuntime
}

type tenantEmbeddingConfigReader interface {
	GetTenantEmbeddingConfig(context.Context, string) (*meta.TenantEmbeddingConfig, error)
}

type embeddingConfigCacheEntry struct {
	resolved  resolvedEmbeddingConfig
	err       error
	expiresAt time.Time
}

type embeddingConfigResolver struct {
	store tenantEmbeddingConfigReader
	dec   extractKeyDecryptor
	def   *backend.AppEmbeddingRuntime
	ttl   time.Duration
	now   func() time.Time

	mu         sync.Mutex
	cache      map[string]embeddingConfigCacheEntry
	generation uint64
	closed     bool
	stop       chan struct{}
	done       chan struct{}
	group      singleflight.Group
}

func newEmbeddingConfigResolver(store tenantEmbeddingConfigReader, dec extractKeyDecryptor, def *backend.AppEmbeddingRuntime, ttl time.Duration) *embeddingConfigResolver {
	if ttl <= 0 {
		ttl = defaultEmbeddingConfigCacheTTL
	}
	r := &embeddingConfigResolver{
		store: store, dec: dec, def: def, ttl: ttl, now: time.Now,
		cache: make(map[string]embeddingConfigCacheEntry), stop: make(chan struct{}), done: make(chan struct{}),
	}
	go r.reapLoop()
	return r
}

func (r *embeddingConfigResolver) Resolve(ctx context.Context, tenantID string) (resolvedEmbeddingConfig, error) {
	if tenantID == "" {
		return resolvedEmbeddingConfig{}, fmt.Errorf("embedding config tenant id is required")
	}
	r.mu.Lock()
	if ent, ok := r.cache[tenantID]; ok && r.now().Before(ent.expiresAt) {
		r.mu.Unlock()
		return ent.resolved, ent.err
	}
	startGeneration := r.generation
	r.mu.Unlock()

	value, err, _ := r.group.Do(fmt.Sprintf("%d:%s", startGeneration, tenantID), func() (any, error) {
		resolved, resolveErr := r.resolve(ctx, tenantID)
		r.mu.Lock()
		defer r.mu.Unlock()
		if r.closed {
			return resolvedEmbeddingConfig{}, fmt.Errorf("embedding config resolver is closed")
		}
		if r.generation != startGeneration {
			return resolvedEmbeddingConfig{}, fmt.Errorf("embedding config changed during resolution")
		}
		ttl := r.ttl
		if resolveErr != nil {
			ttl = embeddingConfigFailureBackoff
		}
		r.cache[tenantID] = embeddingConfigCacheEntry{resolved: resolved, err: resolveErr, expiresAt: r.now().Add(ttl)}
		return resolved, resolveErr
	})
	if err != nil {
		return resolvedEmbeddingConfig{}, err
	}
	return value.(resolvedEmbeddingConfig), nil
}

func (r *embeddingConfigResolver) resolve(ctx context.Context, tenantID string) (resolvedEmbeddingConfig, error) {
	if r.store == nil {
		return resolvedEmbeddingConfig{}, fmt.Errorf("embedding config store is not configured")
	}
	row, err := r.store.GetTenantEmbeddingConfig(ctx, tenantID)
	if err != nil {
		if !errors.Is(err, meta.ErrNotFound) {
			return resolvedEmbeddingConfig{}, fmt.Errorf("get embedding config for tenant %s: %w", tenantID, err)
		}
		if r.def == nil || r.def.Client == nil {
			return resolvedEmbeddingConfig{state: embeddingResolveNotConfigured}, nil
		}
		return resolvedEmbeddingConfig{state: embeddingResolveReady, runtime: *r.def}, nil
	}
	if !row.Enabled {
		return resolvedEmbeddingConfig{state: embeddingResolveDisabled}, nil
	}
	if err := backend.ValidateTenantAPIBase(row.APIBase); err != nil {
		return resolvedEmbeddingConfig{}, fmt.Errorf("validate tenant %s embedding api base: %w", tenantID, err)
	}
	if r.dec == nil {
		return resolvedEmbeddingConfig{}, fmt.Errorf("embedding key decryptor is not configured")
	}
	key, err := r.dec.Decrypt(ctx, row.APIKeyCipher)
	if err != nil {
		return resolvedEmbeddingConfig{}, fmt.Errorf("decrypt tenant %s embedding api key: %w", tenantID, err)
	}
	client, err := embedding.NewOpenAIClient(embedding.OpenAIClientConfig{
		BaseURL: row.APIBase, APIKey: string(key), Model: row.Model,
		Dimensions: embedding.Dimensions, Timeout: defaultEmbeddingTimeout,
		Client: backend.NewTenantExtractorHTTPClient(defaultEmbeddingTimeout),
	})
	if err != nil {
		return resolvedEmbeddingConfig{}, fmt.Errorf("build tenant %s embedding client: %w", tenantID, err)
	}
	return resolvedEmbeddingConfig{state: embeddingResolveReady, runtime: backend.AppEmbeddingRuntime{
		Client: client, Generation: row.Generation, ConfigRevision: row.ConfigRevision, TenantManaged: true,
	}}, nil
}

func (r *embeddingConfigResolver) Evict(tenantID string) {
	if r == nil || tenantID == "" {
		return
	}
	r.mu.Lock()
	delete(r.cache, tenantID)
	r.generation++
	r.mu.Unlock()
}

func (r *embeddingConfigResolver) reapLoop() {
	interval := r.ttl / 2
	if interval <= 0 {
		interval = time.Second
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	defer close(r.done)
	for {
		select {
		case <-ticker.C:
			now := r.now()
			r.mu.Lock()
			for tenantID, ent := range r.cache {
				if !now.Before(ent.expiresAt) {
					delete(r.cache, tenantID)
				}
			}
			r.mu.Unlock()
		case <-r.stop:
			return
		}
	}
}

func (r *embeddingConfigResolver) Close() {
	if r == nil {
		return
	}
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return
	}
	r.closed = true
	close(r.stop)
	r.cache = make(map[string]embeddingConfigCacheEntry)
	r.mu.Unlock()
	<-r.done
}

// ConfigureEmbeddingResolver wires app-managed embedding resolution. The
// provider view is accepted only when a real default runtime exists, so GET
// cannot claim a default that workers cannot execute.
func (s *Server) ConfigureEmbeddingResolver(store tenantEmbeddingConfigReader, dec extractKeyDecryptor, def *backend.AppEmbeddingRuntime, provider *EmbeddingDefaultProvider) {
	r := newEmbeddingConfigResolver(store, dec, def, defaultEmbeddingConfigCacheTTL)
	s.embeddingResolver = r
	if s.pool != nil {
		s.pool.SetEmbeddingRuntimeResolver(func(ctx context.Context, tenantID string) (backend.AppEmbeddingRuntime, bool, error) {
			resolved, err := r.Resolve(ctx, tenantID)
			if err != nil {
				return backend.AppEmbeddingRuntime{}, false, err
			}
			return resolved.runtime, resolved.state == embeddingResolveReady, nil
		})
	}
	s.embeddingDefaultView = nil
	s.embeddingDefaultGeneration = 0
	if def != nil && def.Client != nil && provider != nil {
		view := *provider
		s.embeddingDefaultView = &view
		s.embeddingDefaultGeneration = def.Generation
	}
}