// Copyright 2022 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 topdown

import (
	"encoding/json"
	"errors"

	"github.com/open-policy-agent/opa/internal/gojsonschema"
	"github.com/open-policy-agent/opa/v1/ast"
)

// astValueToJSONSchemaLoader converts a value to JSON Loader.
// Value can be ast.String or ast.Object.
func astValueToJSONSchemaLoader(value ast.Value) (gojsonschema.JSONLoader, error) {
	var loader gojsonschema.JSONLoader
	var err error

	// ast.Value type selector.
	switch x := value.(type) {
	case ast.String:
		// In case of string pass it as is as a raw JSON string.
		// Make pre-check that it's a valid JSON at all because gojsonschema won't do that.
		if !json.Valid([]byte(x)) {
			return nil, errors.New("invalid JSON string")
		}
		loader = gojsonschema.NewStringLoader(string(x))
	case ast.Object, *ast.Array:
		// In case of object serialize it to JSON representation.
		var data any
		data, err = ast.JSON(value)
		if err != nil {
			return nil, err
		}
		loader = gojsonschema.NewGoLoader(data)
	default:
		// Any other cases will produce an error.
		return nil, errors.New("wrong type, expected string or object")
	}

	return loader, nil
}

func newResultTerm(valid bool, data *ast.Term) *ast.Term {
	return ast.ArrayTerm(ast.InternedTerm(valid), data)
}

// newPatternValidatingSchemaLoader returns a SchemaLoader configured to
// compile and enforce the "pattern" keyword. This is the variant used by the
// json.verify_schema and json.match_schema built-ins, where runtime pattern
// validation is expected. It is intentionally not shared with the compile-time
// type-checking path, where pattern validation is disabled to tolerate
// schemas containing ECMA-262 regex features that Go's RE2 dialect can't
// compile.
//
// Remote reference fetching is restricted to the hosts in the caller's
// allow_net capability. Schemas reaching these built-ins come from the policy
// or, worse, from input, so an unrestricted loader would let a `$ref` drive
// outbound requests from wherever OPA happens to be deployed. The query's
// context comes along so that those fetches are abandoned when evaluation is
// cancelled.
func newPatternValidatingSchemaLoader(bctx BuiltinContext) *gojsonschema.SchemaLoader {
	sl := gojsonschema.NewSchemaLoader()
	sl.ValidatePatterns = true
	sl.Context = bctx.Context
	if bctx.Capabilities != nil {
		sl.AllowNet = bctx.Capabilities.AllowNet
	}
	return sl
}

// builtinJSONSchemaVerify accepts 1 argument which can be string or object and checks if it is valid JSON schema.
// Returns array [false, <string>] with error string at index 1, or [true, ""] with empty string at index 1 otherwise.
func builtinJSONSchemaVerify(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	// Take first argument and make JSON Loader from it.
	loader, err := astValueToJSONSchemaLoader(operands[0].Value)
	if err != nil {
		return iter(newResultTerm(false, ast.StringTerm("jsonschema: "+err.Error())))
	}

	// Check that schema is correct and parses without errors.
	if _, err = newPatternValidatingSchemaLoader(bctx).Compile(loader); err != nil {
		return iter(newResultTerm(false, ast.StringTerm("jsonschema: "+err.Error())))
	}

	return iter(newResultTerm(true, ast.InternedNullTerm))
}

// builtinJSONMatchSchema accepts 2 arguments both can be string or object and verifies if the document matches the JSON schema.
// Returns an array where first element is a boolean indicating a successful match, and the second is an array of errors that is empty on success and populated on failure.
// In case of internal error returns empty array.
//
// Cached schemas are keyed on the schema value alone. A compiled schema has
// already resolved its remote references, so a cache shared between callers
// with differing allow_net would leak across them. That needs a deliberately
// shared cache, an assumption http.send makes as well.
func builtinJSONMatchSchema(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	var schema *gojsonschema.Schema

	if bctx.InterQueryBuiltinValueCache != nil {
		if val, ok := bctx.InterQueryBuiltinValueCache.Get(operands[1].Value); ok {
			if s, isSchema := val.(*gojsonschema.Schema); isSchema {
				schema = s
			}
		}
	}

	// Take first argument and make JSON Loader from it.
	// This is a JSON document made from Rego JSON string or object.
	documentLoader, err := astValueToJSONSchemaLoader(operands[0].Value)
	if err != nil {
		return err
	}

	if schema == nil {
		// Take second argument and make JSON Loader from it.
		// This is a JSON schema made from Rego JSON string or object.
		schemaLoader, err := astValueToJSONSchemaLoader(operands[1].Value)
		if err != nil {
			return err
		}

		schema, err = newPatternValidatingSchemaLoader(bctx).Compile(schemaLoader)
		if err != nil {
			return err
		}

		if bctx.InterQueryBuiltinValueCache != nil {
			bctx.InterQueryBuiltinValueCache.Insert(operands[1].Value, schema)
		}
	}

	// Use schema to validate document.
	result, err := schema.Validate(documentLoader)
	if err != nil {
		return err
	}

	// In case of validation errors produce Rego array of objects to describe the errors.
	arr := ast.NewArrayWithCapacity(len(result.Errors()))
	for _, re := range result.Errors() {
		o := ast.NewObject(
			[...]*ast.Term{ast.StringTerm("error"), ast.StringTerm(re.String())},
			[...]*ast.Term{ast.StringTerm("type"), ast.StringTerm(re.Type())},
			[...]*ast.Term{ast.StringTerm("field"), ast.StringTerm(re.Field())},
			[...]*ast.Term{ast.StringTerm("desc"), ast.StringTerm(re.Description())},
		)
		arr = arr.Append(ast.NewTerm(o))
	}

	return iter(newResultTerm(result.Valid(), ast.NewTerm(arr)))
}

func init() {
	RegisterBuiltinFunc(ast.JSONSchemaVerify.Name, builtinJSONSchemaVerify)
	RegisterBuiltinFunc(ast.JSONMatchSchema.Name, builtinJSONMatchSchema)
}