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

import (
	"fmt"
	"regexp"
	"regexp/syntax"
	"sync"

	gintersect "github.com/yashtewari/glob-intersection"

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

const (
	regexCacheMaxSize             = 100
	regexInterQueryValueCacheHits = "rego_builtin_regex_interquery_value_cache_hits"
)

var (
	regexpCacheLock = sync.RWMutex{}
	regexpCache     = make(map[string]*regexp.Regexp)
)

func builtinRegexIsValid(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	if s, err := builtins.StringOperand(operands[0].Value, 1); err == nil {
		if _, err = syntax.Parse(string(s), syntax.Perl); err == nil {
			return iter(ast.InternedTerm(true))
		}
	}

	return iter(ast.InternedTerm(false))
}

func builtinRegexMatch(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	s1, err := builtins.StringOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}
	s2, err := builtins.StringOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}
	re, err := getRegexp(bctx, string(s1))
	if err != nil {
		return err
	}
	return iter(ast.InternedTerm(re.MatchString(string(s2))))
}

func builtinRegexMatchTemplate(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	pattern, err := builtins.StringOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}
	match, err := builtins.StringOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}
	start, err := builtins.StringOperand(operands[2].Value, 3)
	if err != nil {
		return err
	}
	end, err := builtins.StringOperand(operands[3].Value, 4)
	if err != nil {
		return err
	}
	if len(start) != 1 {
		return fmt.Errorf("start delimiter has to be exactly one character long but is %d long", len(start))
	}
	if len(end) != 1 {
		return fmt.Errorf("end delimiter has to be exactly one character long but is %d long", len(start))
	}
	re, err := getRegexpTemplate(string(pattern), string(start)[0], string(end)[0])
	if err != nil {
		return err
	}
	return iter(ast.InternedTerm(re.MatchString(string(match))))
}

func builtinRegexSplit(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	s1, err := builtins.StringOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}
	s2, err := builtins.StringOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}
	re, err := getRegexp(bctx, string(s1))
	if err != nil {
		return err
	}

	elems := re.Split(string(s2), -1)
	arr := make([]*ast.Term, len(elems))
	for i := range elems {
		arr[i] = ast.StringTerm(elems[i])
	}
	return iter(ast.ArrayTerm(arr...))
}

func getRegexp(bctx BuiltinContext, pat string) (*regexp.Regexp, error) {
	if bctx.InterQueryBuiltinValueCache != nil {
		// TODO: Use named cache
		var key ast.Value = ast.String(pat)
		val, ok := bctx.InterQueryBuiltinValueCache.Get(key)
		if ok {
			res, valid := val.(*regexp.Regexp)
			if !valid {
				// The cache key may exist for a different value type (eg. glob).
				// In this case, we calculate the regex and return the result w/o updating the cache.
				return regexp.Compile(pat)
			}

			bctx.Metrics.Counter(regexInterQueryValueCacheHits).Incr()
			return res, nil
		}

		re, err := regexp.Compile(pat)
		if err != nil {
			return nil, err
		}
		bctx.InterQueryBuiltinValueCache.Insert(key, re)
		return re, nil
	}

	regexpCacheLock.RLock()
	re, ok := regexpCache[pat]
	numCached := len(regexpCache)
	regexpCacheLock.RUnlock()
	if !ok {
		var err error
		re, err = regexp.Compile(pat)
		if err != nil {
			return nil, err
		}

		regexpCacheLock.Lock()
		if numCached >= regexCacheMaxSize {
			// Delete a (semi-)random key to make room for the new one.
			for k := range regexpCache {
				delete(regexpCache, k)
				break
			}
		}
		regexpCache[pat] = re
		regexpCacheLock.Unlock()
	}
	return re, nil
}

func getRegexpTemplate(pat string, delimStart, delimEnd byte) (*regexp.Regexp, error) {
	regexpCacheLock.RLock()
	re, ok := regexpCache[pat]
	regexpCacheLock.RUnlock()
	if !ok {
		var err error
		re, err = compileRegexTemplate(pat, delimStart, delimEnd)
		if err != nil {
			return nil, err
		}
		regexpCacheLock.Lock()
		regexpCache[pat] = re
		regexpCacheLock.Unlock()
	}
	return re, nil
}

