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

import (
	"io"
	"reflect"
	"strconv"
	"time"

	"github.com/cockroachdb/errors"
	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dprocedures"
	"github.com/dolthub/go-mysql-server/sql"
	"github.com/dolthub/go-mysql-server/sql/expression"
	"github.com/dolthub/go-mysql-server/sql/plan"
	"github.com/dolthub/go-mysql-server/sql/types"

	"github.com/dolthub/doltgresql/core"
	"github.com/dolthub/doltgresql/server/auth"
	"github.com/dolthub/doltgresql/server/functions/framework"
	pgtypes "github.com/dolthub/doltgresql/server/types"
)

var (
	ErrDoltProcedurePermissionDenied = errors.New("permission denied for Dolt procedure")
	ErrDoltProcedureSelectOnly       = errors.New("Dolt stored procedure may only be invoked using SELECT")
)

func initDoltProcedures() {
	for _, procDef := range dprocedures.DoltProcedures {
		p, err := resolveExternalStoredProcedure(nil, procDef)
		if err != nil {
			panic(err)
		}

		funcVal := reflect.ValueOf(procDef.Function)
		varArgCallable := varArgCallableForDoltProcedure(p, funcVal)
		noArgCallable := noArgCallableForDoltProcedure(p, funcVal)
		framework.RegisterFunction(framework.Function1{
			Name:       procDef.Name,
			Return:     pgtypes.TextArray,
			Parameters: [1]*pgtypes.DoltgresType{pgtypes.TextArray},
			Variadic:   true,
			Callable:   varArgCallable,
		})
		framework.RegisterFunction(framework.Function0{
			Name:     procDef.Name,
			Return:   pgtypes.TextArray,
			Callable: noArgCallable,
		})
	}
}

// varArgCallableForDoltProcedure creates a callable function that takes in a variadic number of parameters. This is
// equivalent to calling "DOLT_PROC_NAME('abc', ...)".
func varArgCallableForDoltProcedure(p *plan.ExternalProcedure, funcVal reflect.Value) func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) {
	funcType := funcVal.Type()

	return func(ctx *sql.Context, paramsAndReturn [2]*pgtypes.DoltgresType, val1 any) (any, error) {
		err := checkDoltProcedureAccess(ctx, p)
		if err != nil {
			return nil, err
		}

		values, ok := val1.([]any)
		if !ok {
			return nil, sql.ErrExternalProcedureInvalidParamType.New(reflect.TypeOf(val1).String())
		}

		funcParams := make([]reflect.Value, len(values)+1)
		funcParams[0] = reflect.ValueOf(ctx)

		for i := range values {
			paramDefinition := p.ParamDefinitions[0]
			var funcParamType reflect.Type
			if paramDefinition.Variadic {
				funcParamType = funcType.In(funcType.NumIn() - 1).Elem()
			} else {
				// TODO: support non-variadic procedures
				return nil, sql.ErrExternalProcedureInvalidParamType.New(funcType.String())
			}

			// Grab the passed-in variable and convert it to the type we expect
			exprParamVal, _, err := paramDefinition.Type.Convert(ctx, values[i])
			if err != nil {
				return nil, err
			}

			funcParams[i+1], err = p.ProcessParam(ctx, funcParamType, exprParamVal)
			if err != nil {
				return nil, err
			}
		}

		out := funcVal.Call(funcParams)
		if err, ok := out[1].Interface().(error); ok { // Only evaluates to true when error is not nil
			return nil, err
		}

		var rowIter sql.RowIter
		if iter, ok := out[0].Interface().(sql.RowIter); ok {
			rowIter = iter
		} else {
			rowIter = sql.RowsToRowIter()
		}

		return drainRowIter(ctx, rowIter)
	}
}

// noArgCallableForDoltProcedure creates a callable function that does not take any parameters. This is equivalent to
// calling "DOLT_PROC_NAME()".
func noArgCallableForDoltProcedure(p *plan.ExternalProcedure, funcVal reflect.Value) func(ctx *sql.Context) (any, error) {
	return func(ctx *sql.Context) (any, error) {
		err := checkDoltProcedureAccess(ctx, p)
		if err != nil {
			return nil, err
		}

		funcParams := []reflect.Value{reflect.ValueOf(ctx)}
		out := funcVal.Call(funcParams)
		if err, ok := out[1].Interface().(error); ok { // Only evaluates to true when error is not nil
			return nil, err
		}

		var rowIter sql.RowIter
		if iter, ok := out[0].Interface().(sql.RowIter); ok {
			rowIter = iter
		} else {
			rowIter = sql.RowsToRowIter()
		}
		return drainRowIter(ctx, rowIter)
	}
}

