// 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 (
	"fmt"

	"github.com/dolthub/go-mysql-server/sql"
	"github.com/dolthub/go-mysql-server/sql/types"

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

// initJsonB registers the functions to the catalog.
func initJsonB() {
	framework.RegisterFunction(jsonb_in)
	framework.RegisterFunction(jsonb_out)
	framework.RegisterFunction(jsonb_recv)
	framework.RegisterFunction(jsonb_send)
	framework.RegisterFunction(jsonb_cmp)
	framework.RegisterFunction(jsonb_build_array)
	framework.RegisterFunction(jsonb_build_object)

}

// jsonb_in represents the PostgreSQL function of jsonb type IO input.
var jsonb_in = framework.Function1{
	Name:       "jsonb_in",
	Return:     pgtypes.JsonB,
	Parameters: [1]*pgtypes.DoltgresType{pgtypes.Cstring},
	Strict:     true,
	Callable:   json_in_callable,
}

// jsonb_out represents the PostgreSQL function of jsonb type IO output.
var jsonb_out = framework.Function1{
	Name:       "jsonb_out",
	Return:     pgtypes.Cstring,
	Parameters: [1]*pgtypes.DoltgresType{pgtypes.JsonB},
	Strict:     true,
	Callable:   jsonb_out_callable,
}

// jsonb_recv represents the PostgreSQL function of jsonb type IO receive.
var jsonb_recv = framework.Function1{
	Name:       "jsonb_recv",
	Return:     pgtypes.JsonB,
	Parameters: [1]*pgtypes.DoltgresType{pgtypes.Internal},
	Strict:     true,
	Callable: func(ctx *sql.Context, t [2]*pgtypes.DoltgresType, val any) (any, error) {
		data := val.([]byte)
		if data == nil {
			return nil, nil
		}
		if len(data) <= 1 {
			return "", nil
		}
		return t[1].IoInput(ctx, string(data[1:]))
	},
}

// jsonb_send represents the PostgreSQL function of jsonb type IO send.
var jsonb_send = framework.Function1{
	Name:       "jsonb_send",
	Return:     pgtypes.Bytea,
	Parameters: [1]*pgtypes.DoltgresType{pgtypes.JsonB},
	Strict:     true,
	Callable: func(ctx *sql.Context, t [2]*pgtypes.DoltgresType, val any) (any, error) {
		if wrapper, ok := val.(sql.AnyWrapper); ok {
			var err error
			val, err = wrapper.UnwrapAny(ctx)
			if err != nil {
				return nil, err
			}
			if val == nil {
				return nil, nil
			}
		}
		textVal, err := t[0].SQL(ctx, nil, val)
		if err != nil {
			return nil, err
		}
		writer := utils.NewWireWriter()
		writer.WriteUint8(1)
		writer.WriteBytes(textVal.ToBytes())
		return writer.BufferData(), nil
	},
}

// jsonb_cmp represents the PostgreSQL function of jsonb type compare.
var jsonb_cmp = framework.Function2{
	Name:       "jsonb_cmp",
	Return:     pgtypes.Int32,
	Parameters: [2]*pgtypes.DoltgresType{pgtypes.JsonB, pgtypes.JsonB},
	Strict:     true,
	Callable: func(ctx *sql.Context, _ [3]*pgtypes.DoltgresType, val1, val2 any) (any, error) {
		w1, ok1 := val1.(sql.JSONWrapper)
		w2, ok2 := val2.(sql.JSONWrapper)
		if !ok1 || !ok2 {
			return nil, fmt.Errorf("jsonb_cmp: unexpected types %T, %T", val1, val2)
		}
		res, err := types.CompareJSON(ctx, w1, w2)
		return int32(res), err
	},
}

// jsonb_build_array represents the PostgreSQL function jsonb_build_array.
var jsonb_build_array = framework.Function1{
	Name:       "jsonb_build_array",
	Return:     pgtypes.JsonB,
	Parameters: [1]*pgtypes.DoltgresType{pgtypes.AnyArray},
	Variadic:   true,
	Callable:   json_build_array_callable,
}

// jsonb_build_object represents the PostgreSQL function jsonb_build_object.
var jsonb_build_object = framework.Function1{
	Name:       "jsonb_build_object",
	Return:     pgtypes.JsonB,
	Parameters: [1]*pgtypes.DoltgresType{pgtypes.AnyArray},
	Variadic:   true,
	Callable:   json_build_object_callable,
}