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

import (
	"encoding/json"
	"time"

	"github.com/cockroachdb/apd/v3"
	"github.com/cockroachdb/errors"
	"github.com/dolthub/dolt/go/store/prolly/tree"
	"github.com/dolthub/go-mysql-server/sql"
	"github.com/dolthub/go-mysql-server/sql/expression"
	"github.com/dolthub/go-mysql-server/sql/types"
	"github.com/dolthub/vitess/go/vt/proto/query"

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

// GMSCast handles the conversion from a GMS expression's type to its Doltgres type that is most similar.
type GMSCast struct {
	sqlChild sql.Expression
}

var _ sql.Expression = (*GMSCast)(nil)

// NewGMSCast returns a new *GMSCast.
func NewGMSCast(child sql.Expression) *GMSCast {
	return &GMSCast{
		sqlChild: child,
	}
}

// Children implements the sql.Expression interface.
func (c *GMSCast) Children() []sql.Expression {
	return []sql.Expression{c.sqlChild}
}

// Child returns the child that is being cast.
func (c *GMSCast) Child() sql.Expression {
	return c.sqlChild
}

// DoltgresType returns the DoltgresType that the cast evaluates to. This is the same value that is returned by Type().
func (c *GMSCast) DoltgresType(ctx *sql.Context) *pgtypes.DoltgresType {
	// GMSCast shouldn't receive a DoltgresType, but we shouldn't error if it happens
	if t, ok := c.sqlChild.Type(ctx).(*pgtypes.DoltgresType); ok {
		return t
	}

	return pgtypes.FromGmsType(c.sqlChild.Type(ctx))
}

// Eval implements the sql.Expression interface.
func (c *GMSCast) Eval(ctx *sql.Context, row sql.Row) (any, error) {
	val, err := c.sqlChild.Eval(ctx, row)
	if err != nil {
		return nil, err
	}
	if val == nil {
		return nil, nil
	}
	// GMSCast shouldn't receive a DoltgresType, but we shouldn't error if it happens
	if _, ok := c.sqlChild.Type(ctx).(*pgtypes.DoltgresType); ok {
		return val, nil
	}
	sqlTyp := c.sqlChild.Type(ctx)
	switch sqlTyp.Type() {
	// Boolean types are a special case because of how they are translated on the wire in Postgres. If we identify a
	// boolean result, we want to convert it from an int back to a boolean.
	case query.Type_INT8:
		if sqlTyp == types.Boolean {
			newVal, _, err := types.Int32.Convert(ctx, val)
			if err != nil {
				return nil, err
			}
			if _, ok := newVal.(int32); !ok {
				return nil, errors.Errorf("GMSCast expected type `int32`, got `%T`", val)
			}
			if newVal.(int32) == 0 {
				return false, nil
			} else {
				return true, nil
			}
		}
		fallthrough
		// In Postgres, Int32 is generally the smallest value returned. But we convert int8 and int16 to this type during
		// schema conversion, which means we must do so here as well to avoid runtime panics.
	case query.Type_INT16, query.Type_INT24, query.Type_INT32, query.Type_YEAR:
		newVal, _, err := types.Int32.Convert(ctx, val)
		if err != nil {
			return nil, err
		}
		if _, ok := newVal.(int32); !ok {
			return nil, errors.Errorf("GMSCast expected type `int32`, got `%T`", val)
		}
		return newVal, nil
	case query.Type_INT64, query.Type_BIT, query.Type_UINT8, query.Type_UINT16, query.Type_UINT24, query.Type_UINT32:
		newVal, _, err := types.Int64.Convert(ctx, val)
		if err != nil {
			return nil, err
		}
		if _, ok := newVal.(int64); !ok {
			return nil, errors.Errorf("GMSCast expected type `int64`, got `%T`", val)
		}
		return newVal, nil
	case query.Type_UINT64:
		// Postgres doesn't have a "public" Uint64 type, so we return a Numeric value
		newVal, _, err := types.InternalDecimalType.Convert(ctx, val)
		if err != nil {
			return nil, err
		}
		dec, ok := newVal.(*apd.Decimal)
		if !ok {
			return nil, errors.Errorf("GMSCast expected type `*apd.Decimal`, got `%T`", val)
		}
		return dec, nil
	case query.Type_FLOAT32:
		newVal, _, err := types.Float32.Convert(ctx, val)
		if err != nil {
			return nil, err
		}
		if _, ok := newVal.(float32); !ok {
			return nil, errors.Errorf("GMSCast expected type `float32`, got `%T`", val)
		}
		return newVal, nil
	case query.Type_FLOAT64:
		newVal, _, err := types.Float64.Convert(ctx, val)
		if err != nil {
			return nil, err
		}
		if _, ok := newVal.(float64); !ok {
			return nil, errors.Errorf("GMSCast expected type `float64`, got `%T`", val)
		}
		return newVal, nil
	case query.Type_DECIMAL:
		newVal, _, err := types.InternalDecimalType.Convert(ctx, val)
		if err != nil {
			return nil, err
		}
		dec, ok := newVal.(*apd.Decimal)
		if !ok {
			return nil, errors.Errorf("GMSCast expected type `*apd.Decimal`, got `%T`", val)
		}
		return dec, nil
	case query.Type_DATE, query.Type_DATETIME, query.Type_TIMESTAMP:
		if val, ok := val.(time.Time); ok {
			return val, nil
		}
		return nil, errors.Errorf("GMSCast expected type `Time`, got `%T`", val)
	case query.Type_TIME:
		if val, ok := val.(types.Timespan); ok {
			return val.String(), nil
		}
		return nil, errors.Errorf("GMSCast expected type `Timespan`, got `%T`", val)
	case query.Type_CHAR, query.Type_VARCHAR, query.Type_TEXT, query.Type_BINARY, query.Type_VARBINARY, query.Type_BLOB, query.Type_SET, query.Type_ENUM:
		newVal, _, err := types.LongText.Convert(ctx, val)
		if err != nil {
			return nil, err
		}
		switch newVal := newVal.(type) {
		case string:
			return newVal, nil
		case sql.StringWrapper:
			return newVal.Unwrap(ctx)
		default:
			return nil, errors.Errorf("GMSCast expected type `string`, got `%T`", val)
		}
	case query.Type_JSON:
		switch val := val.(type) {
		case types.JSONDocument:
			return val.JSONString()
		case tree.IndexedJsonDocument:
			return val.String(), nil
		default:
			// TODO: there are particular dolt tables (dolt_constraint_violations) that return json-marshallable structs
			//  that we need to handle here like this
			bytes, err := json.Marshal(val)
			return string(bytes), err
		}
	case query.Type_NULL_TYPE:
		return nil, nil
	case query.Type_GEOMETRY:
		return nil, errors.Errorf("GMS geometry types are not supported")
	default:
		return nil, errors.Errorf("GMS type `%s` is not supported", c.sqlChild.Type(ctx).String())
	}
}

// IsNullable implements the sql.Expression interface.
func (c *GMSCast) IsNullable(ctx *sql.Context) bool {
	return true
}

// Resolved implements the sql.Expression interface.
func (c *GMSCast) Resolved() bool {
	return c.sqlChild.Resolved()
}

// String implements the sql.Expression interface.
func (c *GMSCast) String() string {
	if gf, ok := c.sqlChild.(*expression.GetField); ok {
		return gf.Name()
	}
	return c.sqlChild.String()
}

// Type implements the sql.Expression interface.
func (c *GMSCast) Type(ctx *sql.Context) sql.Type {
	return c.DoltgresType(ctx)
}

// WithChildren implements the sql.Expression interface.
func (c *GMSCast) WithChildren(ctx *sql.Context, children ...sql.Expression) (sql.Expression, error) {
	if len(children) != 1 {
		return nil, sql.ErrInvalidChildrenNumber.New(c, len(children), 1)
	}
	return &GMSCast{
		sqlChild: children[0],
	}, nil
}