// Copyright 2016 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

// Package inmem implements an in-memory version of the policy engine's storage
// layer.
//
// The in-memory store is used as the default storage layer implementation. The
// in-memory store supports multi-reader/single-writer concurrency with
// rollback.
//
// Callers should assume the in-memory store does not make copies of written
// data. Once data is written to the in-memory store, it should not be modified
// (outside of calling Store.Write). Furthermore, data read from the in-memory
// store should be treated as read-only.
package inmem

import (
	"context"
	"fmt"
	"io"
	"path/filepath"
	"sync"
	"sync/atomic"

	"github.com/open-policy-agent/opa/internal/merge"
	"github.com/open-policy-agent/opa/v1/ast"
	"github.com/open-policy-agent/opa/v1/storage"
	"github.com/open-policy-agent/opa/v1/storage/internal/errors"
	"github.com/open-policy-agent/opa/v1/util"
)

// New returns an empty in-memory store.
func New() storage.Store {
	return NewWithOpts()
}

// NewWithOpts returns an empty in-memory store, with extra options passed.
func NewWithOpts(opts ...Opt) storage.Store {
	s := &store{
		triggers:              map[*handle]storage.TriggerConfig{},
		policies:              map[string][]byte{},
		roundTripOnWrite:      true,
		returnASTValuesOnRead: false,
	}

	for _, opt := range opts {
		opt(s)
	}

	if s.returnASTValuesOnRead {
		s.data = ast.NewObject()
		s.roundTripOnWrite = false
	} else {
		s.data = map[string]any{}
	}

	return s
}

// NewFromObject returns a new in-memory store from the supplied data object.
func NewFromObject(data map[string]any) storage.Store {
	return NewFromObjectWithOpts(data)
}

// NewFromObjectWithOpts returns a new in-memory store from the supplied data object, with the
// options passed.
func NewFromObjectWithOpts(data map[string]any, opts ...Opt) storage.Store {
	db := NewWithOpts(opts...)
	ctx := context.Background()
	txn, err := db.NewTransaction(ctx, storage.WriteParams)
	if err != nil {
		panic(err)
	}
	if err := db.Write(ctx, txn, storage.AddOp, storage.RootPath, data); err != nil {
		panic(err)
	}
	if err := db.Commit(ctx, txn); err != nil {
		panic(err)
	}
	return db
}

// NewFromReader returns a new in-memory store from a reader that produces a
// JSON serialized object. This function is for test purposes.
func NewFromReader(r io.Reader) storage.Store {
	return NewFromReaderWithOpts(r)
}

// NewFromReader returns a new in-memory store from a reader that produces a
// JSON serialized object, with extra options. This function is for test purposes.
func NewFromReaderWithOpts(r io.Reader, opts ...Opt) storage.Store {
	var data map[string]any
	if err := util.NewJSONDecoder(r).Decode(&data); err != nil {
		panic(err)
	}
	return NewFromObjectWithOpts(data, opts...)
}

// NewFromASTObject returns a new in-memory store from the supplied AST object, with
// [OptReturnASTValuesOnRead] enabled. This allows avoiding the overhead of an extra AST
// -> map[string]any -> AST round trip for callers whose data already exists in AST form.
// Note that data passed is **not** copied and it is the responsibility of the caller to
// ensure either that ownership of the data is transferred fully  to the store, or when
// that's not possible, that a deep copy of the original data is passed.
func NewFromASTObject(data ast.Object) storage.Store {
	return &store{
		data:                  data,
		triggers:              map[*handle]storage.TriggerConfig{},
		policies:              map[string][]byte{},
		returnASTValuesOnRead: true,
	}
}

type store struct {
	rmu      sync.RWMutex                      // reader-writer lock
	wmu      sync.Mutex                        // writer lock
	xid      atomic.Uint64                     // last generated transaction id
	data     any                               // raw or AST data
	policies map[string][]byte                 // raw policies
	triggers map[*handle]storage.TriggerConfig // registered triggers

	// roundTripOnWrite, if true, means that every call to Write round trips the
	// data through JSON before adding the data to the store. Defaults to true.
	roundTripOnWrite bool

	// returnASTValuesOnRead, if true, means that the store will eagerly convert data to AST values,
	// and return them on Read.
	// FIXME: naming(?)
	returnASTValuesOnRead bool
}

type handle struct {
	db *store
}

