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

import (
	"strconv"
	"strings"

	"github.com/cockroachdb/errors"
	"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/transform"
	vitess "github.com/dolthub/vitess/go/vt/sqlparser"
	"github.com/jackc/pgx/v5/pgproto3"
	"github.com/lib/pq/oid"

	"github.com/dolthub/doltgresql/core/dataloader"
	"github.com/dolthub/doltgresql/core/id"
	pgexprs "github.com/dolthub/doltgresql/server/expression"
	"github.com/dolthub/doltgresql/server/functions/framework"
	"github.com/dolthub/doltgresql/server/node"
	pgtypes "github.com/dolthub/doltgresql/server/types"
)

// ErrorResponseSeverity represents the severity of an ErrorResponse message.
type ErrorResponseSeverity string

const (
	ErrorResponseSeverity_Error   ErrorResponseSeverity = "ERROR"
	ErrorResponseSeverity_Fatal   ErrorResponseSeverity = "FATAL"
	ErrorResponseSeverity_Panic   ErrorResponseSeverity = "PANIC"
	ErrorResponseSeverity_Warning ErrorResponseSeverity = "WARNING"
	ErrorResponseSeverity_Notice  ErrorResponseSeverity = "NOTICE"
	ErrorResponseSeverity_Debug   ErrorResponseSeverity = "DEBUG"
	ErrorResponseSeverity_Info    ErrorResponseSeverity = "INFO"
	ErrorResponseSeverity_Log     ErrorResponseSeverity = "LOG"
)

// ReadyForQueryTransactionIndicator indicates the state of the transaction related to the query.
type ReadyForQueryTransactionIndicator byte

const (
	ReadyForQueryTransactionIndicator_Idle                   ReadyForQueryTransactionIndicator = 'I'
	ReadyForQueryTransactionIndicator_TransactionBlock       ReadyForQueryTransactionIndicator = 'T'
	ReadyForQueryTransactionIndicator_FailedTransactionBlock ReadyForQueryTransactionIndicator = 'E'
)

// ConvertedQuery represents a query that has been converted from the Postgres representation to the Vitess
// representation. String may contain the string version of the converted query. AST will contain the tree
// version of the converted query, and is the recommended form to use. If AST is nil, then use the String version,
// otherwise always prefer to AST.
type ConvertedQuery struct {
	String       string
	AST          vitess.Statement
	StatementTag string
}

// copyFromStdinState tracks the metadata for an import of data into a table using a COPY FROM STDIN statement. When
// this statement is processed, the server accepts COPY DATA messages from the client with chunks of data to load
// into a table.
type copyFromStdinState struct {
	// copyFromStdinNode stores the original CopyFrom statement that initiated the CopyData message sequence. This
	// node is used to look at what parameters were specified, such as which table to load data into, file format,
	// delimiters, etc.
	copyFromStdinNode *node.CopyFrom
	// insertNode stores the analyzed insert node that will be used to load the data into the target table. This node
	// only needs to be built once, and can be reused with updates to its underlying data loader.
	insertNode sql.Node
	// dataLoader is the implementation of DataLoader that is used to load each individual CopyData chunk into the
	// target table.
	dataLoader dataloader.DataLoader
	// copyErr stores any error that was returned while processing a CopyData message and loading a chunk of data
	// to the target table. The server needs to keep track of any errors that were encountered while processing chunks
	// so that it can avoid sending a CommandComplete message if an error was encountered after the client already
	// sent a CopyDone message to the server.
	copyErr error
}

type PortalData struct {
	Query        ConvertedQuery
	IsEmptyQuery bool
	Fields       []pgproto3.FieldDescription
	BoundPlan    sql.Node
	FormatCodes  []int16
}

type PreparedStatementData struct {
	Query        ConvertedQuery
	ReturnFields []pgproto3.FieldDescription
	BindVarTypes []uint32
}

