// Copyright 2025 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 plpgsql

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/cockroachdb/errors"
	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess"
	"github.com/dolthub/go-mysql-server/sql"
	"github.com/jackc/pgx/v5/pgproto3"

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

// InterpretedFunction is an interface that essentially mirrors the implementation of InterpretedFunction in the
// framework package.
type InterpretedFunction interface {
	ApplyBindings(ctx *sql.Context, stack InterpreterStack, stmt string, bindings []string, enforceType bool) (newStmt string, varFound bool, err error)
	GetParameters() []*pgtypes.DoltgresType
	GetParameterNames() []string
	GetReturn() *pgtypes.DoltgresType
	GetStatements() []InterpreterOperation
	QueryMultiReturn(ctx *sql.Context, stack InterpreterStack, stmt string, bindings []string) (schema sql.Schema, rows []sql.Row, err error)
	QuerySingleReturn(ctx *sql.Context, stack InterpreterStack, stmt string, targetType *pgtypes.DoltgresType, bindings []string) (val any, err error)
	// IsSRF returns whether the function is a set returning function, meaning whether the
	// function returns one or more rows as a result.
	IsSRF() bool
}

// GetTypesCollectionFromContext is declared within the core package, but is assigned to this variable to work around
// import cycles.
var GetTypesCollectionFromContext func(ctx *sql.Context, database string) (*typecollection.TypeCollection, error)

// Call runs the contained operations on the given runner.
func Call(ctx *sql.Context, iFunc InterpretedFunction, runner sql.StatementRunner, paramsAndReturn []*pgtypes.DoltgresType, vals []any) (any, error) {
	// Set up the initial state of the function
	stack := NewInterpreterStack(runner)
	// Add the parameters
	parameterTypes := iFunc.GetParameters()
	parameterNames := iFunc.GetParameterNames()
	if len(vals) != len(parameterTypes) {
		return nil, fmt.Errorf("parameter count mismatch: expected %d got %d", len(parameterTypes), len(vals))
	}
	for i := range vals {
		stack.NewVariableWithValue(parameterNames[i], parameterTypes[i], vals[i])
	}
	return call(ctx, iFunc, stack)
}

// TriggerCall runs the contained trigger operations on the given runner.
func TriggerCall(ctx *sql.Context, iFunc InterpretedFunction, runner sql.StatementRunner, sch sql.Schema, oldRow sql.Row, newRow sql.Row, trigVars map[string]any) (any, error) {
	// Set up the initial state of the function
	stack := NewInterpreterStack(runner)
	// Add the special variables
	stack.NewRecord("OLD", sch, oldRow)
	stack.NewRecord("NEW", sch, newRow)
	for varName, val := range trigVars {
		varType, ok := triggerSpecialVariables[varName]
		if !ok {
			return nil, fmt.Errorf("unknown variable %s for trigger", varName)
		}
		stack.NewVariableWithValue(varName, varType, val)
	}
	return call(ctx, iFunc, stack)
}

