// 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 types contains request/response types and codes for the server.
package types

import (
	"bytes"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"

	"github.com/open-policy-agent/opa/v1/ast"
	"github.com/open-policy-agent/opa/v1/topdown"
	"github.com/open-policy-agent/opa/v1/util"
)

// jsonFields holds the set of JSON field names declared on registered types.
// It is populated at package init time (see RegisterJSONFields) and is
// read-only thereafter, so concurrent reads need no synchronization.
var jsonFields = map[reflect.Type]map[string]bool{}

// RegisterJSONFields records the JSON field names declared on T so that
// UnmarshalExtras and MarshalExtras can distinguish known fields from extras.
// Call from an init() function in the package that defines T.
func RegisterJSONFields[T any]() {
	t := reflect.TypeFor[T]()
	if t.Kind() != reflect.Struct {
		panic(fmt.Sprintf("types: RegisterJSONFields[%s]: not a struct", t))
	}
	f := make(map[string]bool, t.NumField())
	for i := range t.NumField() {
		sf := t.Field(i)
		tag := sf.Tag.Get("json")
		if tag == "-" {
			continue
		}
		name, _, _ := strings.Cut(tag, ",")
		if name == "" {
			name = sf.Name
		}
		f[name] = true
	}
	jsonFields[t] = f
}

func knownJSONFields[T any]() map[string]bool {
	t := reflect.TypeFor[T]()
	f, ok := jsonFields[t]
	if !ok {
		panic(fmt.Sprintf("types: %s not registered with RegisterJSONFields", t))
	}
	return f
}

// UnmarshalExtras decodes data into into (typically a *alias of T to avoid
// recursing through T's UnmarshalJSON), then returns any top-level JSON keys
// not declared on T. T must have been registered with RegisterJSONFields.
func UnmarshalExtras[T any](data []byte, into any) (map[string]any, error) {
	if err := util.UnmarshalJSON(data, into); err != nil {
		return nil, err
	}
	var raw map[string]json.RawMessage
	if err := util.UnmarshalJSON(data, &raw); err != nil {
		return nil, err
	}
	known := knownJSONFields[T]()
	var extra map[string]any
	for k, val := range raw {
		if known[k] {
			continue
		}
		var v any
		if err := util.UnmarshalJSON(val, &v); err != nil {
			return nil, err
		}
		if extra == nil {
			extra = make(map[string]any)
		}
		extra[k] = v
	}
	return extra, nil
}

// MarshalExtras marshals v (typically an alias of T to avoid recursing through
// T's MarshalJSON) and merges extra into the resulting JSON object, dropping
// any extra keys that collide with a JSON field declared on T. T must have
// been registered with RegisterJSONFields.
func MarshalExtras[T any](v any, extra map[string]any) ([]byte, error) {
	data, err := json.Marshal(v)
	if err != nil {
		return nil, err
	}
	if len(extra) == 0 {
		return data, nil
	}
	known := knownJSONFields[T]()
	var base map[string]any
	if err := json.Unmarshal(data, &base); err != nil {
		return nil, err
	}
	for k, val := range extra {
		if known[k] {
			continue
		}
		base[k] = val
	}
	return json.Marshal(base)
}

func init() {
	RegisterJSONFields[DataRequestV1]()
	RegisterJSONFields[DataResponseV1]()
}

// Error codes returned by OPA's REST API.
const (
	CodeInternal          = "internal_error"
	CodeEvaluation        = "evaluation_error"
	CodeUnauthorized      = "unauthorized"
	CodeInvalidParameter  = "invalid_parameter"
	CodeInvalidOperation  = "invalid_operation"
	CodeResourceNotFound  = "resource_not_found"
	CodeResourceConflict  = "resource_conflict"
	CodeUndefinedDocument = "undefined_document"
)

// ErrorV1 models an error response sent to the client.
type ErrorV1 struct {
	Code    string  `json:"code"`
	Message string  `json:"message"`
	Errors  []error `json:"errors,omitempty"`
}

// NewErrorV1 returns a new ErrorV1 object.
func NewErrorV1(code, f string, a ...any) *ErrorV1 {
	return &ErrorV1{
		Code:    code,
		Message: fmt.Sprintf(f, a...),
	}
}

// This shall only used for debugging purpose.
func (e *ErrorV1) Error() string {
	return fmt.Sprintf("%s: %s", e.Code, e.Message)
}

