package ztype

import (
	"database/sql/driver"
	"fmt"
	"github.com/goccy/go-json"
	pgxtype "github.com/jackc/pgx/v5/pgtype"
	"reflect"
)

type JSONCodec struct{}

func (JSONCodec) FormatSupported(format int16) bool {
	return format == pgxtype.TextFormatCode || format == pgxtype.BinaryFormatCode
}

func (JSONCodec) PreferredFormat() int16 {
	return pgxtype.TextFormatCode
}

func (c JSONCodec) PlanEncode(m *pgxtype.Map, oid uint32, format int16, value any) pgxtype.EncodePlan {
	switch value.(type) {
	case string:
		return encodePlanJSONCodecEitherFormatString{}
	case []byte:
		return encodePlanJSONCodecEitherFormatByteSlice{}
	}

	// Because anything can be marshalled the normal wrapping in Map.PlanScan doesn't get a chance to run. So try the
	// appropriate wrappers here.
	for _, f := range []pgxtype.TryWrapEncodePlanFunc{
		pgxtype.TryWrapDerefPointerEncodePlan,
		pgxtype.TryWrapFindUnderlyingTypeEncodePlan,
	} {
		if wrapperPlan, nextValue, ok := f(value); ok {
			if nextPlan := c.PlanEncode(m, oid, format, nextValue); nextPlan != nil {
				wrapperPlan.SetNext(nextPlan)
				return wrapperPlan
			}
		}
	}

	return encodePlanJSONCodecEitherFormatMarshal{}
}

type encodePlanJSONCodecEitherFormatString struct{}

func (encodePlanJSONCodecEitherFormatString) Encode(value any, buf []byte) (newBuf []byte, err error) {
	jsonString := value.(string)
	buf = append(buf, jsonString...)
	return buf, nil
}

type encodePlanJSONCodecEitherFormatByteSlice struct{}

func (encodePlanJSONCodecEitherFormatByteSlice) Encode(value any, buf []byte) (newBuf []byte, err error) {
	jsonBytes := value.([]byte)
	if jsonBytes == nil {
		return nil, nil
	}

	buf = append(buf, jsonBytes...)
	return buf, nil
}

type encodePlanJSONCodecEitherFormatMarshal struct{}

func (encodePlanJSONCodecEitherFormatMarshal) Encode(value any, buf []byte) (newBuf []byte, err error) {
	jsonBytes, err := json.Marshal(value)
	if err != nil {
		return nil, err
	}

	buf = append(buf, jsonBytes...)
	return buf, nil
}

func (JSONCodec) PlanScan(m *pgxtype.Map, oid uint32, format int16, target any) pgxtype.ScanPlan {
	switch target.(type) {
	case *string:
		return scanPlanAnyToString{}
	case *[]byte:
		return scanPlanJSONToByteSlice{}
	case pgxtype.BytesScanner:
		return scanPlanBinaryBytesToBytesScanner{}
	}

	return scanPlanJSONToJSONUnmarshal{}
}

type scanPlanAnyToString struct{}

func (scanPlanAnyToString) Scan(src []byte, dst any) error {
	p := dst.(*string)
	*p = string(src)
	return nil
}

type scanPlanJSONToByteSlice struct{}

func (scanPlanJSONToByteSlice) Scan(src []byte, dst any) error {
	dstBuf := dst.(*[]byte)
	if src == nil {
		*dstBuf = nil
		return nil
	}

	*dstBuf = make([]byte, len(src))
	copy(*dstBuf, src)
	return nil
}

type scanPlanJSONToBytesScanner struct{}

func (scanPlanJSONToBytesScanner) Scan(src []byte, dst any) error {
	scanner := (dst).(pgxtype.BytesScanner)
	return scanner.ScanBytes(src)
}

type scanPlanJSONToJSONUnmarshal struct{}

func (scanPlanJSONToJSONUnmarshal) Scan(src []byte, dst any) error {
	if src == nil {
		dstValue := reflect.ValueOf(dst)
		if dstValue.Kind() == reflect.Ptr {
			el := dstValue.Elem()
			switch el.Kind() {
			case reflect.Ptr, reflect.Slice, reflect.Map:
				el.Set(reflect.Zero(el.Type()))
				return nil
			}
		}

		return fmt.Errorf("cannot scan NULL into %T", dst)
	}

	elem := reflect.ValueOf(dst).Elem()
	elem.Set(reflect.Zero(elem.Type()))

	return json.Unmarshal(src, dst)
}

func (c JSONCodec) DecodeDatabaseSQLValue(m *pgxtype.Map, oid uint32, format int16, src []byte) (driver.Value, error) {
	if src == nil {
		return nil, nil
	}

	dstBuf := make([]byte, len(src))
	copy(dstBuf, src)
	return dstBuf, nil
}

func (c JSONCodec) DecodeValue(m *pgxtype.Map, oid uint32, format int16, src []byte) (any, error) {
	if src == nil {
		return nil, nil
	}

	var dst any
	err := json.Unmarshal(src, &dst)
	return dst, err
}

type scanPlanBinaryBytesToBytesScanner struct{}

func (scanPlanBinaryBytesToBytesScanner) Scan(src []byte, dst any) error {
	scanner := (dst).(pgxtype.BytesScanner)
	return scanner.ScanBytes(src)
}