func (db *store) NewTransaction(_ context.Context, params ...storage.TransactionParams) (storage.Transaction, error) {
	txn := &transaction{
		xid: db.xid.Add(1),
		db:  db,
	}

	if len(params) > 0 {
		txn.write = params[0].Write
		txn.context = params[0].Context
	}

	if txn.write {
		db.wmu.Lock()
	} else {
		db.rmu.RLock()
	}

	return txn, nil
}

// Truncate implements the storage.Store interface. This method must be called within a transaction.
func (db *store) Truncate(ctx context.Context, txn storage.Transaction, params storage.TransactionParams, it storage.Iterator) error {
	var update *storage.Update
	var err error

	underlying, err := db.underlying(txn)
	if err != nil {
		return err
	}

	mergedData := map[string]any{}

	for {
		if update, err = it.Next(); err != nil {
			break
		}

		if update.IsPolicy {
			err = underlying.UpsertPolicy(update.Path.PolicyID(), update.Value)
			if err != nil {
				return err
			}
		} else {
			var value any
			if err = util.Unmarshal(update.Value, &value); err != nil {
				return err
			}

			// Do not round trip via String() to avoid URL encoding.
			key := []string(update.Path)

			if value != nil {
				obj, err := mktree(key, value)
				if err != nil {
					return err
				}

				merged, ok := merge.InterfaceMaps(mergedData, obj)
				if !ok {
					return fmt.Errorf("failed to insert data file from path %s", filepath.Join(key...))
				}
				mergedData = merged
			}
		}
	}

	// err is known not to be nil at this point, as it getting assigned
	// a non-nil value is the only way the loop above can exit.
	if err != io.EOF {
		return err
	}

	// For backwards compatibility, check if `RootOverwrite` was configured.
	if params.RootOverwrite {
		return underlying.Write(storage.AddOp, storage.RootPath, mergedData)
	}

	for _, root := range params.BasePaths {
		newPath, ok := storage.ParsePathEscaped("/" + root)
		if !ok {
			return fmt.Errorf("storage path invalid: %v", newPath)
		}

		if value, ok := lookup(newPath, mergedData); ok {
			if len(newPath) > 0 {
				if err := storage.MakeDir(ctx, db, txn, newPath[:len(newPath)-1]); err != nil {
					return err
				}
			}
			if err := underlying.Write(storage.AddOp, newPath, value); err != nil {
				return err
			}
		}
	}
	return nil
}

func (db *store) Commit(ctx context.Context, txn storage.Transaction) error {
	underlying, err := db.underlying(txn)
	if err != nil {
		return err
	}
	if underlying.write {
		db.rmu.Lock()
		event := underlying.Commit()
		db.runOnCommitTriggers(ctx, txn, event)
		// Mark the transaction stale after executing triggers, so they can
		// perform store operations if needed.
		underlying.stale = true
		db.rmu.Unlock()
		db.wmu.Unlock()
	} else {
		db.rmu.RUnlock()
	}
	return nil
}

func (db *store) Abort(_ context.Context, txn storage.Transaction) {
	underlying, err := db.underlying(txn)
	if err != nil {
		panic(err)
	}
	underlying.stale = true
	if underlying.write {
		db.wmu.Unlock()
	} else {
		db.rmu.RUnlock()
	}
}

func (db *store) ListPolicies(_ context.Context, txn storage.Transaction) ([]string, error) {
	underlying, err := db.underlying(txn)
	if err != nil {
		return nil, err
	}
	return underlying.ListPolicies(), nil
}

func (db *store) GetPolicy(_ context.Context, txn storage.Transaction, id string) ([]byte, error) {
	underlying, err := db.underlying(txn)
	if err != nil {
		return nil, err
	}
	return underlying.GetPolicy(id)
}

func (db *store) UpsertPolicy(_ context.Context, txn storage.Transaction, id string, bs []byte) error {
	underlying, err := db.underlying(txn)
	if err != nil {
		return err
	}
	return underlying.UpsertPolicy(id, bs)
}

func (db *store) DeletePolicy(_ context.Context, txn storage.Transaction, id string) error {
	underlying, err := db.underlying(txn)
	if err != nil {
		return err
	}
	if _, err := underlying.GetPolicy(id); err != nil {
		return err
	}
	return underlying.DeletePolicy(id)
}

func (db *store) Register(_ context.Context, txn storage.Transaction, config storage.TriggerConfig) (storage.TriggerHandle, error) {
	underlying, err := db.underlying(txn)
	if err != nil {
		return nil, err
	}
	if !underlying.write {
		return nil, &storage.Error{
			Code:    storage.InvalidTransactionErr,
			Message: "triggers must be registered with a write transaction",
		}
	}
	h := &handle{db}
	db.triggers[h] = config
	return h, nil
}