// WithError updates e to include a detailed error.
func (e *ErrorV1) WithError(err error) *ErrorV1 {
	e.Errors = append(e.Errors, err)
	return e
}

// WithASTErrors updates e to include detailed AST errors.
func (e *ErrorV1) WithASTErrors(errors []*ast.Error) *ErrorV1 {
	e.Errors = make([]error, len(errors))
	for i := range e.Errors {
		e.Errors[i] = errors[i]
	}
	return e
}

// Bytes marshals e with indentation for readability.
func (e *ErrorV1) Bytes() []byte {
	bs, _ := json.MarshalIndent(e, "", "  ")
	return bs
}

// Messages included in error responses.
const (
	MsgCompileModuleError         = "error(s) occurred while compiling module(s)"
	MsgParseQueryError            = "error(s) occurred while parsing query"
	MsgCompileQueryError          = "error(s) occurred while compiling query"
	MsgEvaluationError            = "error(s) occurred while evaluating query"
	MsgUnauthorizedUndefinedError = "authorization policy missing or undefined"
	MsgUnauthorizedError          = "request rejected by administrative policy"
	MsgUndefinedError             = "document missing or undefined"
	MsgMissingError               = "document missing"
	MsgFoundUndefinedError        = "document undefined"
	MsgPluginConfigError          = "error(s) occurred while configuring plugin(s)"
	MsgDecodingLimitError         = "request body too large"
	MsgDecodingGzipLimitError     = "compressed request body too large"
)

// PatchV1 models a single patch operation against a document.
type PatchV1 struct {
	Op    string `json:"op"`
	Path  string `json:"path"`
	Value any    `json:"value"`
}

// PolicyListResponseV1 models the response message for the Policy API list operation.
type PolicyListResponseV1 struct {
	Result []PolicyV1 `json:"result"`
}

// PolicyGetResponseV1 models the response message for the Policy API get operation.
type PolicyGetResponseV1 struct {
	Result PolicyV1 `json:"result"`
}

// PolicyPutResponseV1 models the response message for the Policy API put operation.
type PolicyPutResponseV1 struct {
	Metrics MetricsV1 `json:"metrics,omitempty"`
}

// PolicyDeleteResponseV1 models the response message for the Policy API delete operation.
type PolicyDeleteResponseV1 struct {
	Metrics MetricsV1 `json:"metrics,omitempty"`
}

// PolicyV1 models a policy module in OPA.
type PolicyV1 struct {
	ID  string      `json:"id"`
	Raw string      `json:"raw"`
	AST *ast.Module `json:"ast"`
}

// Equal returns true if p is equal to other.
func (p PolicyV1) Equal(other PolicyV1) bool {
	return p.ID == other.ID && p.Raw == other.Raw && p.AST.Equal(other.AST)
}

// ProvenanceV1 models a collection of build/version information.
type ProvenanceV1 struct {
	Version   string                        `json:"version"`
	Vcs       string                        `json:"build_commit"`
	Timestamp string                        `json:"build_timestamp"`
	Hostname  string                        `json:"build_hostname"`
	Revision  string                        `json:"revision,omitempty"` // Deprecated: Prefer `Bundles`
	Bundles   map[string]ProvenanceBundleV1 `json:"bundles,omitempty"`
}

// ProvenanceBundleV1 models a bundle at some point in time
type ProvenanceBundleV1 struct {
	Revision string `json:"revision"`
}

// DataRequestV1 models the request message for Data API POST operations.
type DataRequestV1 struct {
	Input *any `json:"input"`

	// Metadata holds any additional top-level fields not defined in this struct.
	// These fields are preserved during JSON unmarshaling, allowing wrapping
	// projects to pass through custom key/value pairs.
	//
	// Future OPA versions may introduce new top-level keys. To avoid conflicts,
	// use a unique namespaced key, e.g. "com.example.opa/metadata": { ... }.
	Metadata map[string]any `json:"-"`
}

func (r *DataRequestV1) UnmarshalJSON(data []byte) error {
	type alias DataRequestV1
	extra, err := UnmarshalExtras[DataRequestV1](data, (*alias)(r))
	r.Metadata = extra
	return err
}

