/*
 * Copyright 2024 CloudWeGo Authors
 *
 * 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 thrift

import (
	"context"
	"encoding/binary"
	"fmt"
	"sync"

	"github.com/cloudwego/gopkg/bufiox"
	"github.com/cloudwego/gopkg/protocol/thrift"
	"github.com/cloudwego/gopkg/protocol/thrift/apache"
	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/pkg/remote"
)

func apacheCodecAvailable(data interface{}) bool {
	return apache.CheckTStruct(data) == nil
}

// skipThriftStruct is used to read a struct bytes from trans when dataLen = 0,
// so that we can use frugalUnmarshal or fastUnmarshal instead of apacheUnmarshal
func skipThriftStruct(trans bufiox.Reader) ([]byte, error) {
	// NOTE: If using go net, the assert will fail and execute the slow path.
	p, ok := trans.(interface {
		NetpollReader() netpoll.Reader
	})
	if !ok {
		return skipThriftStructSlow(trans)
	}
	r := p.NetpollReader()
	if r == nil {
		return skipThriftStructSlow(trans)
	}
	x := newNetpollSkipDecoder(r)
	n, err := x.SkipStruct()
	x.Release()
	if err != nil {
		return nil, err
	}
	return trans.Next(n)
}

func skipThriftStructSlow(trans bufiox.Reader) ([]byte, error) {
	sd := thrift.NewSkipDecoder(trans)
	defer sd.Release()
	buf, err := sd.Next(thrift.STRUCT)
	if err != nil {
		return nil, remote.NewTransError(remote.ProtocolError, err).AppendMessage("caught in skipThriftStruct")
	}
	return buf, nil
}

func apacheMarshal(out bufiox.Writer, ctx context.Context, method string, msgType remote.MessageType, seqID int32, data interface{}) error {
	b, err := out.Malloc(thrift.Binary.MessageBeginLength(method))
	if err != nil {
		return err
	}
	_ = thrift.Binary.WriteMessageBegin(b, method, thrift.TMessageType(msgType), seqID)
	if err := apache.ThriftWrite(out, data); err != nil {
		return err
	}
	return nil
}

func apacheUnmarshal(trans bufiox.Reader, data interface{}) error {
	err := apache.ThriftRead(trans, data)
	if err != nil {
		// use remote.NewTransError to return err back to remote
		// instead of EOF
		return remote.NewTransError(remote.ProtocolError, err)
	}
	return nil
}

// NOTE: only used by `marshalThriftData`
func apacheMarshalData(data interface{}) (buf []byte, err error) {
	// XXX: apache thrift, may have performance issue
	buf = make([]byte, 0, marshalThriftBufferSize)
	bw := bufiox.NewBytesWriter(&buf)
	if err := apache.ThriftWrite(bw, data); err != nil {
		return nil, err
	}
	_ = bw.Flush()
	return buf, nil
}

type netpollSkipDecoder struct {
	n int
	b []byte
	r netpoll.Reader
}

var skipdecoderPool = sync.Pool{
	New: func() any { return &netpollSkipDecoder{} },
}

func newNetpollSkipDecoder(r netpoll.Reader) *netpollSkipDecoder {
	p := skipdecoderPool.Get().(*netpollSkipDecoder)
	p.Reset(r)
	return p
}

func (p *netpollSkipDecoder) Reset(r netpoll.Reader) {
	*p = netpollSkipDecoder{r: r}
}

func (p *netpollSkipDecoder) Release() {
	skipdecoderPool.Put(p)
}

// skipn is small enough that can be inlined.
//
// XXX: as a trade off,
// caller have to use lastbyte or lastbytes after calling skipn without err
func (p *netpollSkipDecoder) skipn(n int) error {
	if p.n+n <= len(p.b) {
		p.n += n
		return nil
	}
	return p.skipnSlow(n)
}

func (p *netpollSkipDecoder) lastbyte() byte {
	return p.b[p.n-1]
}

func (p *netpollSkipDecoder) lastbytes(n int) (b []byte) {
	return p.b[p.n-n:]
}

func (p *netpollSkipDecoder) skipnSlow(n int) error {
	sz := p.n + n

	// trigger underlying conn to read more
	if l := p.r.Len(); l > sz {
		// read as much as possible, luckily, we will have a full buffer
		// then we no need to call p.Peek many times
		sz = l
	}
	b, err := p.r.Peek(sz)
	if err != nil {
		return err
	}
	p.b = b
	p.n += n
	return nil
}

func (p *netpollSkipDecoder) SkipStruct() (int, error) {
	const defaultRecursionDepth = 64
	if err := p.skipType(thrift.STRUCT, defaultRecursionDepth); err != nil {
		return 0, err
	}
	return p.n, nil
}

var (
	errDepthLimitExceeded = thrift.NewProtocolException(
		thrift.DEPTH_LIMIT, "depth limit exceeded")

	errDataLength = thrift.NewProtocolException(
		thrift.INVALID_DATA, "invalid data length")
)

var typeToSize = [256]int8{
	thrift.BOOL:   1,
	thrift.BYTE:   1,
	thrift.DOUBLE: 8,
	thrift.I16:    2,
	thrift.I32:    4,
	thrift.I64:    8,
}

func (p *netpollSkipDecoder) skipType(t thrift.TType, maxdepth int) error {
	if maxdepth == 0 {
		return errDepthLimitExceeded
	}
	if sz := typeToSize[t]; sz > 0 {
		return p.skipn(int(sz))
	}
	switch t {
	case thrift.STRING:
		if err := p.skipn(4); err != nil {
			return err
		}
		sz := int(int32(binary.BigEndian.Uint32(p.lastbytes(4))))
		if sz < 0 {
			return errDataLength
		}
		if err := p.skipn(sz); err != nil {
			return err
		}
	case thrift.STRUCT:
		for {
			if err := p.skipn(1); err != nil { // TType
				return err
			}
			tp := thrift.TType(p.lastbyte())
			if tp == thrift.STOP {
				break
			}
			if sz := typeToSize[tp]; sz > 0 {
				// fastpath
				// Field ID + Value
				if err := p.skipn(2 + int(sz)); err != nil {
					return err
				}
				continue
			}

			// Field ID
			if err := p.skipn(2); err != nil {
				return err
			}
			// Field Value
			if err := p.skipType(tp, maxdepth-1); err != nil {
				return err
			}
		}
	case thrift.MAP:
		// 1 byte key TType, 1 byte value TType, 4 bytes Len
		if err := p.skipn(6); err != nil {
			return err
		}
		b := p.lastbytes(6)
		kt, vt, sz := thrift.TType(b[0]), thrift.TType(b[1]), int32(binary.BigEndian.Uint32(b[2:]))
		if sz < 0 {
			return errDataLength
		}
		ksz, vsz := int(typeToSize[kt]), int(typeToSize[vt])
		if ksz > 0 && vsz > 0 {
			return p.skipn(int(sz) * (ksz + vsz))
		}
		var err error
		for i := int32(0); i < sz; i++ {
			if ksz > 0 {
				err = p.skipn(ksz)
			} else {
				err = p.skipType(kt, maxdepth-1)
			}
			if err != nil {
				return err
			}
			if vsz > 0 {
				err = p.skipn(vsz)
			} else {
				err = p.skipType(vt, maxdepth-1)
			}
			if err != nil {
				return err
			}
		}
	case thrift.SET, thrift.LIST:
		// 1 byte value type, 4 bytes Len
		if err := p.skipn(5); err != nil {
			return err
		}
		b := p.lastbytes(5)
		vt, sz := thrift.TType(b[0]), int32(binary.BigEndian.Uint32(b[1:]))
		if sz < 0 {
			return errDataLength
		}
		if vsz := typeToSize[vt]; vsz > 0 {
			return p.skipn(int(sz) * int(vsz))
		}
		for i := int32(0); i < sz; i++ {
			if err := p.skipType(vt, maxdepth-1); err != nil {
				return err
			}
		}
	default:
		return thrift.NewProtocolException(
			thrift.INVALID_DATA, fmt.Sprintf("unknown data type %d", t))
	}
	return nil
}
