// Copyright 2017 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 (
	"fmt"
	"strings"

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

// TypeEnv contains type info for static analysis such as type checking.
type TypeEnv struct {
	tree *typeTreeNode
	// vars is the tree the types inferred for variables are stored in, which is
	// tree except in the environments created for with modifiers.
	vars       *typeTreeNode
	next       *TypeEnv
	newChecker func() *typeChecker
}

// newTypeEnv returns an empty TypeEnv. The constructor is not exported because
// type environments should only be created by the type checker.
func newTypeEnv(f func() *typeChecker) *TypeEnv {
	tree := newTypeTree()
	return &TypeEnv{
		tree:       tree,
		vars:       tree,
		newChecker: f,
	}
}

// Get returns the type of x.
//
// Deprecated: Use GetByValue or GetByRef instead, as they are more efficient.
func (env *TypeEnv) Get(x any) types.Type {
	if term, ok := x.(*Term); ok {
		x = term.Value
	}

	if v, ok := x.(Value); ok {
		return env.GetByValue(v)
	}

	panic("unreachable")
}

// GetByValue returns the type of v.
func (env *TypeEnv) GetByValue(v Value) types.Type {
	switch x := v.(type) {

	// Scalars.
	case Null:
		return types.Nl
	case Boolean:
		return types.B
	case Number:
		return types.N
	case String, *TemplateString:
		return types.S

	// Composites.
	case *Array:
		static := make([]types.Type, x.Len())
		for i := range static {
			static[i] = env.GetByValue(x.Elem(i).Value)
		}

		return types.NewArray(static, nil)

	case *lazyObj:
		return env.GetByValue(x.force())
	case *object:
		static := []*types.StaticProperty{}
		var dynamic *types.DynamicProperty

		x.Foreach(func(k, v *Term) {
			if IsConstant(k.Value) {
				if kjson, err := JSON(k.Value); err == nil {
					static = append(static, types.NewStaticProperty(kjson, env.GetByValue(v.Value)))
					return
				}
			}
			// Can't handle it as a static property, fallback to dynamic
			dynamic = types.NewDynamicProperty(env.GetByValue(k.Value), env.GetByValue(v.Value))
		})

		return types.NewObject(static, dynamic)

	case *set:
		var tpe types.Type
		x.Foreach(func(elem *Term) {
			tpe = types.Or(tpe, env.GetByValue(elem.Value))
		})
		return types.NewSet(tpe)

	// Comprehensions.
	case *ArrayComprehension:
		cpy, errs := env.newChecker().CheckBody(env, x.Body)
		if len(errs) == 0 {
			return types.NewArray(nil, cpy.GetByValue(x.Term.Value))
		}
		return nil
	case *ObjectComprehension:
		cpy, errs := env.newChecker().CheckBody(env, x.Body)
		if len(errs) == 0 {
			return types.NewObject(nil, types.NewDynamicProperty(cpy.GetByValue(x.Key.Value), cpy.GetByValue(x.Value.Value)))
		}
		return nil
	case *SetComprehension:
		cpy, errs := env.newChecker().CheckBody(env, x.Body)
		if len(errs) == 0 {
			return types.NewSet(cpy.GetByValue(x.Term.Value))
		}
		return nil

	// Refs.
	case Ref:
		return env.GetByRef(x)

	// Vars.
	case Var:
		if node := env.tree.Child(v); node != nil {
			return node.Value()
		}
		if env.next != nil {
			return env.next.GetByValue(v)
		}
		return nil

	// Calls.
	case Call:
		return nil
	}

	return env.Get(v)
}

// GetByRef returns the type of the value referred to by ref.
func (env *TypeEnv) GetByRef(ref Ref) types.Type {
	node := env.tree.Child(ref[0].Value)
	if node == nil {
		return env.getRefFallback(ref)
	}

	return env.getRefRec(node, ref, ref[1:])
}

