// Copyright 2020 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 ast

import (
	_ "embed"
	"encoding/json"
	"fmt"
	"io"
	"io/fs"
	"os"
	"slices"
	"strings"
	"sync"

	"github.com/open-policy-agent/opa/internal/semver"
	"github.com/open-policy-agent/opa/internal/wasm/sdk/opa/capabilities"
	"github.com/open-policy-agent/opa/v1/ast/internal/tokens"
	caps "github.com/open-policy-agent/opa/v1/capabilities"
	"github.com/open-policy-agent/opa/v1/util"
)

// In the compiler, we used this to check that we're OK working with ref heads.
// If this isn't present, we'll fail. This is to ensure that older versions of
// OPA can work with policies that we're compiling -- if they don't know ref
// heads, they wouldn't be able to parse them.
const (
	FeatureRefHeadStringPrefixes = "rule_head_ref_string_prefixes"
	FeatureRefHeads              = "rule_head_refs"
	FeatureRegoV1                = "rego_v1"
	FeatureRegoV1Import          = "rego_v1_import"
	FeatureKeywordsInRefs        = "keywords_in_refs"
	FeatureTemplateStrings       = "template_strings"
)

var (
	// Features carries the default features supported by this version of OPA.
	// Use RegisterFeatures to add to them.
	Features = []string{
		FeatureRegoV1,
		FeatureKeywordsInRefs,
		FeatureTemplateStrings,
	}
	v0v1compatFeatures = []string{
		FeatureRefHeadStringPrefixes,
		FeatureRefHeads,
		FeatureRegoV1Import,
		FeatureRegoV1, // Included in v0 capabilities to allow v1 bundles in --v0-compatible mode
		FeatureKeywordsInRefs,
	}
	// NOTE(tsandall): this file is generated by internal/cmd/genversionindex/main.go
	// and run as part of go:generate. We generate the version index as part of the
	// build process because it's relatively expensive to build (it takes ~500ms on
	// my machine) and never changes.
	//
	//go:embed version_index.json
	versionIndexBs []byte
	// init only on demand, as JSON unmarshalling comes with some cost, and contributes
	// noise to things like pprof stats
	minVersionIndexOnce = sync.OnceValue(minVersionIndex)
)

// VersonIndex contains an index from built-in function name, language feature,
// and future rego keyword to version number. During the build, this is used to
// create an index of the minimum version required for the built-in/feature/kw.
type VersionIndex struct {
	Builtins map[string]semver.Version `json:"builtins"`
	Features map[string]semver.Version `json:"features"`
	Keywords map[string]semver.Version `json:"keywords"`
}

// RegisterFeatures lets applications wrapping OPA register features, to be
// included in `ast.CapabilitiesForThisVersion()`.
func RegisterFeatures(fs ...string) {
	for i := range fs {
		if !slices.Contains(Features, fs[i]) {
			Features = append(Features, fs[i])
		}
	}
}

// Capabilities defines a structure containing data that describes the capabilities
// or features supported by a particular version of OPA.
type Capabilities struct {
	Builtins        []*Builtin       `json:"builtins,omitempty"`
	FutureKeywords  []string         `json:"future_keywords,omitempty"`
	WasmABIVersions []WasmABIVersion `json:"wasm_abi_versions,omitempty"`

	// Features is a bit of a mixed bag for checking that an older version of OPA
	// is able to do what needs to be done.
	// TODO(sr): find better words ^^
	Features []string `json:"features,omitempty"`

	// allow_net is an array of hostnames or IP addresses, that an OPA instance is
	// allowed to connect to.
	// If omitted, ANY host can be connected to. If empty, NO host can be connected to.
	// As of now, this only controls fetching remote refs for using JSON Schemas in
	// the type checker.
	// TODO(sr): support ports to further restrict connection peers
	AllowNet []string `json:"allow_net,omitempty"`
}

// WasmABIVersion captures the Wasm ABI version. Its `Minor` version is indicating
// backwards-compatible changes.
type WasmABIVersion struct {
	Version int `json:"version"`
	Minor   int `json:"minor_version"`
}

type CapabilitiesOptions struct {
	regoVersion          RegoVersion
	experimentalKeywords bool
}

func newCapabilitiesOptions(opts []CapabilitiesOption) CapabilitiesOptions {
	if len(opts) == 0 {
		return CapabilitiesOptions{}
	}
	co := CapabilitiesOptions{}
	for _, opt := range opts {
		opt(&co)
	}
	return co
}

type CapabilitiesOption func(*CapabilitiesOptions)

func CapabilitiesRegoVersion(regoVersion RegoVersion) CapabilitiesOption {
	return func(o *CapabilitiesOptions) {
		o.regoVersion = regoVersion
	}
}

// CapabilitiesExperimentalKeywords controls whether experimental future
// keywords are included in the returned capabilities. When enabled, the
// future_keywords list will include keywords that are recognized by the
// parser but are not yet ready for general use.
//
// Experimental keywords are subject to change or removal without notice and
// are not covered by OPA's compatibility guarantees. Enable this option only
// when you intentionally want to opt in to in-progress language features
// (for example, when writing tests for those features).
func CapabilitiesExperimentalKeywords(yes bool) CapabilitiesOption {
	return func(o *CapabilitiesOptions) {
		o.experimentalKeywords = yes
	}
}

