// Copyright 2024 Dolthub, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package framework

import (
	"fmt"
	"strings"

	"github.com/cockroachdb/errors"
	"github.com/dolthub/go-mysql-server/sql"
	"github.com/dolthub/go-mysql-server/sql/analyzer"
	"github.com/dolthub/go-mysql-server/sql/expression/function"
	"github.com/dolthub/vitess/go/vt/sqlparser"

	"github.com/dolthub/doltgresql/postgres/parser/parser"
	pgtypes "github.com/dolthub/doltgresql/server/types"
)

// Catalog contains all of the PostgreSQL functions.
var Catalog = map[string][]FunctionInterface{}

// AggregateCatalog contains all of the PostgreSQL aggregate functions.
var AggregateCatalog = map[string][]AggregateFunctionInterface{}

// initializedFunctions simply states whether Initialize has been called yet.
var initializedFunctions = false

// RegisterFunction registers the given function, so that it will be usable from a running server. This should be called
// from within an init().
func RegisterFunction(f FunctionInterface) {
	if initializedFunctions {
		panic("attempted to register a function after the init() phase")
	}
	switch f := f.(type) {
	case Function0:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case Function1:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case Function1N:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case Function2:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case Function2N:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case Function3:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case Function4:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case Function5:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case Function6:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case Function7:
		name := strings.ToLower(f.Name)
		Catalog[name] = append(Catalog[name], f)
	case InterpretedFunction:
		name := strings.ToLower(f.ID.FunctionName())
		Catalog[name] = append(Catalog[name], f)
	default:
		panic("unhandled function type")
	}
}

// RegisterAggregateFunction registers the given function, so that it will be usable from a running server. This should be called
// from within an init().
func RegisterAggregateFunction(f AggregateFunctionInterface) {
	if initializedFunctions {
		panic("attempted to register a function after the init() phase")
	}
	switch f := f.(type) {
	case Func1Aggregate:
		name := strings.ToLower(f.Name)
		AggregateCatalog[name] = append(AggregateCatalog[name], f)
	default:
		panic(fmt.Sprintf("unhandled function type %T", f))
	}
}

// Initialize handles the initialization of the catalog by overwriting the built-in GMS functions, since they do not
// apply to PostgreSQL (and functions of the same name often have different behavior).
func Initialize(astConvert func(parser.Statement) (sqlparser.Statement, error)) {
	// This should only be called once. We don't use sync.Once since we also want to panic if someone attempts to
	// register a function after initialization.
	if initializedFunctions {
		return
	}
	initializedFunctions = true

	convertToVitess = astConvert
	pgtypes.LoadFunctionFromCatalog = getQuickFunctionForTypes
	analyzer.ExternalFunctionProvider = &FunctionProvider{}
	replaceGmsBuiltIns()
	validateFunctions()
	compileFunctions()
	compileAggs()
}

// replaceGmsBuiltIns replaces all GMS built-ins that have conflicting names with PostgreSQL functions.
func replaceGmsBuiltIns() {
	functionNames := make(map[string]struct{})
	for name := range Catalog {
		functionNames[strings.ToLower(name)] = struct{}{}
	}
	var newBuiltIns []sql.Function
	for _, f := range function.BuiltIns {
		if _, ok := functionNames[strings.ToLower(f.FunctionName())]; !ok {
			newBuiltIns = append(newBuiltIns, f)
		}
	}
	function.BuiltIns = newBuiltIns
}

// validateFunctions panics if any functions are defined incorrectly or ambiguously.
func validateFunctions() {
	for funcName, overloads := range Catalog {
		if err := validateFunction(funcName, overloads); err != nil {
			panic(err)
		}
	}
}