func (env *TypeEnv) getRefFallback(ref Ref) types.Type {
	if env.next != nil {
		return env.next.GetByRef(ref)
	}

	if RootDocumentNames.Contains(ref[0]) {
		// types.A is an empty types.Any
		// this is used to represent a potential non-local reference
		return types.A
	}

	return nil
}

func (env *TypeEnv) getRefRec(node *typeTreeNode, ref, tail Ref) types.Type {
	if len(tail) == 0 {
		return env.getRefRecExtent(node)
	}

	if node.Leaf() {
		if node.children.Len() > 0 {
			if child := node.Child(tail[0].Value); child != nil {
				return env.getRefRec(child, ref, tail[1:])
			}
		}
		return selectRef(node.Value(), tail)
	}

	if !IsConstant(tail[0].Value) {
		return selectRef(env.getRefRecExtent(node), tail)
	}

	child := node.Child(tail[0].Value)
	if child == nil {
		return env.getRefFallback(ref)
	}

	return env.getRefRec(child, ref, tail[1:])
}

func (env *TypeEnv) getRefRecExtent(node *typeTreeNode) types.Type {

	if node.Leaf() {
		return node.Value()
	}

	children := []*types.StaticProperty{}

	node.Children().Iter(func(key Value, child *typeTreeNode) bool {
		tpe := env.getRefRecExtent(child)

		// NOTE(sr): Converting to Golang-native types here is an extension of what we did
		// before -- only supporting strings. But since we cannot differentiate sets and arrays
		// that way, we could reconsider.
		switch key.(type) {
		case String, Number, Boolean: // skip anything else
			propKey, err := JSON(key)
			if err != nil {
				panic(fmt.Errorf("unreachable, ValueToInterface: %w", err))
			}
			children = append(children, types.NewStaticProperty(propKey, tpe))
		}
		return false
	})

	// TODO(tsandall): for now, these objects can have any dynamic properties
	// because we don't have schema for base docs. Once schemas are supported
	// we can improve this.
	return types.NewObject(children, types.NewDynamicProperty(types.S, types.A))
}

func (env *TypeEnv) wrap() *TypeEnv {
	cpy := *env
	cpy.next = env
	cpy.tree = newTypeTree()
	cpy.vars = cpy.tree
	return &cpy
}

// wrapWith returns a TypeEnv for checking a single expression carrying with
// modifiers: unlike wrap, the types it infers for variables are kept in the
// enclosing environment, as those variables outlive the expression.
func (env *TypeEnv) wrapWith() *TypeEnv {
	cpy := env.wrap()
	cpy.vars = env.vars
	return cpy
}

// typeTreeNode is used to store type information in a tree.
type typeTreeNode struct {
	key      Value
	value    types.Type
	children *util.HasherMap[Value, *typeTreeNode]
}

func newTypeTree() *typeTreeNode {
	return &typeTreeNode{
		key:      nil,
		value:    nil,
		children: util.NewHasherMap[Value, *typeTreeNode](ValueEqual),
	}
}

func (n *typeTreeNode) Child(key Value) *typeTreeNode {
	value, ok := n.children.Get(key)
	if !ok {
		return nil
	}
	return value
}

func (n *typeTreeNode) Children() *util.HasherMap[Value, *typeTreeNode] {
	return n.children
}

func (n *typeTreeNode) Get(path Ref) types.Type {
	curr := n
	for _, term := range path {
		child, ok := curr.children.Get(term.Value)
		if !ok {
			return nil
		}
		curr = child
	}
	return curr.Value()
}

func (n *typeTreeNode) Leaf() bool {
	return n.value != nil
}

func (n *typeTreeNode) PutOne(key Value, tpe types.Type) {
	c, ok := n.children.Get(key)

	var child *typeTreeNode
	if !ok {
		child = newTypeTree()
		child.key = key
		n.children.Put(key, child)
	} else {
		child = c
	}

	child.value = tpe
}

func (n *typeTreeNode) Put(path Ref, tpe types.Type) {
	curr := n
	for _, term := range path {
		child, ok := curr.children.Get(term.Value)
		if !ok {
			child = newTypeTree()
			child.key = term.Value
			curr.children.Put(child.key, child)
		}

		curr = child
	}
	curr.value = tpe
}

