package probe

import (
	"context"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"time"
)

const (
	icmpEchoReply = 0
	icmpEcho      = 8
)

// PingResult contains the outcome of one ICMP echo request.
type PingResult struct {
	Reachable bool
	RTT       time.Duration
	Err       error
}

// Ping sends one ICMP echo request to an IPv4 host.
func Ping(ctx context.Context, host string, timeout time.Duration) (result PingResult) {
	start := time.Now()
	ip, err := resolveIPv4(ctx, host)
	if err != nil {
		return PingResult{Err: err}
	}

	conn, err := net.ListenPacket("ip4:icmp", "0.0.0.0")
	if err != nil {
		return PingResult{Err: fmt.Errorf("opening icmp socket: %w", err)}
	}
	defer func() {
		if closeErr := conn.Close(); closeErr != nil && result.Err == nil {
			result = PingResult{Err: fmt.Errorf("closing icmp socket: %w", closeErr)}
		}
	}()

	deadline := time.Now().Add(timeout)
	if ctxDeadline, ok := ctx.Deadline(); ok && ctxDeadline.Before(deadline) {
		deadline = ctxDeadline
	}
	if err := conn.SetDeadline(deadline); err != nil {
		return PingResult{Err: fmt.Errorf("setting icmp deadline: %w", err)}
	}

	id, seq, err := echoIdentifiers()
	if err != nil {
		return PingResult{Err: err}
	}
	packet := echoPacket(id, seq)
	if _, err := conn.WriteTo(packet, &net.IPAddr{IP: ip}); err != nil {
		return PingResult{Err: fmt.Errorf("sending icmp echo to %s: %w", host, err)}
	}

	buffer := make([]byte, 1500)
	for {
		n, addr, err := conn.ReadFrom(buffer)
		if err != nil {
			if ctxErr := ctx.Err(); ctxErr != nil {
				return PingResult{Err: fmt.Errorf("ping canceled: %w", ctxErr)}
			}
			return PingResult{Err: fmt.Errorf("reading icmp reply from %s: %w", host, err)}
		}
		if !sameIP(addr, ip) {
			continue
		}
		if echoReplyMatches(buffer[:n], id, seq) {
			return PingResult{Reachable: true, RTT: time.Since(start)}
		}
	}
}

func resolveIPv4(ctx context.Context, host string) (net.IP, error) {
	if parsed := net.ParseIP(host); parsed != nil {
		if ipv4 := parsed.To4(); ipv4 != nil {
			return ipv4, nil
		}
		return nil, fmt.Errorf("%s is not an IPv4 address", host)
	}
	addrs, err := net.DefaultResolver.LookupIPAddr(ctx, host)
	if err != nil {
		return nil, fmt.Errorf("resolving ping host %s: %w", host, err)
	}
	for _, addr := range addrs {
		if ipv4 := addr.IP.To4(); ipv4 != nil {
			return ipv4, nil
		}
	}
	return nil, fmt.Errorf("resolving ping host %s: no IPv4 address found", host)
}

func echoIdentifiers() (uint16, uint16, error) {
	var data [4]byte
	if _, err := rand.Read(data[:]); err != nil {
		return 0, 0, fmt.Errorf("generating icmp identifiers: %w", err)
	}
	return binary.BigEndian.Uint16(data[0:2]), binary.BigEndian.Uint16(data[2:4]), nil
}

func echoPacket(id uint16, seq uint16) []byte {
	payload := []byte("statuspanel")
	packet := make([]byte, 8+len(payload))
	packet[0] = icmpEcho
	binary.BigEndian.PutUint16(packet[4:6], id)
	binary.BigEndian.PutUint16(packet[6:8], seq)
	copy(packet[8:], payload)
	checksum := icmpChecksum(packet)
	binary.BigEndian.PutUint16(packet[2:4], checksum)
	return packet
}

func echoReplyMatches(packet []byte, id uint16, seq uint16) bool {
	if len(packet) < 8 {
		return false
	}
	if packet[0] == 69 && len(packet) >= 28 {
		packet = packet[20:]
	}
	return len(packet) >= 8 &&
		packet[0] == icmpEchoReply &&
		binary.BigEndian.Uint16(packet[4:6]) == id &&
		binary.BigEndian.Uint16(packet[6:8]) == seq
}

func sameIP(addr net.Addr, ip net.IP) bool {
	ipAddr, ok := addr.(*net.IPAddr)
	if !ok {
		return false
	}
	return ipAddr.IP.Equal(ip)
}

func icmpChecksum(data []byte) uint16 {
	var sum uint32
	for len(data) > 1 {
		sum += uint32(binary.BigEndian.Uint16(data[:2]))
		data = data[2:]
	}
	if len(data) == 1 {
		sum += uint32(data[0]) << 8
	}
	for sum>>16 != 0 {
		sum = (sum & 0xffff) + (sum >> 16)
	}
	return ^uint16(sum)
}

var errUnsupportedGateway = errors.New("gateway discovery is unsupported on this platform")