plugin.go

package natsstore

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/open-policy-agent/opa/v1/ast"
	"github.com/open-policy-agent/opa/v1/logging"
	"github.com/open-policy-agent/opa/v1/plugins"
	"github.com/open-policy-agent/opa/v1/rego"
	"github.com/open-policy-agent/opa/v1/storage"
	"github.com/open-policy-agent/opa/v1/types"
	"github.com/tidwall/gjson"
	"github.com/tidwall/sjson"
)

const (
	PluginName             = "nats"
	WatchBucketBuiltinName = "nats.kv.watch_bucket"
	GetDataBuiltinName     = "nats.kv.get_data"
)

func getDataCacheKey(bucketName string) string {
	return fmt.Sprintf("nats.kv.%s", bucketName)
}

// PluginFactory creates and manages NATS K/V store plugin instances.
type PluginFactory struct {
	logger        logging.Logger
	originalStore storage.Store
	// Keep a reference to the current plugin's bucket data manager for builtin functions
	currentPlugin *Plugin
}

// NewPluginFactory creates a new NATS K/V store plugin factory.
func NewPluginFactory() *PluginFactory {
	factory := &PluginFactory{}
	factory.RegisterBuiltin()
	return factory
}

func (f *PluginFactory) watchBucketBuiltin(bctx rego.BuiltinContext, inputTerm *ast.Term) (*ast.Term, error) {
	bucketNameValue, ok := inputTerm.Value.(ast.String)
	if !ok {
		return nil, fmt.Errorf("expected string bucket name, got %T", inputTerm.Value)
	}
	bucketName := string(bucketNameValue)

	// If no bucket data manager is available (e.g., during tests), return false
	if f == nil || f.getBucketDataManager() == nil {
		return ast.BooleanTerm(false), nil
	}

	// Check if bucket is already being watched
	if f.getBucketDataManager().watcherManager.HasWatcher(bucketName) {
		return ast.BooleanTerm(true), nil
	}

	// Bucket not watched yet - load data into cache and start watching
	bucketGjson, err := f.loadTenantAsGJSON(bctx.Context, bucketName)
	if err != nil {
		// Surface the error instead of failing open to "not watched". With the
		// single muxed bucket, an absent *tenant* yields zero keys (handled by
		// loadTenantAsGJSON returning a Null result, no error), so a
		// nats.ErrBucketNotFound here means the *configured* bucket is missing —
		// a deployment misconfiguration. Returning false would silently evaluate
		// policy against missing data; this is consistent with get_data, which
		// also returns the error.
		return nil, fmt.Errorf("failed to load bucket data: %w", err)
	}

	// Store gjson data in ND builtin cache for this query
	cacheKey := getDataCacheKey(bucketName)
	bctx.Cache.Put(cacheKey, bucketGjson)

	// Start watching asynchronously (no OPA store writes during query)
	go func() {
		bdm := f.getBucketDataManager()
		if bdm != nil {
			if _, err := bdm.watcherManager.GetOrCreateWatcher(bucketName, f.originalStore); err != nil {
				// Log error but don't fail the query
				f.logger.Error("Failed to start watcher for bucket %s: %v", bucketName, err)
			}
		}
	}()

	return ast.BooleanTerm(false), nil
}

func (f *PluginFactory) getDataBuiltin(bctx rego.BuiltinContext, bucketTerm *ast.Term, keyTerm *ast.Term) (*ast.Term, error) {
	bucketNameValue, ok := bucketTerm.Value.(ast.String)
	if !ok {
		return nil, fmt.Errorf("expected string bucket name, got %T", bucketTerm.Value)
	}
	bucketName := string(bucketNameValue)

	keyValue, ok := keyTerm.Value.(ast.String)
	if !ok {
		return nil, fmt.Errorf("expected string key, got %T", keyTerm.Value)
	}
	dotNotationKey := string(keyValue)

	// Try to get gjson data from cache first
	cacheKey := getDataCacheKey(bucketName)
	if cachedValue, ok := bctx.Cache.Get(cacheKey); ok {
		if bucketGjson, ok := cachedValue.(*gjson.Result); ok {
			result := bucketGjson.Get(dotNotationKey)
			if result.Exists() {
				return f.gjsonResultToASTTerm(result), nil
			}
			return ast.NullTerm(), nil
		}
	}

	// Cache miss - load directly from NATS with warning
	f.logger.Warn("Warning: Cache miss for bucket %s, key %s. Loading directly from NATS.\n", bucketName, dotNotationKey)

	bucketGjson, err := f.loadTenantAsGJSON(bctx.Context, bucketName)
	if err != nil {
		return nil, fmt.Errorf("failed to load bucket data: %w", err)
	}

	// Cache the gjson data for subsequent calls in this query
	bctx.Cache.Put(cacheKey, bucketGjson)

	// Get the requested value using gjson
	result := bucketGjson.Get(dotNotationKey)
	if result.Exists() {
		return f.gjsonResultToASTTerm(result), nil
	}

	// Key not found
	return ast.NullTerm(), nil
}

