// 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 (
	"bytes"
	_ "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"
	caps "github.com/open-policy-agent/opa/v1/capabilities"
	"github.com/open-policy-agent/opa/v1/util"
)

// 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"`
}

// 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
var versionIndexBs []byte

// init only on demand, as JSON unmarshalling comes with some cost, and contributes
// noise to things like pprof stats
var minVersionIndexOnce = sync.OnceValue(func() VersionIndex {
	var vi VersionIndex
	if err := json.Unmarshal(versionIndexBs, &vi); err != nil {
		panic(err)
	}
	return vi
})

// 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"
const FeatureRefHeads = "rule_head_refs"
const FeatureRegoV1 = "rego_v1"
const FeatureRegoV1Import = "rego_v1_import"
const FeatureKeywordsInRefs = "keywords_in_refs"
const FeatureTemplateStrings = "template_strings"

// Features carries the default features supported by this version of OPA.
// Use RegisterFeatures to add to them.
var Features = []string{
	FeatureRegoV1,
	FeatureKeywordsInRefs,
	FeatureTemplateStrings,
}

// 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]) {
			continue
		}
		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 {
	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{}

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

	f.Builtins = make([]*Builtin, len(Builtins))
	copy(f.Builtins, Builtins)

	slices.SortFunc(f.Builtins, func(a, b *Builtin) int {
		return strings.Compare(a.Name, b.Name)
	})

	switch co.regoVersion {
	case RegoV0, RegoV0CompatV1:
		for kw := range allFutureKeywords {
			if _, internal := experimentalFutureKeywords[kw]; internal && !co.experimentalKeywords {
				continue
			}
			f.FutureKeywords = append(f.FutureKeywords, kw)
		}

		f.Features = []string{
			FeatureRefHeadStringPrefixes,
			FeatureRefHeads,
			FeatureRegoV1Import,
			FeatureRegoV1, // Included in v0 capabilities to allow v1 bundles in --v0-compatible mode
			FeatureKeywordsInRefs,
		}
	default:
		for kw := range futureKeywords {
			if _, internal := experimentalFutureKeywords[kw]; internal && !co.experimentalKeywords {
				continue
			}
			f.FutureKeywords = append(f.FutureKeywords, kw)
		}

		f.Features = make([]string, len(Features))
		copy(f.Features, Features)
	}

	slices.Sort(f.FutureKeywords)
	slices.Sort(f.Features)

	return f
}

// LoadCapabilitiesJSON loads a JSON serialized capabilities structure from the reader r.
func LoadCapabilitiesJSON(r io.Reader) (*Capabilities, error) {
	d := util.NewJSONDecoder(r)
	var c Capabilities
	return &c, d.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
	}

	for _, cv := range cvs {
		if cv == version {
			cont, err := caps.FS.ReadFile(cv + ".json")
			if err != nil {
				return nil, err
			}

			return LoadCapabilitiesJSON(bytes.NewReader(cont))
		}

	}
	return nil, fmt.Errorf("no capabilities version found %v", 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) {
	ents, err := caps.FS.ReadDir(".")
	if err != nil {
		return nil, err
	}

	return util.SortedStableFunc(util.Map(ents, 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 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)
}