// checkDoltProcedureAccess ensures the current user is authorized as a SUPERUSER if the given |procedure| requires
// admin.
func checkDoltProcedureAccess(ctx *sql.Context, procedure *plan.ExternalProcedure) error {
	if !procedure.AdminOnly {
		return nil
	}

	var userRole auth.Role
	auth.LockRead(func() {
		userRole = auth.GetRole(ctx.Client().User)
	})

	if !userRole.IsValid() || !userRole.IsSuperUser {
		return ErrDoltProcedurePermissionDenied
	}

	return nil
}

func drainRowIter(ctx *sql.Context, rowIter sql.RowIter) (any, error) {
	defer rowIter.Close(ctx)

	row, err := rowIter.Next(ctx)
	if err == io.EOF {
		return nil, nil
	} else if err != nil {
		return nil, err
	}
	castsColl, err := core.GetCastsCollectionFromContext(ctx, "")
	if err != nil {
		return nil, err
	}

	// The conversion to []text needs []any, not sql.Row
	rowSlice := make([]any, len(row))
	for i := range row {
		sourceType, err := typeForElement(row[i])
		if err != nil {
			return nil, err
		}
		cast, err := castsColl.GetExplicitCast(ctx, sourceType, pgtypes.Text)
		if err != nil {
			return nil, err
		}
		textVal, err := cast.Eval(ctx, row[i], sourceType, pgtypes.Text)
		if err != nil {
			return nil, err
		}

		rowSlice[i] = textVal
	}
	return rowSlice, nil
}

func typeForElement(v any) (*pgtypes.DoltgresType, error) {
	switch x := v.(type) {
	case int64:
		return pgtypes.Int64, nil
	case int32:
		return pgtypes.Int32, nil
	case int16, int8:
		return pgtypes.Int16, nil
	case string:
		return pgtypes.Text, nil
	default:
		return nil, errors.Errorf("dolt_procedures: unsupported type %T", x)
	}
}

var (
	// ctxType is the reflect.Type of a *sql.Context.
	ctxType = reflect.TypeOf((*sql.Context)(nil))
	// ctxType is the reflect.Type of a sql.RowIter.
	rowIterType = reflect.TypeOf((*sql.RowIter)(nil)).Elem()
	// ctxType is the reflect.Type of an error.
	errorType = reflect.TypeOf((*error)(nil)).Elem()
	// externalStoredProcedurePointerTypes maps a non-pointer type to a sql.Type for external stored procedures.
	externalStoredProcedureTypes = map[reflect.Type]sql.Type{
		reflect.TypeOf(int(0)):      types.Int64,
		reflect.TypeOf(int8(0)):     types.Int8,
		reflect.TypeOf(int16(0)):    types.Int16,
		reflect.TypeOf(int32(0)):    types.Int32,
		reflect.TypeOf(int64(0)):    types.Int64,
		reflect.TypeOf(uint(0)):     types.Uint64,
		reflect.TypeOf(uint8(0)):    types.Uint8,
		reflect.TypeOf(uint16(0)):   types.Uint16,
		reflect.TypeOf(uint32(0)):   types.Uint32,
		reflect.TypeOf(uint64(0)):   types.Uint64,
		reflect.TypeOf(float32(0)):  types.Float32,
		reflect.TypeOf(float64(0)):  types.Float64,
		reflect.TypeOf(bool(false)): types.Int8,
		reflect.TypeOf(string("")):  types.LongText,
		reflect.TypeOf([]byte{}):    types.LongBlob,
		reflect.TypeOf(time.Time{}): types.DatetimeMaxPrecision,
	}
	// externalStoredProcedurePointerTypes maps a pointer type to a sql.Type for external stored procedures.
	externalStoredProcedurePointerTypes = map[reflect.Type]sql.Type{
		reflect.TypeOf((*int)(nil)):       types.Int64,
		reflect.TypeOf((*int8)(nil)):      types.Int8,
		reflect.TypeOf((*int16)(nil)):     types.Int16,
		reflect.TypeOf((*int32)(nil)):     types.Int32,
		reflect.TypeOf((*int64)(nil)):     types.Int64,
		reflect.TypeOf((*uint)(nil)):      types.Uint64,
		reflect.TypeOf((*uint8)(nil)):     types.Uint8,
		reflect.TypeOf((*uint16)(nil)):    types.Uint16,
		reflect.TypeOf((*uint32)(nil)):    types.Uint32,
		reflect.TypeOf((*uint64)(nil)):    types.Uint64,
		reflect.TypeOf((*float32)(nil)):   types.Float32,
		reflect.TypeOf((*float64)(nil)):   types.Float64,
		reflect.TypeOf((*bool)(nil)):      types.Int8,
		reflect.TypeOf((*string)(nil)):    types.LongText,
		reflect.TypeOf((*[]byte)(nil)):    types.LongBlob,
		reflect.TypeOf((*time.Time)(nil)): types.DatetimeMaxPrecision,
	}
)

