package utils

import (
	"fmt"
	"net"

	"github.com/vmihailenco/msgpack/v5"
)

type UDPPayload struct {
	MagicNumber int    `msgpack:"magicNumber"`
	Data        string `msgpack:"data"`
}

const (
	BUFFER_SIZE  int = 1024
	MAGIC_NUMBER int = 114514
)

func ListenUDP(
	port string,
	doneCh chan struct{},
) (string, error) {
	addr, err := net.ResolveUDPAddr("udp", ":"+port)
	if err != nil {
		return "", fmt.Errorf("failed to resolve address, %w", err)
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return "", fmt.Errorf("failed to listen on addr '%v', %w", addr, err)
	}
	defer conn.Close()

	dataCh := make(chan string)
	go func() {
		for {
			buffer := make([]byte, BUFFER_SIZE)

			n, _, err := conn.ReadFromUDP(buffer)
			if err != nil {
				continue
			}

			var payload UDPPayload
			err = msgpack.Unmarshal(buffer[:n], &payload)
			if err == nil && payload.MagicNumber == MAGIC_NUMBER {
				dataCh <- payload.Data
				return
			}
		}
	}()

	select {
	case <-doneCh:
		return "", nil
	case data := <-dataCh:
		return data, nil
	}

	return "", fmt.Errorf("failed to receive data")
}

func SendPayload(endpoint, data string) error {
	conn, err := net.Dial("udp", endpoint)
	if err != nil {
		return fmt.Errorf("failed to dial endpoint, %w", err)
	}
	defer conn.Close()

	b, err := msgpack.Marshal(UDPPayload{
		MagicNumber: MAGIC_NUMBER,
		Data:        data,
	})
	if err != nil {
		return fmt.Errorf("failed to marshal payload, %w", err)
	}

	_, err = conn.Write(b)
	if err != nil {
		return fmt.Errorf("failed to send payload, %w", err)
	}

	return nil
}