// call runs the contained operations on the given runner.
func call(ctx *sql.Context, iFunc InterpretedFunction, stack InterpreterStack) (any, error) {
	// We increment before accessing, so start at -1
	counter := -1
	// Run the statements
	statements := iFunc.GetStatements()
	for {
		counter++
		if counter >= len(statements) {
			break
		} else if counter < 0 {
			panic("negative function counter")
		}

		operation := statements[counter]
		switch operation.OpCode {
		case OpCode_Alias:
			iv := stack.GetVariable(operation.PrimaryData)
			if iv.Type == nil {
				return nil, fmt.Errorf("variable `%s` could not be found", operation.PrimaryData)
			}
			stack.NewVariableAlias(operation.Target, operation.PrimaryData)
		case OpCode_Assign:
			iv := stack.GetVariable(operation.Target)
			if iv.Type == nil {
				return nil, fmt.Errorf("variable `%s` could not be found", operation.Target)
			}
			retVal, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, iv.Type, operation.SecondaryData)
			if err != nil {
				return nil, err
			}
			err = stack.SetVariable(ctx, operation.Target, retVal)
			if err != nil {
				return nil, err
			}
		case OpCode_Declare:
			typeCollection, err := GetTypesCollectionFromContext(ctx, "")
			if err != nil {
				return nil, err
			}

			// pg_query_go sets PrimaryData for implicit CASE statement variables to
			// `pg_catalog."integer"`, so we remove double-quotes and extract the schema name.
			typeName := operation.PrimaryData
			typeName = strings.ReplaceAll(typeName, `"`, "")
			schemaName := "pg_catalog"
			if strings.Contains(typeName, ".") {
				parts := strings.Split(typeName, ".")
				schemaName = parts[0]
				typeName = parts[1]
				// Check the NonKeyword type names to see if we're looking at
				// an alias of a type if we're in the pg_catalog schema.
				// Skip array types (names starting with "_") since their internal
				// lookup key uses the "_typename" form, not the "typename[]" form
				// that TypeForNonKeywordTypeName returns.
				if schemaName == "pg_catalog" && !strings.HasPrefix(typeName, "_") {
					typ, ok, _ := types.TypeForNonKeywordTypeName(typeName)
					if ok && typ != nil {
						typeName = typ.Name()
					}
				}
			}
			resolvedType, err := typeCollection.GetType(ctx, id.NewType(schemaName, typeName))
			if err != nil {
				return nil, err
			}
			if resolvedType == nil {
				return nil, pgtypes.ErrTypeDoesNotExist.New(operation.PrimaryData)
			}
			if len(operation.SecondaryData) != 0 {
				defVal := operation.SecondaryData[0]
				// Default value can be a literal value or a reference to parameter
				isParam := false
				for _, param := range iFunc.GetParameterNames() {
					if param == defVal {
						isParam = true
						break
					}
				}
				if isParam {
					ivr := stack.GetVariable(defVal)
					if ivr.Value != nil {
						stack.NewVariableWithValue(operation.Target, resolvedType, *ivr.Value)
					} else {
						stack.NewVariable(operation.Target, resolvedType)
					}
				} else {
					val, err := resolvedType.IoInput(ctx, strings.Trim(operation.SecondaryData[0], "'"))
					if err != nil {
						return nil, err
					}
					stack.NewVariableWithValue(operation.Target, resolvedType, val)
				}
			} else {
				stack.NewVariable(operation.Target, resolvedType)
			}
		case OpCode_DeleteInto:
			// TODO: implement
		case OpCode_Exception:
			// TODO: implement
		case OpCode_Execute:
			if len(operation.Target) > 0 {
				if vars := strings.Split(operation.Target, ","); len(vars) > 1 {
					// multiple column row result
					sch, rows, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
					if err != nil {
						return nil, err
					}
					if len(rows) > 1 {
						return nil, errors.New("query returned more than one row")
					}
					for i, row := range rows {
						if len(row) != len(vars) {
							return nil, errors.New("number of row values does not match number of schema columns")
						}
						target := stack.GetVariable(vars[i])
						if target.Type == nil {
							return nil, fmt.Errorf("variable `%s` could not be found", operation.Target)
						}
						if sch[i].Type.(*pgtypes.DoltgresType).ID != target.Type.ID {
							return nil, fmt.Errorf("variable type `%s` does not match `%s`", sch[i].Type.String(), target.Type.String())
						}
						err = stack.SetVariable(ctx, vars[i], rows[0][i])
						if err != nil {
							return nil, err
						}
					}
				} else {
					// single column
					target := stack.GetVariable(operation.Target)
					if target.Type == nil {
						return nil, fmt.Errorf("variable `%s` could not be found", operation.Target)
					}
					retVal, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, target.Type, operation.SecondaryData)
					if err != nil {
						return nil, err
					}
					err = stack.SetVariable(ctx, operation.Target, retVal)
					if err != nil {
						return nil, err
					}
				}
			} else {
				_, _, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
				if err != nil {
					return nil, err
				}
			}
		case OpCode_Get:
			// TODO: implement
		case OpCode_Goto:
			// We must compare to the index - 1, so that the increment hits our target
			if counter <= operation.Index {
				for ; counter < operation.Index-1; counter++ {
					switch statements[counter].OpCode {
					case OpCode_ScopeBegin:
						stack.PushScope()
					case OpCode_ScopeEnd:
						stack.PopScope()
					}
				}
			} else {
				for ; counter > operation.Index-1; counter-- {
					switch statements[counter].OpCode {
					case OpCode_ScopeBegin:
						stack.PopScope()
					case OpCode_ScopeEnd:
						stack.PushScope()
					}
				}
			}
		case OpCode_If:
			retVal, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, pgtypes.Bool, operation.SecondaryData)
			if err != nil {
				return nil, err
			}
			if retVal.(bool) {
				// We're never changing the scope, so we can just assign it directly.
				// Also, we must assign to index-1, so that the increment hits our target.
				counter = operation.Index - 1
			}
		case OpCode_InsertInto:
			// TODO: implement
		case OpCode_Perform:
			_, _, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
			if err != nil {
				return nil, err
			}
		case OpCode_Raise:
			// TODO: Use the client_min_messages config param to determine which
			//       notice levels to send to the client.
			// https://www.postgresql.org/docs/current/runtime-config-client.html#GUC-CLIENT-MIN-MESSAGES

			message, err := evaluteNoticeMessage(ctx, iFunc, operation, stack)
			if err != nil {
				return nil, err
			}

			if operation.PrimaryData == "EXCEPTION" {
				// TODO: Notices at the EXCEPTION level should also abort the current tx.
				return nil, errors.New(message)
			} else {
				noticeResponse := &pgproto3.NoticeResponse{
					Severity: operation.PrimaryData,
					Message:  message,
				}
				if err = applyNoticeOptions(ctx, noticeResponse, operation.Options); err != nil {
					return nil, err
				}
				sess := dsess.DSessFromSess(ctx.Session)
				sess.Notice(noticeResponse)
			}
		case OpCode_Return:
			// If RETURN QUERY results are being buffered, return those
			if len(stack.ReturnQueryResults()) > 0 {
				records := stack.ReturnQueryResults()

				rows := make([]sql.Row, len(records))
				for i, record := range records {
					rows[i] = sql.Row{record}
				}

				return sql.RowsToRowIter(rows...), nil
			}

			if len(operation.PrimaryData) == 0 {
				return nil, nil
			}

			// TODO: handle record types properly, we'll special case triggers for now
			if iFunc.GetReturn().ID == pgtypes.Trigger.ID && len(operation.SecondaryData) == 1 {
				normalized := strings.ReplaceAll(strings.ToLower(operation.PrimaryData), " ", "")
				if normalized == "select$1;" {
					if strings.EqualFold(operation.SecondaryData[0], "new") {
						return *stack.GetVariable("NEW").Value, nil
					} else if strings.EqualFold(operation.SecondaryData[0], "old") {
						return *stack.GetVariable("OLD").Value, nil
					}
				}
			}
			val, err := iFunc.QuerySingleReturn(ctx, stack, operation.PrimaryData, iFunc.GetReturn(), operation.SecondaryData)
			if err != nil {
				return nil, err
			}

			// If this is a set returning function, then we need to return a RowIter and wrap
			// the composite value in a sql.Row.
			if iFunc.IsSRF() {
				return sql.RowsToRowIter(sql.Row{val}), nil
			}
			return val, err
		case OpCode_ForQueryInit:
			schema, rows, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
			if err != nil {
				return nil, err
			}
			stack.InitCursor(operation.Target, schema, rows)
		case OpCode_ForQueryNext:
			schema, row, ok := stack.AdvanceCursor(operation.PrimaryData)
			if !ok {
				stack.CloseCursor(operation.PrimaryData)
				// Jump forward past the loop body and back-goto, same mechanism as OpCode_If.
				counter = operation.Index - 1
			} else {
				if err := stack.UpdateRecord(operation.Target, schema, row); err != nil {
					return nil, err
				}
			}
		case OpCode_ReturnQuery:
			schema, rows, err := iFunc.QueryMultiReturn(ctx, stack, operation.PrimaryData, operation.SecondaryData)
			if err != nil {
				return nil, err
			}
			records, err := convertRowsToRecords(schema, rows)
			if err != nil {
				return nil, err
			}
			stack.BufferReturnQueryResults(records)
		case OpCode_ScopeBegin:
			stack.PushScope()
		case OpCode_ScopeEnd:
			stack.PopScope()
		case OpCode_SelectInto:
			// TODO: implement
		case OpCode_UpdateInto:
			// TODO: implement
		default:
			panic("unimplemented opcode")
		}
	}
	return nil, nil
}