// Insert inserts tpe at path in the tree, but also merges the value into any types.Object present along that path.
// If a types.Object is inserted, any leafs already present further down the tree are merged into the inserted object.
// path must be ground.
func (n *typeTreeNode) Insert(path Ref, tpe types.Type, env *TypeEnv) {
	curr := n
	for i, term := range path {
		child, ok := curr.children.Get(term.Value)
		if !ok {
			child = newTypeTree()
			child.key = term.Value
			curr.children.Put(child.key, child)
		} else if child.value != nil && i+1 < len(path) {
			// If child has an object value, merge the new value into it.
			cv := child.value
			if r, ok := cv.(*types.Recursive); ok {
				cv = r.Unwrap()
			}
			if o, ok := cv.(*types.Object); ok {
				var err error
				child.value, err = insertIntoObject(o, path[i+1:], tpe, env)
				if err != nil {
					panic(fmt.Errorf("unreachable, insertIntoObject: %w", err))
				}
			}
		}

		curr = child
	}

	curr.value = mergeTypes(curr.value, tpe)

	_, isObj := tpe.(*types.Object)
	if !isObj {
		if r, ok := tpe.(*types.Recursive); ok {
			_, isObj = r.Unwrap().(*types.Object)
		}
	}
	if isObj && curr.children.Len() > 0 {
		// merge all leafs into the inserted object
		cv := curr.value
		if r, ok := cv.(*types.Recursive); ok {
			cv = r.Unwrap()
		}
		for p, t := range curr.Leafs() {
			var err error
			cv, err = insertIntoObject(cv.(*types.Object), *p, t, env)
			if err != nil {
				panic(fmt.Errorf("unreachable, insertIntoObject: %w", err))
			}
		}
		curr.value = cv
	}
}

// mergeTypes merges the types of 'a' and 'b'. If both are sets, their 'of' types are joined with an types.Or.
// If both are objects, the key types of their dynamic properties are joined with types.Or:s, and their value types
// are recursively merged (using mergeTypes).
// If 'a' and 'b' are both objects, and at least one of them have static properties, they are joined
// with an types.Or, instead of being merged.
// If 'a' is an Any containing an Object, and 'b' is an Object (or vice versa); AND both objects have no
// static properties, they are merged.
// If either object has neither static nor dynamic properties, it is the empty object type, and the other
// type is returned unchanged.
// If 'a' and 'b' are different types, they are joined with an types.Or.
func mergeTypes(a, b types.Type) types.Type {
	if a == nil {
		return b
	}

	if b == nil {
		return a
	}

	// Unwrap recursive types so they merge as their underlying type.
	if r, ok := a.(*types.Recursive); ok {
		a = r.Unwrap()
	}
	if r, ok := b.(*types.Recursive); ok {
		b = r.Unwrap()
	}

	switch a := a.(type) {
	case *types.Object:
		aDynProps := a.DynamicProperties()
		if bObj, ok := b.(*types.Object); ok && len(a.StaticProperties()) == 0 && len(bObj.StaticProperties()) == 0 {
			bDynProps := bObj.DynamicProperties()

			// An object type with neither static nor dynamic properties is the
			// empty object, which the other object type already covers.
			if aDynProps == nil {
				return bObj
			}
			if bDynProps == nil {
				return a
			}

			dynProps := types.NewDynamicProperty(
				types.Or(aDynProps.Key, bDynProps.Key),
				mergeTypes(aDynProps.Value, bDynProps.Value),
			)
			return types.NewObject(nil, dynProps)
		} else if bAny, ok := b.(types.Any); ok && len(a.StaticProperties()) == 0 && aDynProps != nil {
			// If a is an object type with no static components ...
			for _, t := range bAny {
				if tObj, ok := t.(*types.Object); ok && len(tObj.StaticProperties()) == 0 {
					// ... and b is a types.Any containing an object with no static components, we merge them.
					tDynProps := tObj.DynamicProperties()
					if tDynProps == nil {
						continue
					}
					tDynProps.Key = types.Or(tDynProps.Key, aDynProps.Key)
					tDynProps.Value = types.Or(tDynProps.Value, aDynProps.Value)
					return bAny
				}
			}
		}
	case *types.Set:
		if bSet, ok := b.(*types.Set); ok {
			return types.NewSet(types.Or(a.Of(), bSet.Of()))
		}
	case types.Any:
		if _, ok := b.(types.Any); !ok {
			return mergeTypes(b, a)
		}
	}

	return types.Or(a, b)
}

