package dht

import (
	"bytes"
	"fmt"

	"github.com/jackpal/bencode-go"
)

// Msg represents a KRPC message.
type Msg struct {
	T string                 `bencode:"t"`
	Y string                 `bencode:"y"`
	Q string                 `bencode:"q,omitempty"`
	A map[string]interface{} `bencode:"a,omitempty"`
	R map[string]interface{} `bencode:"r,omitempty"`
	E []interface{}          `bencode:"e,omitempty"`
}

// EncodeMsg marshals a KRPC message to bencode.
func EncodeMsg(msg Msg) ([]byte, error) {
	m := make(map[string]interface{})
	m["t"] = msg.T
	m["y"] = msg.Y
	if msg.Q != "" {
		m["q"] = msg.Q
	}
	if msg.A != nil {
		m["a"] = msg.A
	}
	if msg.R != nil {
		m["r"] = msg.R
	}
	if msg.E != nil {
		m["e"] = msg.E
	}

	var buf bytes.Buffer
	err := bencode.Marshal(&buf, m)
	return buf.Bytes(), err
}

// DecodeMsg unmarshals a KRPC message from bencode.
func DecodeMsg(data []byte) (Msg, error) {
	val, err := bencode.Decode(bytes.NewReader(data))
	if err != nil {
		return Msg{}, err
	}
	m, ok := val.(map[string]interface{})
	if !ok {
		return Msg{}, fmt.Errorf("expected map[string]interface{}, got %T", val)
	}

	var msg Msg
	if t, ok := m["t"].(string); ok {
		msg.T = t
	}
	if y, ok := m["y"].(string); ok {
		msg.Y = y
	}
	if q, ok := m["q"].(string); ok {
		msg.Q = q
	}
	if a, ok := m["a"].(map[string]interface{}); ok {
		msg.A = a
	}
	if r, ok := m["r"].(map[string]interface{}); ok {
		msg.R = r
	}
	if e, ok := m["e"].([]interface{}); ok {
		msg.E = e
	}
	return msg, nil
}

// NewQuery creates a KRPC query message.
func NewQuery(t string, q string, a map[string]interface{}) Msg {
	return Msg{
		T: t,
		Y: "q",
		Q: q,
		A: a,
	}
}