func (r DataRequestV1) MarshalJSON() ([]byte, error) {
	type alias DataRequestV1
	return MarshalExtras[DataRequestV1](alias(r), r.Metadata)
}

// DataResponseV1 models the response message for Data API read operations.
type DataResponseV1 struct {
	DecisionID  string        `json:"decision_id,omitempty"`
	Provenance  *ProvenanceV1 `json:"provenance,omitempty"`
	Explanation TraceV1       `json:"explanation,omitempty"`
	Metrics     MetricsV1     `json:"metrics,omitempty"`
	Result      *any          `json:"result,omitempty"`
	Warning     *Warning      `json:"warning,omitempty"`

	// Metadata holds any additional top-level fields not defined in this struct.
	// These fields are preserved during JSON marshaling/unmarshaling, allowing
	// wrapping projects to pass through custom key/value pairs.
	//
	// Future OPA versions may introduce new top-level keys. To avoid conflicts,
	// use a unique namespaced key, e.g. "com.example.opa/metadata": { ... }.
	Metadata map[string]any `json:"-"`
}

func (r *DataResponseV1) UnmarshalJSON(data []byte) error {
	type alias DataResponseV1
	extra, err := UnmarshalExtras[DataResponseV1](data, (*alias)(r))
	r.Metadata = extra
	return err
}

func (r DataResponseV1) MarshalJSON() ([]byte, error) {
	type alias DataResponseV1
	return MarshalExtras[DataResponseV1](alias(r), r.Metadata)
}

// Warning models DataResponse warnings
type Warning struct {
	Code    string `json:"code,omitempty"`
	Message string `json:"message,omitempty"`
}

// Warning Codes
const CodeAPIUsageWarn = "api_usage_warning"

// Warning Messages
const MsgInputKeyMissing = "'input' key missing from the request"

// NewWarning returns a new Warning object
func NewWarning(code, message string) *Warning {
	return &Warning{Code: code, Message: message}
}

// MetricsV1 models a collection of performance metrics.
type MetricsV1 map[string]any

// QueryResponseV1 models the response message for Query API operations.
type QueryResponseV1 struct {
	Explanation TraceV1               `json:"explanation,omitempty"`
	Metrics     MetricsV1             `json:"metrics,omitempty"`
	Result      AdhocQueryResultSetV1 `json:"result,omitempty"`
}

// AdhocQueryResultSetV1 models the result of a Query API query.
type AdhocQueryResultSetV1 []map[string]any

// ExplainModeV1 defines supported values for the "explain" query parameter.
type ExplainModeV1 string

// Explanation mode enumeration.
const (
	ExplainOffV1   ExplainModeV1 = "off"
	ExplainFullV1  ExplainModeV1 = "full"
	ExplainNotesV1 ExplainModeV1 = "notes"
	ExplainFailsV1 ExplainModeV1 = "fails"
	ExplainDebugV1 ExplainModeV1 = "debug"
)

// TraceV1 models the trace result returned for queries that include the
// "explain" parameter.
type TraceV1 json.RawMessage

// MarshalJSON unmarshals the TraceV1 to a JSON representation.
func (t TraceV1) MarshalJSON() ([]byte, error) {
	return t, nil
}

// UnmarshalJSON unmarshals the TraceV1 from a JSON representation.
func (t *TraceV1) UnmarshalJSON(b []byte) error {
	*t = TraceV1(b)
	return nil
}

// TraceV1Raw models the trace result returned for queries that include the
// "explain" parameter. The trace is modelled as series of trace events that
// identify the expression, local term bindings, query hierarchy, etc.
type TraceV1Raw []TraceEventV1

// UnmarshalJSON unmarshals the TraceV1Raw from a JSON representation.
func (t *TraceV1Raw) UnmarshalJSON(b []byte) error {
	var trace []TraceEventV1
	if err := json.Unmarshal(b, &trace); err != nil {
		return err
	}
	*t = TraceV1Raw(trace)
	return nil
}

// TraceV1Pretty models the trace result returned for queries that include the "explain"
// parameter. The trace is modelled as a human readable array of strings representing the
// evaluation of the query.
type TraceV1Pretty []string

// UnmarshalJSON unmarshals the TraceV1Pretty from a JSON representation.
func (t *TraceV1Pretty) UnmarshalJSON(b []byte) error {
	var s []string
	if err := json.Unmarshal(b, &s); err != nil {
		return err
	}
	*t = TraceV1Pretty(s)
	return nil
}

