// Package ping implements the "ping" check (ICMP echo through an unprivileged datagram socket) and the "tcp"
// check (sequential TCP connects), which is also the fallback when ICMP sockets are not available.
//
// Unprivileged ICMP: golang.org/x/net/icmp "udp4"/"udp6" endpoints work out of the box on macOS; on Linux
// they require the probe's group to be within net.ipv4.ping_group_range (see README). When the socket cannot be
// opened the check falls back to kind "tcp" and reports the reason "icmp_unavailable".
package ping

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"math"
	"net"
	"os"
	"runtime"
	"sort"
	"time"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"

	"internetpressure.io/probe-agent/internal/protocol"
)

const (
	// Count is the number of echoes / connects per check.
	Count = 5
	// Spacing between consecutive echoes.
	Spacing = 200 * time.Millisecond
	// EchoTimeout is the wait for each reply.
	EchoTimeout = 2 * time.Second
	// ResolveTimeout bounds the one hostname lookup per run.
	ResolveTimeout = 3 * time.Second

	protoICMP   = 1  // iana.ProtocolICMP
	protoICMPv6 = 58 // iana.ProtocolIPv6ICMP
)

// Checker runs ping / tcp checks.
type Checker struct {
	// ICMP is true when an unprivileged ICMP socket could be opened at start-up (see Detect).
	ICMP bool
}

// Detect reports whether an unprivileged ICMPv4 datagram socket can be opened on this host.
func Detect() bool {
	c, err := icmp.ListenPacket("udp4", "0.0.0.0")
	if err != nil {
		return false
	}
	c.Close()
	return true
}

// Resolve picks the address to ping: the pinned IP if any, else the first IPv4 (then IPv6) of the hostname,
// resolved once with the system resolver.
func Resolve(ctx context.Context, t protocol.Target) (net.IP, error) {
	if ip := t.FixedIP(); ip != "" {
		parsed := net.ParseIP(ip)
		if parsed == nil {
			return nil, errors.New("invalid pinned ip")
		}
		return parsed, nil
	}
	ctx, cancel := context.WithTimeout(ctx, ResolveTimeout)
	defer cancel()
	addrs, err := net.DefaultResolver.LookupIPAddr(ctx, t.Hostname)
	if err != nil {
		return nil, err
	}
	for _, a := range addrs {
		if a.IP.To4() != nil {
			return a.IP, nil
		}
	}
	if len(addrs) > 0 {
		return addrs[0].IP, nil
	}
	return nil, errors.New("no address")
}

// Run performs the ping check for t: 5 ICMP echoes 200 ms apart, or the tcp fallback.
func (c *Checker) Run(ctx context.Context, t protocol.Target) protocol.Measurement {
	m := protocol.Measurement{TS: protocol.FormatTime(time.Now()), TargetID: t.TargetID, Kind: "ping"}
	ip, err := Resolve(ctx, t)
	if err != nil {
		m.Error = protocol.ErrDNSFail
		m.Sent, m.Received, m.PacketLoss = protocol.I(0), protocol.I(0), protocol.F(1)
		return m
	}
	m.ResolvedIP = ip.String()

	rtts, sent, icmpErr := echo(ctx, ip)
	if icmpErr != nil {
		// Socket could not be opened → tcp fallback with the reason kept in error.
		port := t.Port
		if port == 0 {
			port = 443
		}
		rtts, sent = connects(ctx, ip, port)
		m.Kind = "tcp"
		summarise(&m, rtts, sent)
		if m.OK {
			m.Error = protocol.ErrICMPUnavailable
		}
		return m
	}
	summarise(&m, rtts, sent)
	return m
}

// RunTCP performs the "tcp" check (5 sequential connects to t.Port, default 443).
func (c *Checker) RunTCP(ctx context.Context, t protocol.Target) protocol.Measurement {
	m := protocol.Measurement{TS: protocol.FormatTime(time.Now()), TargetID: t.TargetID, Kind: "tcp"}
	ip, err := Resolve(ctx, t)
	if err != nil {
		m.Error = protocol.ErrDNSFail
		m.Sent, m.Received, m.PacketLoss = protocol.I(0), protocol.I(0), protocol.F(1)
		return m
	}
	m.ResolvedIP = ip.String()
	port := t.Port
	if port == 0 {
		port = 443
	}
	rtts, sent := connects(ctx, ip, port)
	summarise(&m, rtts, sent)
	return m
}

// summarise fills sent/received/loss/rtt stats/jitter and ok/error.
func summarise(m *protocol.Measurement, rtts []time.Duration, sent int) {
	m.Sent = protocol.I(sent)
	m.Received = protocol.I(len(rtts))
	if sent > 0 {
		m.PacketLoss = protocol.F(math.Round((1-float64(len(rtts))/float64(sent))*1000) / 1000)
	} else {
		m.PacketLoss = protocol.F(1)
	}
	if len(rtts) == 0 {
		m.Error = protocol.ErrUnreachable
		return
	}
	m.OK = true
	// Jitter = mean absolute successive difference, in arrival order.
	var jitter float64
	for i := 1; i < len(rtts); i++ {
		jitter += math.Abs(protocol.Ms(rtts[i]) - protocol.Ms(rtts[i-1]))
	}
	if len(rtts) > 1 {
		jitter /= float64(len(rtts) - 1)
	}
	m.JitterMs = protocol.F(round1(jitter))

	sorted := append([]time.Duration(nil), rtts...)
	sort.Slice(sorted, func(i, j int) bool { return sorted[i] < sorted[j] })
	var sum time.Duration
	for _, r := range sorted {
		sum += r
	}
	m.RTTMinMs = protocol.F(protocol.Ms(sorted[0]))
	m.RTTMaxMs = protocol.F(protocol.Ms(sorted[len(sorted)-1]))
	m.RTTAvgMs = protocol.F(round1(protocol.Ms(sum) / float64(len(sorted))))
}