func (n *typeTreeNode) String() string {
	b := &strings.Builder{}

	key := "-"
	if k := n.key; k != nil {
		key = k.String()
	}

	b.WriteString(key)
	if v := n.value; v != nil {
		b.WriteString(": ")
		b.WriteString(v.String())
	}

	n.children.Iter(func(_ Value, child *typeTreeNode) bool {
		b.WriteString("\n\t+ ")
		b.WriteString(strings.ReplaceAll(child.String(), "\n", "\n\t"))

		return false
	})

	return b.String()
}

func insertIntoObject(o *types.Object, path Ref, tpe types.Type, env *TypeEnv) (*types.Object, error) {
	if len(path) == 0 {
		return o, nil
	}

	key := env.GetByValue(path[0].Value)

	if len(path) == 1 {
		var dynamicProps *types.DynamicProperty
		if dp := o.DynamicProperties(); dp != nil {
			dynamicProps = types.NewDynamicProperty(types.Or(o.DynamicProperties().Key, key), types.Or(o.DynamicProperties().Value, tpe))
		} else {
			dynamicProps = types.NewDynamicProperty(key, tpe)
		}
		return types.NewObject(o.StaticProperties(), dynamicProps), nil
	}

	child, err := insertIntoObject(types.NewObject(nil, nil), path[1:], tpe, env)
	if err != nil {
		return nil, err
	}

	var dynamicProps *types.DynamicProperty
	if dp := o.DynamicProperties(); dp != nil {
		dynamicProps = types.NewDynamicProperty(types.Or(o.DynamicProperties().Key, key), types.Or(o.DynamicProperties().Value, child))
	} else {
		dynamicProps = types.NewDynamicProperty(key, child)
	}
	return types.NewObject(o.StaticProperties(), dynamicProps), nil
}

func (n *typeTreeNode) Leafs() map[*Ref]types.Type {
	leafs := map[*Ref]types.Type{}
	n.children.Iter(func(_ Value, v *typeTreeNode) bool {
		collectLeafs(v, nil, leafs)
		return false
	})
	return leafs
}

func collectLeafs(n *typeTreeNode, path Ref, leafs map[*Ref]types.Type) {
	nPath := append(path, NewTerm(n.key))
	if n.Leaf() {
		npc := nPath // copy of else nPath escapes to heap even if !n.Leaf()
		leafs[&npc] = n.Value()
		return
	}
	n.children.Iter(func(_ Value, v *typeTreeNode) bool {
		collectLeafs(v, nPath, leafs)
		return false
	})
}

func (n *typeTreeNode) Value() types.Type {
	return n.value
}

// selectConstant returns the attribute of the type referred to by the term. If
// the attribute type cannot be determined, nil is returned.
func selectConstant(tpe types.Type, term *Term) types.Type {
	x, err := JSON(term.Value)
	if err == nil {
		return types.Select(tpe, x)
	}
	return nil
}

// selectRef returns the type of the nested attribute referred to by ref. If
// the attribute type cannot be determined, nil is returned. If the ref
// contains vars or refs, then the returned type will be a union of the
// possible types.
func selectRef(tpe types.Type, ref Ref) types.Type {
	if tpe == nil || len(ref) == 0 {
		return tpe
	}

	head, tail := ref[0], ref[1:]

	switch head.Value.(type) {
	case Var, Ref, *Array, Object, Set:
		return selectRef(types.Values(tpe), tail)
	default:
		return selectRef(selectConstant(tpe, head), tail)
	}
}