func (db *store) MakeDir(_ context.Context, txn storage.Transaction, path storage.Path) error {
	underlying, err := db.underlying(txn)
	if err != nil {
		return err
	}
	return underlying.makeDir(path)
}

func (db *store) Read(_ context.Context, txn storage.Transaction, path storage.Path) (any, error) {
	underlying, err := db.underlying(txn)
	if err != nil {
		return nil, err
	}

	return underlying.Read(path)
}

func (db *store) Write(_ context.Context, txn storage.Transaction, op storage.PatchOp, path storage.Path, value any) error {
	underlying, err := db.underlying(txn)
	if err != nil {
		return err
	}

	if db.returnASTValuesOnRead || !util.NeedsRoundTrip(value) {
		// Fast path when value is nil, bool, string or json.Number.
		return underlying.Write(op, path, value)
	}

	val := util.Reference(value)
	if db.roundTripOnWrite {
		if err := util.RoundTripFast(val); err != nil {
			return err
		}
	}

	return underlying.Write(op, path, *val)
}

func (h *handle) Unregister(_ context.Context, txn storage.Transaction) {
	underlying, err := h.db.underlying(txn)
	if err != nil {
		panic(err)
	}
	if !underlying.write {
		panic(&storage.Error{
			Code:    storage.InvalidTransactionErr,
			Message: "triggers must be unregistered with a write transaction",
		})
	}
	delete(h.db.triggers, h)
}

func (db *store) runOnCommitTriggers(ctx context.Context, txn storage.Transaction, event storage.TriggerEvent) {
	// While it's unlikely, the API allows one trigger to be configured to want
	// data conversion, and another that doesn't. So let's handle that properly.
	var wantsDataConversion bool
	if db.returnASTValuesOnRead && len(event.Data) > 0 {
		for _, t := range db.triggers {
			if !t.SkipDataConversion {
				wantsDataConversion = true
				break
			}
		}
	}

	var converted storage.TriggerEvent
	if wantsDataConversion {
		converted = storage.TriggerEvent{
			Policy:  event.Policy,
			Data:    make([]storage.DataEvent, 0, len(event.Data)),
			Context: event.Context,
		}

		for _, dataEvent := range event.Data {
			if astData, ok := dataEvent.Data.(ast.Value); ok {
				jsn, err := ast.ValueToInterface(astData, illegalResolver{})
				if err != nil {
					panic(err)
				}
				converted.Data = append(converted.Data, storage.DataEvent{
					Path:    dataEvent.Path,
					Data:    jsn,
					Removed: dataEvent.Removed,
				})
			}
		}
	}

	for _, t := range db.triggers {
		if wantsDataConversion && !t.SkipDataConversion {
			t.OnCommit(ctx, txn, converted)
		} else {
			t.OnCommit(ctx, txn, event)
		}
	}
}

type illegalResolver struct{}

func (illegalResolver) Resolve(ref ast.Ref) (any, error) {
	return nil, fmt.Errorf("illegal value: %v", ref)
}

func (db *store) underlying(txn storage.Transaction) (*transaction, error) {
	underlying, ok := txn.(*transaction)
	if !ok {
		return nil, &storage.Error{
			Code:    storage.InvalidTransactionErr,
			Message: fmt.Sprintf("unexpected transaction type %T", txn),
		}
	}
	if underlying.db != db {
		return nil, &storage.Error{
			Code:    storage.InvalidTransactionErr,
			Message: "unknown transaction",
		}
	}
	if underlying.stale {
		return nil, &storage.Error{
			Code:    storage.InvalidTransactionErr,
			Message: "stale transaction",
		}
	}
	return underlying, nil
}

func mktree(path []string, value any) (map[string]any, error) {
	if len(path) == 0 {
		// For 0 length path the value is the full tree.
		obj, ok := value.(map[string]any)
		if !ok {
			return nil, errors.RootMustBeObjectErr
		}
		return obj, nil
	}

	dir := map[string]any{}
	for i := len(path) - 1; i > 0; i-- {
		dir[path[i]] = value
		value = dir
		dir = map[string]any{}
	}
	dir[path[0]] = value

	return dir, nil
}

func lookup(path storage.Path, data map[string]any) (any, bool) {
	if len(path) == 0 {
		return data, true
	}
	for i := range len(path) - 1 {
		value, ok := data[path[i]]
		if !ok {
			return nil, false
		}
		obj, ok := value.(map[string]any)
		if !ok {
			return nil, false
		}
		data = obj
	}
	value, ok := data[path[len(path)-1]]
	return value, ok
}