func round1(v float64) float64 { return math.Round(v*1000) / 1000 }

// echo sends Count ICMP echo requests, one at a time, waiting up to EchoTimeout for each reply and keeping
// Spacing between sends. It returns (rtts of received replies, number sent, socket error).
func echo(ctx context.Context, dst net.IP) ([]time.Duration, int, error) {
	network, laddr, proto := "udp4", "0.0.0.0", protoICMP
	var typ, replyType icmp.Type = ipv4.ICMPTypeEcho, ipv4.ICMPTypeEchoReply
	if dst.To4() == nil {
		network, laddr, proto = "udp6", "::", protoICMPv6
		typ, replyType = ipv6.ICMPTypeEchoRequest, ipv6.ICMPTypeEchoReply
	}
	conn, err := icmp.ListenPacket(network, laddr)
	if err != nil {
		return nil, 0, err
	}
	defer conn.Close()

	var token [8]byte
	_, _ = rand.Read(token[:])
	id := os.Getpid() & 0xffff
	peer := &net.UDPAddr{IP: dst}
	buf := make([]byte, 1500)
	var rtts []time.Duration
	sent := 0

	for seq := 1; seq <= Count; seq++ {
		if ctx.Err() != nil {
			break
		}
		payload := make([]byte, 16)
		copy(payload, token[:])
		binary.BigEndian.PutUint64(payload[8:], uint64(time.Now().UnixNano()))
		msg := icmp.Message{Type: typ, Code: 0, Body: &icmp.Echo{ID: id, Seq: seq, Data: payload}}
		wb, err := msg.Marshal(nil)
		if err != nil {
			return rtts, sent, nil
		}
		sendAt := time.Now()
		if _, err := conn.WriteTo(wb, peer); err != nil {
			sent++
			sleepUntil(ctx, sendAt.Add(Spacing))
			continue
		}
		sent++
		deadline := sendAt.Add(EchoTimeout)
		if d, ok := ctx.Deadline(); ok && d.Before(deadline) {
			deadline = d
		}
		for {
			_ = conn.SetReadDeadline(deadline)
			n, from, err := conn.ReadFrom(buf)
			if err != nil {
				break // timeout for this echo
			}
			if !sameIP(from, dst) {
				continue
			}
			body, ok := parseReply(buf[:n], proto, replyType)
			if !ok || body.Seq != seq || len(body.Data) < 8 || string(body.Data[:8]) != string(token[:]) {
				continue
			}
			rtts = append(rtts, time.Since(sendAt))
			break
		}
		if seq < Count {
			sleepUntil(ctx, sendAt.Add(Spacing))
		}
	}
	return rtts, sent, nil
}

// parseReply decodes an ICMP packet, tolerating the IPv4 header that macOS prepends on udp4 sockets.
func parseReply(b []byte, proto int, want icmp.Type) (*icmp.Echo, bool) {
	try := func(p []byte) (*icmp.Echo, bool) {
		msg, err := icmp.ParseMessage(proto, p)
		if err != nil || msg.Type != want {
			return nil, false
		}
		e, ok := msg.Body.(*icmp.Echo)
		return e, ok
	}
	if e, ok := try(b); ok {
		return e, true
	}
	if proto == protoICMP && len(b) > ipv4.HeaderLen && b[0]>>4 == 4 {
		hl := int(b[0]&0x0f) * 4
		if hl >= ipv4.HeaderLen && len(b) > hl {
			return try(b[hl:])
		}
	}
	return nil, false
}

func sameIP(a net.Addr, ip net.IP) bool {
	switch v := a.(type) {
	case *net.UDPAddr:
		return v.IP.Equal(ip)
	case *net.IPAddr:
		return v.IP.Equal(ip)
	}
	return false
}

// connects performs Count sequential TCP connects to ip:port and returns connect durations.
func connects(ctx context.Context, ip net.IP, port int) ([]time.Duration, int) {
	addr := net.JoinHostPort(ip.String(), itoa(port))
	var rtts []time.Duration
	sent := 0
	d := &net.Dialer{Timeout: EchoTimeout}
	for i := 0; i < Count; i++ {
		if ctx.Err() != nil {
			break
		}
		start := time.Now()
		sent++
		conn, err := d.DialContext(ctx, "tcp", addr)
		if err == nil {
			rtts = append(rtts, time.Since(start))
			conn.Close()
		}
		if i < Count-1 {
			sleepUntil(ctx, start.Add(Spacing))
		}
	}
	return rtts, sent
}

func sleepUntil(ctx context.Context, t time.Time) {
	d := time.Until(t)
	if d <= 0 {
		return
	}
	select {
	case <-ctx.Done():
	case <-time.After(d):
	}
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var b [20]byte
	i := len(b)
	for n > 0 {
		i--
		b[i] = byte('0' + n%10)
		n /= 10
	}
	return string(b[i:])
}

// LinuxHint is logged when ICMP is unavailable on Linux.
func LinuxHint() string {
	if runtime.GOOS == "linux" {
		return "set `sysctl -w net.ipv4.ping_group_range=\"0 2147483647\"` to enable unprivileged ICMP"
	}
	return ""
}