package main

import (
	"context"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/gorcon/rcon"
)

const (
	authID = 0
	cmdID  = 1
	endID  = 2
)

func authenticate(conn net.Conn, password string) error {
	if _, err := rcon.NewPacket(rcon.SERVERDATA_AUTH, authID, password).WriteTo(conn); err != nil {
		return err
	}
	var packet rcon.Packet
	if _, err := packet.ReadFrom(conn); err != nil {
		return err
	}
	if packet.Type != rcon.SERVERDATA_AUTH_RESPONSE {
		return rcon.ErrInvalidAuthResponse
	}
	if packet.ID == -1 {
		return rcon.ErrAuthFailed
	}
	if packet.ID != authID {
		return rcon.ErrInvalidPacketID
	}
	return nil
}

func withRcon(ctx context.Context, address, password string, action func(net.Conn) (string, error)) (string, error) {
	dialer := net.Dialer{Timeout: rcon.DefaultDialTimeout}
	conn, err := dialer.DialContext(ctx, "tcp", address)
	if err != nil {
		return "", err
	}
	defer conn.Close()
	cancel := context.AfterFunc(ctx, func() { _ = conn.Close() })
	defer cancel()
	if err := conn.SetDeadline(time.Now().Add(rcon.DefaultDeadline)); err != nil {
		return "", err
	}
	if err := authenticate(conn, password); err != nil {
		return "", err
	}
	if err := conn.SetDeadline(time.Now().Add(rcon.DefaultDeadline)); err != nil {
		return "", err
	}
	return action(conn)
}

func checkRcon(ctx context.Context, address, password string) error {
	_, err := withRcon(ctx, address, password, func(net.Conn) (string, error) { return "", nil })
	return err
}

func execute(ctx context.Context, address, password, command string) (string, error) {
	return withRcon(ctx, address, password, func(conn net.Conn) (string, error) {
		if _, err := rcon.NewPacket(rcon.SERVERDATA_EXECCOMMAND, cmdID, command).WriteTo(conn); err != nil {
			return "", err
		}
		var first rcon.Packet
		if _, err := first.ReadFrom(conn); err != nil {
			return "", err
		}
		if first.ID != cmdID || first.Type != rcon.SERVERDATA_RESPONSE_VALUE {
			return "", fmt.Errorf("unexpected RCON response")
		}
		// Wait for the command response before sending the delimiter; Minecraft drops pipelined requests.
		if _, err := rcon.NewPacket(rcon.SERVERDATA_RESPONSE_VALUE, endID, "").WriteTo(conn); err != nil {
			return "", err
		}
		var output strings.Builder
		output.WriteString(first.Body())
		for {
			var packet rcon.Packet
			if _, err := packet.ReadFrom(conn); err != nil {
				return "", err
			}
			if packet.ID == endID {
				return output.String(), nil
			}
			if packet.ID != cmdID || packet.Type != rcon.SERVERDATA_RESPONSE_VALUE {
				return "", fmt.Errorf("unexpected RCON response")
			}
			output.WriteString(packet.Body())
		}
	})
}