package protocol

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
)

const (
	MsgRegister    byte = 1
	MsgRegisterAck byte = 2
	MsgPing        byte = 3
	MsgPong        byte = 4
	MsgOpenStream  byte = 5
	MsgStreamReady byte = 6
	MsgClose       byte = 7
)

const (
	ConnTypeTCP  = "tcp"
	ConnTypeUDP  = "udp"
	ConnTypeHTTP = "http"
)

type RegisterPayload struct {
	VerifyKey string `json:"verify_key"`
	Version   string `json:"version"`
}

type RegisterAckPayload struct {
	ClientID uint   `json:"client_id"`
	Message  string `json:"message"`
	OK       bool   `json:"ok"`
}

type LinkMeta struct {
	ConnType   string `json:"conn_type"`
	Target     string `json:"target"`
	RemoteAddr string `json:"remote_addr"`
	StreamID   uint32 `json:"stream_id"`
}

func WriteMessage(w io.Writer, msgType byte, payload interface{}) error {
	var body []byte
	var err error
	if payload != nil {
		body, err = json.Marshal(payload)
		if err != nil {
			return err
		}
	}
	header := make([]byte, 5)
	header[0] = msgType
	binary.BigEndian.PutUint32(header[1:], uint32(len(body)))
	if _, err := w.Write(header); err != nil {
		return err
	}
	if len(body) > 0 {
		_, err = w.Write(body)
	}
	return err
}

func ReadMessage(r io.Reader) (byte, []byte, error) {
	header := make([]byte, 5)
	if _, err := io.ReadFull(r, header); err != nil {
		return 0, nil, err
	}
	msgType := header[0]
	length := binary.BigEndian.Uint32(header[1:])
	if length == 0 {
		return msgType, nil, nil
	}
	if length > 1<<20 {
		return 0, nil, fmt.Errorf("message too large: %d", length)
	}
	body := make([]byte, length)
	if _, err := io.ReadFull(r, body); err != nil {
		return 0, nil, err
	}
	return msgType, body, nil
}

func ParseRegister(body []byte) (*RegisterPayload, error) {
	var p RegisterPayload
	if err := json.Unmarshal(body, &p); err != nil {
		return nil, err
	}
	return &p, nil
}

func ParseLinkMeta(body []byte) (*LinkMeta, error) {
	var p LinkMeta
	if err := json.Unmarshal(body, &p); err != nil {
		return nil, err
	}
	return &p, nil
}