// buildTenantJSON builds a tenant-relative JSON document from muxed keys.
// Each key is "<tenant>.<sub...>"; the "<tenant>." prefix is stripped so the
// resulting tree is tenant-relative (e.g. "t1.members" -> "members"). Keys
// outside the tenant's subtree are skipped.
//
// get returns the raw value bytes for a full key, or nil to signal the key
// should be SKIPPED (e.g. its Get failed) — a skipped key is omitted entirely
// rather than written as an explicit null, since a key flipping to null can
// change an authz decision. A non-JSON value is stored as a JSON string,
// mirroring loadSingleKey, so this read path and the bulk-load path agree.
func buildTenantJSON(tenant string, keys []string, get func(key string) []byte) []byte {
	prefix := tenant + "."
	jsonBytes := []byte("{}")
	for _, k := range keys {
		sub, ok := strings.CutPrefix(k, prefix)
		if !ok || sub == "" {
			continue // not this tenant's key, or the bare tenant token
		}
		value := get(k)
		if value == nil {
			continue // get failed for this key — skip, don't inject null
		}
		if !json.Valid(value) {
			// Non-JSON stored value: keep it as a JSON string instead of
			// corrupting the document with raw bytes.
			if quoted, err := json.Marshal(string(value)); err == nil {
				value = quoted
			} else {
				continue
			}
		}
		next, err := sjson.SetRawBytes(jsonBytes, sub, value)
		if err != nil {
			continue // don't corrupt the document on a single bad key
		}
		jsonBytes = next
	}
	return jsonBytes
}

// loadTenantAsGJSON builds a tenant-relative gjson document by reading ONLY the
// tenant's slice of the muxed bucket (prefix-filtered), with the "<tenant>."
// prefix stripped so callers address keys tenant-relatively.
func (f *PluginFactory) loadTenantAsGJSON(ctx context.Context, tenant string) (*gjson.Result, error) {
	bdm := f.getBucketDataManager()
	if bdm == nil {
		return nil, fmt.Errorf("bucket data manager not available")
	}

	// List ONLY this tenant's keys (filtered watch) — never enumerate the bucket.
	keys, err := bdm.natsClient.tenantKeys(ctx, tenant)
	if err != nil {
		return nil, fmt.Errorf("failed to list keys for tenant %s: %w", tenant, err)
	}
	if len(keys) == 0 {
		return &gjson.Result{Type: gjson.Null}, nil
	}

	kv, err := bdm.natsClient.getBucket()
	if err != nil {
		return nil, fmt.Errorf("failed to get bucket: %w", err)
	}

	jsonBytes := buildTenantJSON(tenant, keys, func(key string) []byte {
		entry, err := kv.Get(key)
		if err != nil {
			// Return nil so buildTenantJSON omits the key instead of writing an
			// explicit null on a transient Get error.
			bdm.logger.Warn("Skipping key %s for tenant %s: get failed: %v", key, tenant, err)
			return nil
		}
		return entry.Value()
	})

	gjsonResult := gjson.ParseBytes(jsonBytes)
	return &gjsonResult, nil
}

func (f *PluginFactory) gjsonResultToASTTerm(result gjson.Result) *ast.Term {
	switch result.Type {
	case gjson.String:
		return ast.StringTerm(result.String())
	case gjson.Number:
		return ast.NumberTerm(json.Number(result.Raw))
	case gjson.True:
		return ast.BooleanTerm(true)
	case gjson.False:
		return ast.BooleanTerm(false)
	case gjson.JSON:
		// Parse as JSON and convert to AST
		var value interface{}
		if err := json.Unmarshal([]byte(result.Raw), &value); err != nil {
			return ast.StringTerm(result.Raw) // Fallback to string
		}
		// Convert interface{} to AST term manually
		return ast.NewTerm(ast.MustInterfaceToValue(value))
	default:
		return ast.NullTerm()
	}
}

func (f *PluginFactory) RegisterBuiltin() {
	// Register nats.kv.watch_bucket builtin
	rego.RegisterBuiltin1(
		&rego.Function{
			Name:             WatchBucketBuiltinName,
			Decl:             types.NewFunction(types.Args(types.S), types.B),
			Memoize:          false,
			Nondeterministic: true,
		},
		f.watchBucketBuiltin,
	)

	// Register nats.kv.get_data builtin
	rego.RegisterBuiltin2(
		&rego.Function{
			Name:             GetDataBuiltinName,
			Decl:             types.NewFunction(types.Args(types.S, types.S), types.A),
			Memoize:          true,
			Nondeterministic: true,
		},
		f.getDataBuiltin,
	)
}

func (f *PluginFactory) Store() storage.Store {
	// Return the original store since we're injecting data into it
	return f.originalStore
}

