package control

import (
	"encoding/json"

	"github.com/pkg/errors"
)

const (
	TypeAuthRequest  MessageType = "auth-req"
	TypeAuthResponse MessageType = "auth-res"
	TypeProxyRequest MessageType = "proxy-req"
	TypeCloseProxy   MessageType = "close-proxy"
)

type MessageType string

type BaseMessage struct {
	Type       MessageType     `json:"t"`
	RawPayload json.RawMessage `json:"p"`
}

type Message struct {
	BaseMessage
	Payload interface{} `json:"p"`
}

func (m *Message) UnmarshalJSON(data []byte) error {
	base := &BaseMessage{}

	if err := json.Unmarshal(data, base); err != nil {
		return errors.WithStack(err)
	}

	payload, err := unmarshalPayload(base.Type, base.RawPayload)
	if err != nil {
		return errors.WithStack(err)
	}

	m.Type = base.Type
	m.Payload = payload

	return nil
}

func NewMessage(mType MessageType, payload interface{}) *Message {
	return &Message{
		BaseMessage: BaseMessage{
			Type: mType,
		},
		Payload: payload,
	}
}

func unmarshalPayload(mType MessageType, data []byte) (interface{}, error) {
	var payload interface{}

	switch mType {
	case TypeAuthRequest:
		payload = &AuthRequestPayload{}
	case TypeAuthResponse:
		payload = &AuthResponsePayload{}
	case TypeProxyRequest:
		payload = &ProxyRequestPayload{}
	case TypeCloseProxy:
		payload = &CloseProxyPayload{}
	default:
		return nil, errors.Wrapf(ErrUnexpectedMessage, "unexpected message type '%s'", mType)
	}

	if err := json.Unmarshal(data, payload); err != nil {
		return nil, errors.WithStack(err)
	}

	return payload, nil
}