// NewTraceV1 returns a new TraceV1 object.
func NewTraceV1(trace []*topdown.Event, pretty bool) (result TraceV1, err error) {
	if pretty {
		return newPrettyTraceV1(trace)
	}
	return newRawTraceV1(trace)
}

func newRawTraceV1(trace []*topdown.Event) (TraceV1, error) {
	result := TraceV1Raw(make([]TraceEventV1, len(trace)))
	for i := range trace {
		result[i] = TraceEventV1{
			Op:       strings.ToLower(string(trace[i].Op)),
			QueryID:  trace[i].QueryID,
			ParentID: trace[i].ParentID,
			Locals:   NewBindingsV1(trace[i].Locals),
			Message:  trace[i].Message,
		}
		if trace[i].Node != nil {
			result[i].Type = ast.TypeName(trace[i].Node)
			result[i].Node = trace[i].Node
		}
	}

	b, err := json.Marshal(result)
	if err != nil {
		return nil, err
	}
	return TraceV1(json.RawMessage(b)), nil
}

func newPrettyTraceV1(trace []*topdown.Event) (TraceV1, error) {
	var buf bytes.Buffer
	topdown.PrettyTraceWithLocation(&buf, trace)

	str := strings.Trim(buf.String(), "\n")
	b, err := json.Marshal(strings.Split(str, "\n"))
	if err != nil {
		return nil, err
	}
	return TraceV1(json.RawMessage(b)), nil
}

// TraceEventV1 represents a step in the query evaluation process.
type TraceEventV1 struct {
	Op       string     `json:"op"`
	QueryID  uint64     `json:"query_id"`
	ParentID uint64     `json:"parent_id"`
	Type     string     `json:"type"`
	Node     any        `json:"node"`
	Locals   BindingsV1 `json:"locals"`
	Message  string     `json:"message,omitempty"`
}

// UnmarshalJSON deserializes a TraceEventV1 object. The Node field is
// deserialized based on the type hint from the type property in the JSON
// object.
func (te *TraceEventV1) UnmarshalJSON(bs []byte) error {

	keys := map[string]json.RawMessage{}

	if err := util.UnmarshalJSON(bs, &keys); err != nil {
		return err
	}

	if err := util.UnmarshalJSON(keys["type"], &te.Type); err != nil {
		return err
	}

	if err := util.UnmarshalJSON(keys["op"], &te.Op); err != nil {
		return err
	}

	if err := util.UnmarshalJSON(keys["query_id"], &te.QueryID); err != nil {
		return err
	}

	if err := util.UnmarshalJSON(keys["parent_id"], &te.ParentID); err != nil {
		return err
	}

	switch te.Type {
	case "body":
		var body ast.Body
		if err := util.UnmarshalJSON(keys["node"], &body); err != nil {
			return err
		}
		te.Node = body
	case "expr":
		var expr ast.Expr
		if err := util.UnmarshalJSON(keys["node"], &expr); err != nil {
			return err
		}
		te.Node = &expr
	case "rule":
		var rule ast.Rule
		if err := util.UnmarshalJSON(keys["node"], &rule); err != nil {
			return err
		}
		te.Node = &rule
	}

	return util.UnmarshalJSON(keys["locals"], &te.Locals)
}

// BindingsV1 represents a set of term bindings.
type BindingsV1 []*BindingV1

// BindingV1 represents a single term binding.
type BindingV1 struct {
	Key   *ast.Term `json:"key"`
	Value *ast.Term `json:"value"`
}

// NewBindingsV1 returns a new BindingsV1 object.
func NewBindingsV1(locals *ast.ValueMap) (result []*BindingV1) {
	result = make([]*BindingV1, 0, locals.Len())
	locals.Iter(func(key, value ast.Value) bool {
		result = append(result, &BindingV1{
			Key:   &ast.Term{Value: key},
			Value: &ast.Term{Value: value},
		})
		return false
	})
	return result
}

// CompileRequestV1 models the request message for Compile API operations.
type CompileRequestV1 struct {
	Input    *any      `json:"input"`
	Query    string    `json:"query"`
	Unknowns *[]string `json:"unknowns"`
	Options  struct {
		DisableInlining          []string `json:"disableInlining,omitempty"`
		NondeterministicBuiltins bool     `json:"nondeterministicBuiltins"`
	} `json:"options"`
}