// validateFunction validates whether functions are defined incorrectly or ambiguously.
func validateFunction(funcName string, overloads []FunctionInterface) error {
	// Verify that each function uses the correct Function overload
	for _, functionOverload := range overloads {
		if functionOverload.GetExpectedParameterCount() >= 0 &&
			len(functionOverload.GetParameters()) != functionOverload.GetExpectedParameterCount() {
			return errors.Errorf("function `%s` should have %d arguments but has %d arguments",
				funcName, functionOverload.GetExpectedParameterCount(), len(functionOverload.GetParameters()))
		}
	}
	// Verify that all overloads are unique
	for functionIndex, f1 := range overloads {
		for _, f2 := range overloads[functionIndex+1:] {
			sameCount := 0
			if f1.GetExpectedParameterCount() == f2.GetExpectedParameterCount() {
				f2Parameters := f2.GetParameters()
				for parameterIndex, f1Parameter := range f1.GetParameters() {
					if f1Parameter.Equals(f2Parameters[parameterIndex]) {
						sameCount++
					}
				}
			}
			if sameCount == f1.GetExpectedParameterCount() && f1.GetExpectedParameterCount() > 0 {
				return errors.Errorf("duplicate function overloads on `%s`", funcName)
			}
		}
	}
	return nil
}

// compileNonOperatorFunction creates a CompiledFunction for each overload of the given function.
func compileNonOperatorFunction(funcName string, overloads []FunctionInterface) {
	overloadTree := NewOverloads()
	for _, functionOverload := range overloads {
		if err := overloadTree.Add(functionOverload); err != nil {
			panic(err)
		}
	}

	// Store the compiled function into the engine's built-in functions
	// TODO: don't do this, use an actual contract for communicating these functions to the engine catalog
	createFunc := func(ctx *sql.Context, params ...sql.Expression) (sql.Expression, error) {
		return NewCompiledFunction(ctx, funcName, params, overloadTree, false), nil
	}
	function.BuiltIns = append(function.BuiltIns, sql.FunctionN{
		Name: funcName,
		Fn:   createFunc,
	})
	compiledCatalog[funcName] = createFunc
}

// compileNonOperatorFunction creates a CompiledFunction for each overload of the given function.
func compileAggFunction(funcName string, overloads []AggregateFunctionInterface) {
	var newBuffer NewBufferFn
	overloadTree := NewOverloads()
	for _, functionOverload := range overloads {
		newBuffer = functionOverload.NewBuffer
		if err := overloadTree.Add(functionOverload); err != nil {
			panic(err)
		}
	}

	// Store the compiled function into the engine's built-in functions
	// TODO: don't do this, use an actual contract for communicating these functions to the engine catalog
	createFunc := func(ctx *sql.Context, params ...sql.Expression) (sql.Expression, error) {
		return NewCompiledAggregateFunction(ctx, funcName, params, overloadTree, newBuffer), nil
	}
	function.BuiltIns = append(function.BuiltIns, sql.FunctionN{
		Name: funcName,
		Fn:   createFunc,
	})
	compiledCatalog[funcName] = createFunc
}

// compileFunctions creates a CompiledFunction for each overload of each function in the catalog.
func compileFunctions() {
	for funcName, overloads := range Catalog {
		compileNonOperatorFunction(funcName, overloads)
	}

	// Build the overload for all unary and binary functions based on their operator. This will be used for fallback if
	// an exact match is not found. Compiled functions (which wrap the overload deducer) handle upcasting and other
	// special rules, so it's far more efficient to reuse it for operators. Operators are also a special case since they
	// all have different names, while standard overload deducers work on a function-name basis.
	for signature, functionOverload := range unaryFunctions {
		overloads, ok := unaryOperatorOverloads[signature.Operator]
		if !ok {
			overloads = NewOverloads()
			unaryOperatorOverloads[signature.Operator] = overloads
		}
		if err := overloads.Add(functionOverload); err != nil {
			panic(err)
		}
	}

	for signature, functionOverload := range binaryFunctions {
		overloads, ok := binaryOperatorOverloads[signature.Operator]
		if !ok {
			overloads = NewOverloads()
			binaryOperatorOverloads[signature.Operator] = overloads
		}
		if err := overloads.Add(functionOverload); err != nil {
			panic(err)
		}
	}

	// Add all permutations for the unary and binary operators
	for operator, overload := range unaryOperatorOverloads {
		unaryOperatorPermutations[operator] = overload.overloadsForParams(1)
	}
	for operator, overload := range binaryOperatorOverloads {
		binaryOperatorPermutations[operator] = overload.overloadsForParams(2)
	}
}

func compileAggs() {
	for funcName, overloads := range AggregateCatalog {
		compileAggFunction(funcName, overloads)
	}
}