package grpctool

import (
	"fmt"

	"google.golang.org/grpc/encoding"
	protoenc "google.golang.org/grpc/encoding/proto"
	"google.golang.org/grpc/mem"
)

var (
	ProtoCodec = encoding.GetCodecV2(protoenc.Name)
)

type RawFrame struct {
	Data mem.BufferSlice
}

// RawCodec is a *raw* encoding.Codec.
// This codec treats a gRPC message frame as raw bytes.
type RawCodec struct {
}

func (c RawCodec) Marshal(v any) (mem.BufferSlice, error) {
	out, ok := v.(*RawFrame)
	if !ok {
		return nil, fmt.Errorf("RawCodec.Marshal(): unexpected source message type: %T", v)
	}
	// The owner of the RawFrame will free the buffer when they don't need it, so we grab our own ref.
	// The caller will free it when it's no longer needed.
	out.Data.Ref()
	return out.Data, nil
}

func (c RawCodec) Unmarshal(data mem.BufferSlice, v any) error {
	dst, ok := v.(*RawFrame)
	if !ok {
		return fmt.Errorf("RawCodec.Unmarshal(): unexpected target message type: %T", v)
	}
	// The caller will free the buffer. We need to grab our own ref since we are retaining the buffer and
	// giving it away as part of the RawFrame.
	data.Ref()
	dst.Data = data
	return nil
}

func (c RawCodec) Name() string {
	// Pretend to be a codec for protobuf.
	return protoenc.Name
}

// RawCodecWithProtoFallback is a *raw* encoding.Codec.
// This codec treats a gRPC message as raw bytes if it's RawFrame and falls back to default proto encoding
// for other message types.
type RawCodecWithProtoFallback struct {
}

func (c RawCodecWithProtoFallback) Marshal(v any) (mem.BufferSlice, error) {
	out, ok := v.(*RawFrame)
	if !ok {
		return ProtoCodec.Marshal(v)
	}
	out.Data.Ref()
	return out.Data, nil
}

func (c RawCodecWithProtoFallback) Unmarshal(data mem.BufferSlice, v any) error {
	dst, ok := v.(*RawFrame)
	if !ok {
		return ProtoCodec.Unmarshal(data, v)
	}
	data.Ref()
	dst.Data = data
	return nil
}

func (c RawCodecWithProtoFallback) Name() string {
	// Pretend to be a codec for protobuf.
	return protoenc.Name
}