// CompileResponseV1 models the response message for Compile API operations.
type CompileResponseV1 struct {
	Result      *any      `json:"result,omitempty"`
	Explanation TraceV1   `json:"explanation,omitempty"`
	Metrics     MetricsV1 `json:"metrics,omitempty"`
}

// PartialEvaluationResultV1 represents the output of partial evaluation and is
// included in Compile API responses.
type PartialEvaluationResultV1 struct {
	Queries []ast.Body    `json:"queries,omitempty"`
	Support []*ast.Module `json:"support,omitempty"`
}

// QueryRequestV1 models the request message for Query API operations.
type QueryRequestV1 struct {
	Query string `json:"query"`
	Input *any   `json:"input"`
}

// ConfigResponseV1 models the response message for Config API operations.
type ConfigResponseV1 struct {
	Result *any `json:"result,omitempty"`
}

// StatusResponseV1 models the response message for Status API (pull) operations.
type StatusResponseV1 struct {
	Result *any `json:"result,omitempty"`
}

// HealthResponseV1 models the response message for Health API operations.
type HealthResponseV1 struct {
	Error string `json:"error,omitempty"`
}

const (
	// ParamQueryV1 defines the name of the HTTP URL parameter that specifies
	// values for the request query.
	ParamQueryV1 = "q"

	// ParamInputV1 defines the name of the HTTP URL parameter that specifies
	// values for the "input" document.
	ParamInputV1 = "input"

	// ParamPrettyV1 defines the name of the HTTP URL parameter that indicates
	// the client wants to receive a pretty-printed version of the response.
	ParamPrettyV1 = "pretty"

	// ParamExplainV1 defines the name of the HTTP URL parameter that indicates the
	// client wants to receive explanations in addition to the result.
	ParamExplainV1 = "explain"

	// ParamMetricsV1 defines the name of the HTTP URL parameter that indicates
	// the client wants to receive performance metrics in addition to the
	// result.
	ParamMetricsV1 = "metrics"

	// ParamInstrumentV1 defines the name of the HTTP URL parameter that
	// indicates the client wants to receive instrumentation data for
	// diagnosing performance issues.
	ParamInstrumentV1 = "instrument"

	// ParamProvenanceV1 defines the name of the HTTP URL parameter that indicates
	// the client wants build and version information in addition to the result.
	ParamProvenanceV1 = "provenance"

	// ParamBundleActivationV1 defines the name of the HTTP URL parameter that
	// indicates the client wants to include bundle activation in the results
	// of the health API.
	//
	// Deprecated: Use ParamBundlesActivationV1 instead.
	ParamBundleActivationV1 = "bundle"

	// ParamBundlesActivationV1 defines the name of the HTTP URL parameter that
	// indicates the client wants to include bundle activation in the results
	// of the health API.
	ParamBundlesActivationV1 = "bundles"

	// ParamPluginsV1 defines the name of the HTTP URL parameter that
	// indicates the client wants to include bundle status in the results
	// of the health API.
	ParamPluginsV1 = "plugins"

	// ParamExcludePluginV1 defines the name of the HTTP URL parameter that
	// indicates the client wants to exclude plugin status in the results
	// of the health API for the specified plugin(s)
	ParamExcludePluginV1 = "exclude-plugin"

	// ParamStrictBuiltinErrors names the HTTP URL parameter that indicates the client
	// wants built-in function errors to be treated as fatal.
	ParamStrictBuiltinErrors = "strict-builtin-errors"
)

// BadRequestErr represents an error condition raised if the caller passes
// invalid parameters.
type BadRequestErr string

// BadPatchOperationErr returns BadRequestErr indicating the patch operation was
// invalid.
func BadPatchOperationErr(op string) error {
	return BadRequestErr(fmt.Sprintf("bad patch operation: %v", op))
}

// BadPatchPathErr returns BadRequestErr indicating the patch path was invalid.
func BadPatchPathErr(path string) error {
	return BadRequestErr(fmt.Sprintf("bad patch path: %v", path))
}

func (err BadRequestErr) Error() string {
	return string(err)
}

// IsBadRequest returns true if err is a BadRequestErr.
func IsBadRequest(err error) bool {
	_, ok := err.(BadRequestErr)
	return ok
}