func init() {
	if strconv.IntSize == 32 {
		externalStoredProcedureTypes[reflect.TypeOf(int(0))] = types.Int32
		externalStoredProcedureTypes[reflect.TypeOf(uint(0))] = types.Uint32
		externalStoredProcedurePointerTypes[reflect.TypeOf((*int)(nil))] = types.Int32
		externalStoredProcedurePointerTypes[reflect.TypeOf((*uint)(nil))] = types.Uint32
	}
}

func resolveExternalStoredProcedure(_ *sql.Context, externalProcedure sql.ExternalStoredProcedureDetails) (*plan.ExternalProcedure, error) {
	funcVal := reflect.ValueOf(externalProcedure.Function)
	funcType := funcVal.Type()
	if funcType.Kind() != reflect.Func {
		return nil, sql.ErrExternalProcedureNonFunction.New(externalProcedure.Function)
	}
	if funcType.NumIn() == 0 {
		return nil, sql.ErrExternalProcedureMissingContextParam.New()
	}
	if funcType.NumOut() != 2 {
		return nil, sql.ErrExternalProcedureReturnTypes.New()
	}
	if funcType.In(0) != ctxType {
		return nil, sql.ErrExternalProcedureMissingContextParam.New()
	}
	if funcType.Out(0) != rowIterType {
		return nil, sql.ErrExternalProcedureFirstReturn.New()
	}
	if funcType.Out(1) != errorType {
		return nil, sql.ErrExternalProcedureSecondReturn.New()
	}
	funcIsVariadic := funcType.IsVariadic()

	paramDefinitions := make([]plan.ProcedureParam, funcType.NumIn()-1)
	paramReferences := make([]*expression.ProcedureParam, len(paramDefinitions))
	for i := 0; i < len(paramDefinitions); i++ {
		funcParamType := funcType.In(i + 1)
		paramName := "A" + strconv.FormatInt(int64(i), 10)
		paramIsVariadic := false
		if funcIsVariadic && i == len(paramDefinitions)-1 {
			paramIsVariadic = true
			funcParamType = funcParamType.Elem()
			if funcParamType.Kind() == reflect.Ptr {
				return nil, sql.ErrExternalProcedurePointerVariadic.New()
			}
		}

		if sqlType, ok := externalStoredProcedureTypes[funcParamType]; ok {
			paramDefinitions[i] = plan.ProcedureParam{
				Direction: plan.ProcedureParamDirection_In,
				Name:      paramName,
				Type:      sqlType,
				Variadic:  paramIsVariadic,
			}
			paramReferences[i] = expression.NewProcedureParam(paramName, sqlType)
		} else if sqlType, ok = externalStoredProcedurePointerTypes[funcParamType]; ok {
			paramDefinitions[i] = plan.ProcedureParam{
				Direction: plan.ProcedureParamDirection_Inout,
				Name:      paramName,
				Type:      sqlType,
				Variadic:  paramIsVariadic,
			}
			paramReferences[i] = expression.NewProcedureParam(paramName, sqlType)
		} else {
			return nil, sql.ErrExternalProcedureInvalidParamType.New(funcParamType.String())
		}
	}

	return &plan.ExternalProcedure{
		ExternalStoredProcedureDetails: externalProcedure,
		ParamDefinitions:               paramDefinitions,
		Params:                         paramReferences,
	}, nil
}