// extractBindVarTypes returns types based on the given query plan.
// This function is used to get bind var types for running our prepared
// tests only. A valid prepared query and execution messages must have
// the types defined.
func extractBindVarTypes(ctx *sql.Context, queryPlan sql.Node) ([]uint32, error) {
	inspectNode := queryPlan

	types := make(map[string]uint32)
	var err error
	var extractBindVars func(ctx *sql.Context, n sql.Node, expr sql.Expression) bool
	extractBindVars = func(ctx *sql.Context, n sql.Node, expr sql.Expression) bool {
		if err != nil {
			return false
		}

		switch e := expr.(type) {
		// Subquery doesn't walk its Node child via Expressions, so we must walk it separately here
		case *plan.Subquery:
			transform.InspectExpressionsWithNode(ctx, e.Query, extractBindVars)
		case *expression.BindVar:
			var typOid uint32
			if doltgresType, ok := e.Type(ctx).(*pgtypes.DoltgresType); ok {
				typOid = id.Cache().ToOID(doltgresType.ID.AsId())
			} else if _, ok := e.Type(ctx).(sql.DeferredType); ok {
				// for a deferred type, we can make a guess to its type based on the containing node
				switch n.(type) {
				case *plan.Limit:
					typOid = uint32(oid.T_int4)
				case *plan.Offset:
					typOid = uint32(oid.T_int4)
				default:
					typOid, err = VitessTypeToObjectID(e.Type(ctx))
					if err != nil {
						err = errors.Errorf("could not determine OID for placeholder %s: %e", e.Name, err)
						return false
					}
				}
			} else {
				// TODO: should remove usage non doltgres type
				typOid, err = VitessTypeToObjectID(e.Type(ctx))
				if err != nil {
					err = errors.Errorf("could not determine OID for placeholder %s: %e", e.Name, err)
					return false
				}
			}
			if existingOid, ok := types[e.Name]; ok {
				err = checkCompatibleTypes(ctx, existingOid, typOid, e.Name)
			}
			types[e.Name] = typOid
		case *pgexprs.ExplicitCast:
			if bindVar, ok := e.Child().(*expression.BindVar); ok {
				var typOid uint32
				if doltgresType, ok := bindVar.Type(ctx).(*pgtypes.DoltgresType); ok {
					typOid = id.Cache().ToOID(doltgresType.ID.AsId())
				} else {
					typOid, err = VitessTypeToObjectID(e.Type(ctx))
					if err != nil {
						err = errors.Errorf("could not determine OID for placeholder %s: %e", bindVar.Name, err)
						return false
					}
				}
				if existingOid, ok := types[bindVar.Name]; ok {
					err = checkCompatibleTypes(ctx, existingOid, typOid, bindVar.Name)
				}
				types[bindVar.Name] = typOid
				return false
			}
		// $1::text and similar get converted to a Convert expression wrapping the bindvar
		case *expression.Convert:
			if bindVar, ok := e.Child.(*expression.BindVar); ok {
				var typOid uint32
				typOid, err = VitessTypeToObjectID(e.Type(ctx))
				if err != nil {
					err = errors.Errorf("could not determine OID for placeholder %s: %e", bindVar.Name, err)
					return false
				}
				if existingOid, ok := types[bindVar.Name]; ok {
					err = checkCompatibleTypes(ctx, existingOid, typOid, bindVar.Name)
				}
				types[bindVar.Name] = typOid
				return false
			}
		}
		return true
	}

	transform.InspectExpressionsWithNode(ctx, inspectNode, extractBindVars)

	// Insert nodes are special, as their source expressions are not returned by the Expressions()
	// interface and must be walked separately.
	switch queryPlan := queryPlan.(type) {
	case *plan.InsertInto:
		transform.InspectExpressionsWithNode(ctx, queryPlan.Source, extractBindVars)
	}

	// above finds types of bindvars in unordered form.
	// the list of types needs to be ordered as v1, v2, v3, etc.
	typesArr := make([]uint32, len(types))
	for i, t := range types {
		idx, err := strconv.ParseInt(strings.TrimPrefix(i, "v"), 10, 32)
		if err != nil {
			return nil, errors.Errorf("could not determine the index of placeholder %s: %e", i, err)
		}
		if int(idx-1) >= len(types) {
			return nil, errors.Errorf("could not determine the index of placeholder %s in slice of %d elements", i, len(types))
		}
		typesArr[idx-1] = t
	}
	return typesArr, err
}

// VitessTypeToObjectID returns a type, as defined by Vitess, into a type as defined by Postgres.
// OIDs can be obtained with the following query: `SELECT oid, typname FROM pg_type ORDER BY 1;`
func VitessTypeToObjectID(typ sql.Type) (uint32, error) {
	doltgresType, err := pgtypes.FromGmsTypeToDoltgresType(typ)
	if err != nil {
		return 0, err
	}
	return id.Cache().ToOID(doltgresType.ID.AsId()), nil
}

// checkCompatibleTypes checks if multiple types for which a parameter are used are compatible.
func checkCompatibleTypes(ctx *sql.Context, existingOid, newOid uint32, newName string) error {
	var err error
	existing := pgtypes.GetTypeByID(id.Type(id.Cache().ToInternal(existingOid)))
	newType := pgtypes.GetTypeByID(id.Type(id.Cache().ToInternal(newOid)))
	if _, _, err = framework.FindCommonType(ctx, []*pgtypes.DoltgresType{existing, newType}); err != nil {
		err = errors.Errorf("parameter %s is used for incompatible types: %s and %s", newName, existing.String(), newType.String())
	}
	return err
}