// convertRowsToRecords iterates overs |rows| and converts each field in each row
// into a RecordValue. |schema| is specified for type information.
func convertRowsToRecords(schema sql.Schema, rows []sql.Row) ([][]pgtypes.RecordValue, error) {
	records := make([][]pgtypes.RecordValue, 0, len(rows))
	for _, row := range rows {
		record := make([]pgtypes.RecordValue, len(row))
		for i, field := range row {
			t := schema[i].Type
			doltgresType, ok := t.(*pgtypes.DoltgresType)
			if !ok {
				// non-Doltgres types are still used in analysis, but we only support disk serialization
				// for Doltgres types, so we must convert the GMS type to the nearest Doltgres type here.
				// TODO: this conversion isn't fully accurate. expression.GMSCast has additional logic in
				//       its Eval() method to handle types more exactly and also handles converting the
				//       value to ensure it is well formed for the returned DoltgresType. We can't
				//       currently use GMSCast directly here though, because of a dependency cycle, so
				//       that conversion logic needs to be extracted into a package both places can import.
				var err error
				doltgresType, err = pgtypes.FromGmsTypeToDoltgresType(t)
				if err != nil {
					return nil, err
				}
			}

			record[i] = pgtypes.RecordValue{
				Value: field,
				Type:  doltgresType,
			}
		}
		records = append(records, record)
	}

	return records, nil
}