func builtinGlobsMatch(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	s1, err := builtins.StringOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}
	s2, err := builtins.StringOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}
	ne, err := gintersect.NonEmpty(string(s1), string(s2))
	if err != nil {
		return err
	}
	return iter(ast.InternedTerm(ne))
}

func builtinRegexFind(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	s1, err := builtins.StringOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}
	s2, err := builtins.StringOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}
	n, err := builtins.IntOperand(operands[2].Value, 3)
	if err != nil {
		return err
	}
	re, err := getRegexp(bctx, string(s1))
	if err != nil {
		return err
	}

	elems := re.FindAllString(string(s2), n)
	arr := make([]*ast.Term, len(elems))
	for i := range elems {
		arr[i] = ast.StringTerm(elems[i])
	}
	return iter(ast.ArrayTerm(arr...))
}

func builtinRegexFindAllStringSubmatch(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	s1, err := builtins.StringOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}
	s2, err := builtins.StringOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}
	n, err := builtins.IntOperand(operands[2].Value, 3)
	if err != nil {
		return err
	}

	re, err := getRegexp(bctx, string(s1))
	if err != nil {
		return err
	}
	matches := re.FindAllStringSubmatch(string(s2), n)

	outer := make([]*ast.Term, len(matches))
	for i := range matches {
		inner := make([]*ast.Term, len(matches[i]))
		for j := range matches[i] {
			inner[j] = ast.StringTerm(matches[i][j])
		}
		outer[i] = ast.ArrayTerm(inner...)
	}

	return iter(ast.ArrayTerm(outer...))
}

func builtinRegexReplace(bctx BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
	base, err := builtins.StringOperand(operands[0].Value, 1)
	if err != nil {
		return err
	}

	pattern, err := builtins.StringOperand(operands[1].Value, 2)
	if err != nil {
		return err
	}

	value, err := builtins.StringOperand(operands[2].Value, 3)
	if err != nil {
		return err
	}

	re, err := getRegexp(bctx, string(pattern))
	if err != nil {
		return err
	}

	// If no cancellation context, use the fast path
	if bctx.Cancel == nil {
		res := re.ReplaceAllString(string(base), string(value))
		if res == string(base) {
			return iter(operands[0])
		}
		return iter(ast.InternedTerm(res))
	}

	// Use sink writer for cancellation-aware replacement
	sink := newSink(ast.RegexReplace.Name, len(base), bctx.Cancel)
	src := []byte(base)
	repl := []byte(value)

	// Find all matches at once to preserve anchor behavior: replace("foo", "^[a-z]", "F") => "Foo"
	allMatches := re.FindAllSubmatchIndex(src, -1)

	lastEnd := 0
	for _, match := range allMatches {
		if _, err := sink.Write(src[lastEnd:match[0]]); err != nil {
			return err
		}

		if _, err := sink.Write(re.Expand(nil, repl, src, match)); err != nil {
			return err
		}

		lastEnd = match[1]
	}

	if _, err := sink.Write(src[lastEnd:]); err != nil {
		return err
	}

	res := sink.String()
	if res == string(base) {
		return iter(operands[0])
	}

	return iter(ast.InternedTerm(res))
}

func init() {
	RegisterBuiltinFunc(ast.RegexIsValid.Name, builtinRegexIsValid)
	RegisterBuiltinFunc(ast.RegexMatch.Name, builtinRegexMatch)
	RegisterBuiltinFunc(ast.RegexMatchDeprecated.Name, builtinRegexMatch)
	RegisterBuiltinFunc(ast.RegexSplit.Name, builtinRegexSplit)
	RegisterBuiltinFunc(ast.GlobsMatch.Name, builtinGlobsMatch)
	RegisterBuiltinFunc(ast.RegexTemplateMatch.Name, builtinRegexMatchTemplate)
	RegisterBuiltinFunc(ast.RegexFind.Name, builtinRegexFind)
	RegisterBuiltinFunc(ast.RegexFindAllStringSubmatch.Name, builtinRegexFindAllStringSubmatch)
	RegisterBuiltinFunc(ast.RegexReplace.Name, builtinRegexReplace)
}