package main

import (
    "fmt"
    "os"
    "bytes"
    "encoding/binary"
    "math/rand"
    "net"
//    "strconv"
    "context"
    "time"
    "github.com/pkg/errors"
)

// Server contains all the information retreived from the server query API.
type Server struct {
	Address    string            `json:"address"`
	Hostname   string            `json:"hostname"`
	Players    int               `json:"players"`
	MaxPlayers int               `json:"max_players"`
	Gamemode   string            `json:"gamemode"`
	Language   string            `json:"language"`
	Password   bool              `json:"password"`
	Rules      map[string]string `json:"rules"`
	Ping       int               `json:"ping"`
}

// QueryType represents a query method from the SA:MP set: i, r, c, d, x, p
type QueryType uint8

const (
	// Info is the 'i' packet type
	Info QueryType = 'i'
	// Rules is the 'r' packet type
	Rules QueryType = 'r'
	// Players is the 'c' packet type
	Players QueryType = 'c'
	// Ping is the 'p' packet type
	Ping QueryType = 'p'
)

// Query stores state for masterlist queries
type Query struct {
	addr *net.UDPAddr
	Data Server
}

// NewQuery creates a new query handler for a server
func NewQuery(host string) (query *Query, err error) {
	query = new(Query)

	query.addr, err = net.ResolveUDPAddr("udp", host)
	if err != nil {
		return nil, errors.Wrap(err, "failed to resolve host")
	}

	return query, nil
}

// Close closes a query manager's connection
func (query *Query) Close() error {
	return nil
}

func (query *Query) SendQuery(ctx context.Context, opcode QueryType) (response []byte, err error) {
	request := new(bytes.Buffer)

	port := [2]byte{
		byte(query.addr.Port & 0xFF),
		byte((query.addr.Port >> 8) & 0xFF),
	}

	if err = binary.Write(request, binary.LittleEndian, []byte("VCMP")); err != nil {
		return
	}
	if err = binary.Write(request, binary.LittleEndian, query.addr.IP.To4()); err != nil {
		return
	}
	if err = binary.Write(request, binary.LittleEndian, port[0]); err != nil {
		return
	}
	if err = binary.Write(request, binary.LittleEndian, port[1]); err != nil {
		return
	}
	if err = binary.Write(request, binary.LittleEndian, opcode); err != nil {
		return
	}
	if opcode == Ping {
		p := make([]byte, 4)
		_, err = rand.Read(p)
		if err != nil {
			return
		}
		if err = binary.Write(request, binary.LittleEndian, p); err != nil {
			return
		}
	}

	conn, err := openConnection(query.addr)
	if err != nil {
		return
	}
	defer conn.Close()

	_, err = conn.Write(request.Bytes())
	if err != nil {
		return nil, errors.Wrap(err, "failed to write")
	}

	type resultData struct {
		data  []byte
		bytes int
		err   error
	}
	waitResult := make(chan resultData, 1)

	go func() {
		response := make([]byte, 2048)
		n, errInner := conn.Read(response)
		if errInner != nil {
			waitResult <- resultData{err: errors.Wrap(errInner, "failed to read response")}
			return
		}
		if n > cap(response) {
			waitResult <- resultData{err: errors.New("read response over buffer capacity")}
			return
		}
		waitResult <- resultData{data: response, bytes: n}
	}()

	var result resultData
	select {
	case <-ctx.Done():
		return nil, errors.New("socket read timed out")

	case result = <-waitResult:
		break
	}

	if result.err != nil {
		return nil, result.err
	}

	return result.data[:result.bytes], nil
}

// GetPing sends and receives a packet to measure ping
func (query *Query) GetPing(ctx context.Context) (ping time.Duration, err error) {
	t := time.Now()
	_, err = query.SendQuery(ctx, Ping)
	if err != nil {
		return 0, err
	}
	ping = time.Now().Sub(t)

	return
}

func openConnection(addr *net.UDPAddr) (conn *net.UDPConn, err error) {
	conn, err = net.DialUDP("udp", nil, addr)
	if err != nil {
		return nil, errors.Wrap(err, "failed to dial")
	}
	return
}

func main() {
    args := os.Args

    if len(os.Args) < 2 {
        fmt.Println("Usage: ./vcmpquery ip:port")
        os.Exit(1)
    }
    qIp := args[1]

    ctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
    defer cancel()
    query, err := NewQuery(qIp)
    if err != nil {
        fmt.Println(err)
        os.Exit(1)
    }

    ping, err := query.GetPing(ctx)
    if err != nil {
        fmt.Println(err)
        os.Exit(1)
    }

    fmt.Println("Server online, ping: ", ping)
    os.Exit(0)
}