// Validate validates the plugin configuration.
func (f *PluginFactory) Validate(manager *plugins.Manager, config []byte) (any, error) {
	logger := manager.Logger()
	if logger == nil {
		logger = logging.New() // Create a default logger for testing scenarios
	}
	f.logger = logger

	// Start with empty config to properly validate required fields
	pluginConfig := &Config{}

	if len(config) > 0 {
		if err := json.Unmarshal(config, pluginConfig); err != nil {
			return nil, fmt.Errorf("failed to unmarshal plugin config: %w", err)
		}
	}

	// Validate required fields first
	if err := pluginConfig.Validate(); err != nil {
		return nil, fmt.Errorf("invalid plugin config: %w", err)
	}

	// Apply defaults for optional fields
	f.applyDefaults(pluginConfig)

	// Validate again with defaults applied
	if err := pluginConfig.ValidateWithDefaults(); err != nil {
		return nil, fmt.Errorf("invalid plugin config: %w", err)
	}

	logger.Info("Validated NATS K/V store plugin config, attempting to connect")

	// Store reference to original store
	f.originalStore = manager.Store

	// Create the bucket data manager during plugin creation
	_, err := NewNATSClient(pluginConfig, logger)
	if err != nil {
		logger.Error("Failed to connect to NATS: %v", err)
		return nil, fmt.Errorf("failed to connect to NATS: %w", err)
	}

	return pluginConfig, nil
}

// applyDefaults applies default values for optional configuration fields.
func (f *PluginFactory) applyDefaults(config *Config) {
	if config.TTL == 0 {
		config.TTL = Duration(10 * time.Minute)
	}
	if config.RefreshInterval == 0 {
		config.RefreshInterval = Duration(30 * time.Second)
	}
	if config.MaxReconnectAttempts == 0 {
		config.MaxReconnectAttempts = 10
	}
	if config.ReconnectWait == 0 {
		config.ReconnectWait = Duration(2 * time.Second)
	}
	if config.MaxBucketsWatchers == 0 {
		config.MaxBucketsWatchers = 10
	}
}

// New creates a new plugin instance.
func (f *PluginFactory) New(manager *plugins.Manager, config any) plugins.Plugin {
	logger := manager.Logger()
	if logger == nil {
		logger = logging.New() // Create a default logger for testing scenarios
	}
	pluginConfig, ok := config.(*Config)
	if !ok {
		logger.Error("Invalid config type for NATS K/V store plugin")
		return nil
	}

	// Create the bucket data manager during plugin creation
	bucketDataManager, err := NewBucketDataManager(pluginConfig, logger)
	if err != nil {
		logger.Error("Failed to create bucket data manager: %v", err)
		return nil
	}

	plugin := &Plugin{
		manager:           manager,
		config:            pluginConfig,
		bucketDataManager: bucketDataManager,
		logger:            logger,
		factory:           f,
	}

	// Set the current plugin reference for builtin functions
	f.currentPlugin = plugin

	return plugin
}

// getBucketDataManager returns the current plugin's bucket data manager for builtin functions
func (f *PluginFactory) getBucketDataManager() *BucketDataManager {
	if f.currentPlugin != nil {
		return f.currentPlugin.bucketDataManager
	}
	return nil
}

// Plugin represents the NATS K/V cache plugin instance.
type Plugin struct {
	manager           *plugins.Manager
	config            *Config
	bucketDataManager *BucketDataManager
	logger            logging.Logger
	factory           *PluginFactory
}

// Start initializes and starts the plugin.
func (p *Plugin) Start(ctx context.Context) error {
	p.logger.Info("Starting NATS K/V data injection plugin")

	if err := p.bucketDataManager.Start(ctx); err != nil {
		return fmt.Errorf("failed to start bucket data manager: %w", err)
	}

	// If we have a root bucket configured, load it immediately
	if p.config.RootTenant != "" {
		p.logger.Info("Loading root tenant: %s", p.config.RootTenant)
		if err := p.bucketDataManager.EnsureBucketLoaded(ctx, p.config.RootTenant, p.manager.Store, true); err != nil {
			return err
		}
	}

	p.logger.Info("NATS plugin started successfully")
	return nil
}

// Stop shuts down the plugin.
func (p *Plugin) Stop(ctx context.Context) {
	p.logger.Info("Stopping NATS K/V data injection plugin")

	if err := p.bucketDataManager.Stop(ctx); err != nil {
		p.logger.Error("Error stopping bucket data manager: %v", err)
	}

	p.logger.Info("NATS plugin stopped")
}

// Reconfigure updates the plugin configuration.
func (p *Plugin) Reconfigure(ctx context.Context, config any) {
	p.logger.Info("Reconfiguring NATS K/V data injection plugin")

	newConfig, ok := config.(*Config)
	if !ok {
		p.logger.Error("Invalid config type for reconfiguration")
		return
	}

	// Stop the current manager
	if err := p.bucketDataManager.Stop(ctx); err != nil {
		p.logger.Error("Error stopping bucket data manager during reconfiguration: %v", err)
	}

	// Update configuration
	p.config = newConfig

	// Create new bucket data manager with updated config
	bucketDataManager, err := NewBucketDataManager(newConfig, p.logger)
	if err != nil {
		p.logger.Error("Failed to create new bucket data manager during reconfiguration: %v", err)
		return
	}

	if err := bucketDataManager.Start(ctx); err != nil {
		p.logger.Error("Failed to start new bucket data manager during reconfiguration: %v", err)
		return
	}

	p.bucketDataManager = bucketDataManager

	p.logger.Info("NATS K/V data injection plugin reconfigured successfully")
}