package deploy

import (
	"encoding/json"
	"net"
	"strconv"
	"strings"
	"time"
)

var RelayScanUDPPorts = []int{53, 51820}

type RelayScanResult struct {
	Host      string         `json:"host"`
	OpenPorts []int          `json:"open_ports"`
	UDPHints  []RelayUDPHint `json:"udp_hints,omitempty"`
}

type RelayUDPHint struct {
	Port    int    `json:"port"`
	Open    bool   `json:"open"`
	Service string `json:"service"`
}

var udpGuessPortsFn func(host string, ports []int) []RelayUDPHint

func SetRelayUDPGuessHook(fn func(host string, ports []int) []RelayUDPHint) { udpGuessPortsFn = fn }

func RunRelayScan(host string, udpGuess bool) RelayScanResult {
	host = strings.TrimSpace(host)
	open := probePorts(host, SubnetReconPorts)
	result := RelayScanResult{Host: host, OpenPorts: append([]int(nil), open...)}
	if udpGuess {
		result.UDPHints = guessRelayUDPPorts(host, RelayScanUDPPorts)
	}
	return result
}

func RunRelayScanJSON(host string, udpGuess bool) (string, error) {
	b, err := json.Marshal(RunRelayScan(host, udpGuess))
	if err != nil {
		return "", err
	}
	return string(b), nil
}

func guessRelayUDPPorts(host string, ports []int) []RelayUDPHint {
	if udpGuessPortsFn != nil {
		return udpGuessPortsFn(host, ports)
	}
	var out []RelayUDPHint
	for _, port := range ports {
		out = append(out, RelayUDPHint{Port: port, Open: probeRelayUDPQuick(host, port), Service: relayUDPServiceLabel(port)})
	}
	return out
}

func relayUDPServiceLabel(port int) string {
	switch port {
	case 53:
		return "dns"
	case 51820:
		return "wireguard"
	default:
		return ""
	}
}

func probeRelayUDPQuick(host string, port int) bool {
	conn, err := net.DialTimeout("udp", net.JoinHostPort(host, strconv.Itoa(port)), 800*time.Millisecond)
	if err != nil {
		return false
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(800 * time.Millisecond))
	if b := relayUDPProbePayload(port); len(b) > 0 {
		_, _ = conn.Write(b)
	}
	buf := make([]byte, 512)
	n, err := conn.Read(buf)
	return err == nil && n > 0
}

func relayUDPProbePayload(port int) []byte {
	switch port {
	case 53:
		return []byte{0x00, 0x01, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
			0x07, 'v', 'e', 'r', 's', 'i', 'o', 'n', 0x04, 'b', 'i', 'n', 'd', 0x00, 0x00, 0x10, 0x00, 0x03}
	case 51820:
		return []byte{0x01, 0x00, 0x00, 0x00}
	default:
		return nil
	}
}