// applyNoticeOptions adds the specified |options| to the |noticeResponse|.
func applyNoticeOptions(ctx *sql.Context, noticeResponse *pgproto3.NoticeResponse, options map[string]string) error {
	for key, value := range options {
		i, err := strconv.Atoi(key)
		if err != nil {
			return err
		}

		switch NoticeOptionType(i) {
		case NoticeOptionTypeErrCode:
			noticeResponse.Code = value
		case NoticeOptionTypeMessage:
			noticeResponse.Message = value
		case NoticeOptionTypeDetail:
			noticeResponse.Detail = value
		case NoticeOptionTypeHint:
			noticeResponse.Hint = value
		case NoticeOptionTypeConstraint:
			noticeResponse.ConstraintName = value
		case NoticeOptionTypeDataType:
			noticeResponse.DataTypeName = value
		case NoticeOptionTypeTable:
			noticeResponse.TableName = value
		case NoticeOptionTypeSchema:
			noticeResponse.SchemaName = value
		default:
			ctx.GetLogger().Warnf("unhandled notice option type: %s", key)
		}
	}
	return nil
}

// evaluteNoticeMessage evaluates the message for a RAISE NOTICE statement, including
// evaluating any specified parameters and plugging them into the message in place of
// the % placeholders.
func evaluteNoticeMessage(ctx *sql.Context, iFunc InterpretedFunction,
	operation InterpreterOperation, stack InterpreterStack) (string, error) {
	message := operation.SecondaryData[0]
	if len(operation.SecondaryData) > 1 {
		params := operation.SecondaryData[1:]
		currentParamIdx := 0

		parts := strings.Split(message, "%%")
		for i, part := range parts {
			for strings.Contains(part, "%") {
				if currentParamIdx >= len(params) {
					return "", errors.New("too few parameters specified for RAISE")
				}
				currentParam := params[currentParamIdx]
				currentParamIdx += 1
				formattedVar, varFound, err := iFunc.ApplyBindings(ctx, stack, "$1", []string{currentParam}, false)
				if varFound {
					if err != nil {
						return "", err
					}
					part = strings.Replace(part, "%", formattedVar, 1)
				} else {
					retVal, err := iFunc.QuerySingleReturn(ctx, stack, fmt.Sprintf("SELECT (%s)::text", currentParam), nil, nil)
					if err != nil {
						return "", err
					}
					stringVal := fmt.Sprintf("%v", retVal) // We should always return a string, but this is just a safety net
					part = strings.Replace(part, "%", stringVal, 1)
				}
			}
			parts[i] = part
		}
		if currentParamIdx < len(params) {
			return "", errors.New("too many parameters specified for RAISE")
		}
		message = strings.Join(parts, "%")
	}
	return message, nil
}

// triggerSpecialVariables are the list of special variables for triggers.
// https://www.postgresql.org/docs/15/plpgsql-trigger.html
// TODO: NEW and OLD variables are handled separately using `InterpreterStack.NewRecord` function.
var triggerSpecialVariables = map[string]*pgtypes.DoltgresType{
	//"NEW":
	//"OLD":
	"TG_NAME":         pgtypes.Name,
	"TG_WHEN":         pgtypes.Text,
	"TG_LEVEL":        pgtypes.Text,
	"TG_OP":           pgtypes.Text,
	"TG_RELID":        pgtypes.Oid,
	"TG_RELNAME":      pgtypes.Name,
	"TG_TABLE_NAME":   pgtypes.Name,
	"TG_TABLE_SCHEMA": pgtypes.Name,
	"TG_NARGS":        pgtypes.Int32,
	"TG_ARGV[]":       pgtypes.TextArray,
}