package protocol

import (
	"encoding/binary"
	"errors"
	"io"

	"ananta-ps/pkg/crypto/chacha8"
)

type Frame struct {
	Mode    byte
	Payload []byte
}

func ReadFrame(r io.Reader, cryptor *chacha8.UxChaCha8) (*Frame, error) {
	head := make([]byte, 5)
	if _, err := io.ReadFull(r, head); err != nil {
		return nil, err
	}

	if cryptor != nil {
		cryptor.XorInPlace(head)
	}

	payloadLength := int(binary.LittleEndian.Uint32(head[0:4]))
	if payloadLength < 0 || payloadLength > 16*1024*1024 {
		return nil, errors.New("invalid frame payload length after decrypt")
	}

	payload := make([]byte, payloadLength)
	if payloadLength > 0 {
		if _, err := io.ReadFull(r, payload); err != nil {
			return nil, err
		}
		if cryptor != nil {
			cryptor.XorInPlace(payload)
		}
	}

	return &Frame{
		Mode:    head[4],
		Payload: payload,
	}, nil
}

func WriteFrame(w io.Writer, mode byte, payload []byte, cryptor *chacha8.UxChaCha8) error {
	buf := make([]byte, 5+len(payload))
	binary.LittleEndian.PutUint32(buf[0:4], uint32(len(payload)))
	buf[4] = mode
	copy(buf[5:], payload)

	if cryptor != nil {
		cryptor.XorInPlace(buf)
	}

	_, err := w.Write(buf)
	return err
}

func ParseRpcPacket(payload []byte) RpcPacket {
	if len(payload) < 5 {
		return RpcPacket{Kind: PacketKindUnknown, Body: payload}
	}

	tag := payload[0]
	methodID := binary.LittleEndian.Uint32(payload[1:5])

	switch tag {
	case 0x03: // Notify
		return RpcPacket{
			Kind:     PacketKindNotify,
			MethodID: methodID,
			Body:     payload[5:],
		}
	case 0x01: // Invoke
		if len(payload) >= 9 {
			invokeID := int32(binary.LittleEndian.Uint32(payload[5:9]))
			return RpcPacket{
				Kind:     PacketKindInvoke,
				MethodID: methodID,
				InvokeID: invokeID,
				Body:     payload[9:],
			}
		}
	case 0x02: // Return
		if len(payload) >= 13 {
			invokeID := int32(binary.LittleEndian.Uint32(payload[5:9]))
			errCode := int32(binary.LittleEndian.Uint32(payload[9:13]))
			return RpcPacket{
				Kind:     PacketKindReturn,
				MethodID: methodID,
				InvokeID: invokeID,
				Error:    errCode,
				Body:     payload[13:],
			}
		}
	}

	return RpcPacket{
		Kind:     PacketKindUnknown,
		MethodID: methodID,
		Body:     payload,
	}
}

func BuildNotify(methodID uint32, body []byte) []byte {
	buf := make([]byte, 5+len(body))
	buf[0] = 0x03
	binary.LittleEndian.PutUint32(buf[1:5], methodID)
	copy(buf[5:], body)
	return buf
}

func BuildInvoke(methodID uint32, invokeID int32, body []byte) []byte {
	buf := make([]byte, 9+len(body))
	buf[0] = 0x01
	binary.LittleEndian.PutUint32(buf[1:5], methodID)
	binary.LittleEndian.PutUint32(buf[5:9], uint32(invokeID))
	copy(buf[9:], body)
	return buf
}

func BuildReturn(methodID uint32, invokeID int32, errCode int32, body []byte) []byte {
	buf := make([]byte, 13+len(body))
	buf[0] = 0x02
	binary.LittleEndian.PutUint32(buf[1:5], methodID)
	binary.LittleEndian.PutUint32(buf[5:9], uint32(invokeID))
	binary.LittleEndian.PutUint32(buf[9:13], uint32(errCode))
	copy(buf[13:], body)
	return buf
}