// CapabilitiesForThisVersion returns the capabilities of this version of OPA.
func CapabilitiesForThisVersion(opts ...CapabilitiesOption) *Capabilities {
	co := newCapabilitiesOptions(opts)
	f := &Capabilities{Builtins: util.SortedFunc(slices.Clone(Builtins), cmpBuiltinName)}

	for _, vers := range capabilities.ABIVersions() {
		f.WasmABIVersions = append(f.WasmABIVersions, WasmABIVersion{Version: vers[0], Minor: vers[1]})
	}

	switch co.regoVersion {
	case RegoV0, RegoV0CompatV1:
		f.FutureKeywords = co.filterFutureKeywords(allFutureKeywords)
		f.Features = util.Sorted(slices.Clone(v0v1compatFeatures))
	default:
		f.FutureKeywords = co.filterFutureKeywords(futureKeywords)
		f.Features = util.Sorted(slices.Clone(Features))
	}

	return f
}

func (co *CapabilitiesOptions) filterFutureKeywords(src map[string]tokens.Token) []string {
	if co.experimentalKeywords {
		return util.KeysSorted(src)
	}
	return slices.DeleteFunc(util.KeysSorted(src), isExperimental)
}

// LoadCapabilitiesJSON loads a JSON serialized capabilities structure from the reader r.
func LoadCapabilitiesJSON(r io.Reader) (*Capabilities, error) {
	var c Capabilities
	return &c, util.NewJSONDecoder(r).Decode(&c)
}

// LoadCapabilitiesVersion loads a JSON serialized capabilities structure from the specific version.
func LoadCapabilitiesVersion(version string) (*Capabilities, error) {
	cvs, err := LoadCapabilitiesVersions()
	if err != nil {
		return nil, err
	}

	if slices.Contains(cvs, version) {
		fd, err := caps.FS.Open(version + ".json")
		if err != nil {
			return nil, err
		}
		defer fd.Close()
		return LoadCapabilitiesJSON(fd)
	}

	return nil, fmt.Errorf("no capabilities version found %s", version)
}

// LoadCapabilitiesFile loads a JSON serialized capabilities structure from a file.
func LoadCapabilitiesFile(file string) (*Capabilities, error) {
	fd, err := os.Open(file)
	if err != nil {
		return nil, err
	}
	defer fd.Close()
	return LoadCapabilitiesJSON(fd)
}

// LoadCapabilitiesVersions loads all capabilities versions
func LoadCapabilitiesVersions() ([]string, error) {
	entries, err := caps.FS.ReadDir(".")
	if err != nil {
		return nil, err
	}
	return util.SortedStableFunc(util.Map(entries, removeJsonSuffix), semver.Compare), nil
}

// MinimumCompatibleVersion returns the minimum compatible OPA version based on
// the built-ins, features, and keywords in c.
func (c *Capabilities) MinimumCompatibleVersion() (string, bool) {
	// this is the oldest OPA release that includes capabilities
	maxVersion := semver.MustParse("0.17.0")
	minVersionIndex := minVersionIndexOnce()

	for _, bi := range c.Builtins {
		v, ok := minVersionIndex.Builtins[bi.Name]
		if !ok {
			return "", false
		}
		if v.Compare(maxVersion) > 0 {
			maxVersion = v
		}
	}

	for _, kw := range c.FutureKeywords {
		v, ok := minVersionIndex.Keywords[kw]
		if !ok {
			return "", false
		}
		if v.Compare(maxVersion) > 0 {
			maxVersion = v
		}
	}

	for _, feat := range c.Features {
		v, ok := minVersionIndex.Features[feat]
		if !ok {
			return "", false
		}
		if v.Compare(maxVersion) > 0 {
			maxVersion = v
		}
	}

	return maxVersion.String(), true
}

func (c *Capabilities) ContainsFeature(feature string) bool {
	return slices.Contains(c.Features, feature)
}

func (c *Capabilities) ContainsBuiltin(name string) bool {
	return slices.ContainsFunc(c.Builtins, func(builtin *Builtin) bool {
		return builtin.Name == name
	})
}

func (c *Capabilities) ContainsFutureKeyword(kw string) bool {
	return slices.Contains(c.FutureKeywords, kw)
}

// addBuiltinSorted inserts a built-in into c in sorted order. An existing built-in with the same name
// will be overwritten.
func (c *Capabilities) addBuiltinSorted(bi *Builtin) {
	i, found := slices.BinarySearchFunc(c.Builtins, bi, cmpBuiltinName)
	if !found {
		c.Builtins = append(c.Builtins, nil)
		copy(c.Builtins[i+1:], c.Builtins[i:])
	}
	c.Builtins[i] = bi
}

func minVersionIndex() (vi VersionIndex) {
	if err := json.Unmarshal(versionIndexBs, &vi); err != nil {
		panic(err)
	}
	return vi
}

func cmpBuiltinName(a, b *Builtin) int {
	return strings.Compare(a.Name, b.Name)
}

func removeJsonSuffix(ent fs.DirEntry) string {
	return strings.Replace(ent.Name(), ".json", "", 1)
}

func isExperimental(kw string) bool {
	_, experimental := experimentalFutureKeywords[kw]
	return experimental
}