package icmp

import (
	"bytes"
	"context"
	cryptorand "crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"math/rand/v2"
	"net"
	"net/netip"
	"strings"

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

var (
	ErrICMPBodyUnsupported  = errors.New("ICMP body type is not supported")
	ErrICMPEchoDataMismatch = errors.New("ICMP data mismatch")
)

type Echoer struct {
	buffer       []byte
	randomSource io.Reader
	logger       Logger
}

func NewEchoer(logger Logger) *Echoer {
	const maxICMPEchoSize = 1500
	buffer := make([]byte, maxICMPEchoSize)
	var seed [32]byte
	_, _ = cryptorand.Read(seed[:])
	randomSource := rand.NewChaCha8(seed)
	return &Echoer{
		buffer:       buffer,
		randomSource: randomSource,
		logger:       logger,
	}
}

var (
	ErrTimedOut     = errors.New("timed out waiting for ICMP echo reply")
	ErrNotPermitted = errors.New("not permitted")
)

func (i *Echoer) Echo(ctx context.Context, ip netip.Addr) (err error) {
	var ipVersion string
	var conn net.PacketConn
	if ip.Is4() {
		ipVersion = "v4"
		conn, err = listenICMPv4(ctx)
	} else {
		ipVersion = "v6"
		conn, err = listenICMPv6(ctx)
	}
	if err != nil {
		if strings.HasSuffix(err.Error(), "socket: operation not permitted") {
			err = fmt.Errorf("%w: you can try adding NET_RAW capability to resolve this", ErrNotPermitted)
		}
		return fmt.Errorf("listening for ICMP packets: %w", err)
	}

	go func() {
		<-ctx.Done()
		conn.Close()
	}()

	const echoDataSize = 32
	id, message := buildMessageToSend(ipVersion, echoDataSize, i.randomSource)

	encodedMessage, err := message.Marshal(nil)
	if err != nil {
		return fmt.Errorf("encoding ICMP message: %w", err)
	}

	_, err = conn.WriteTo(encodedMessage, &net.IPAddr{IP: ip.AsSlice()})
	if err != nil {
		if strings.HasSuffix(err.Error(), "sendto: operation not permitted") {
			err = fmt.Errorf("%w", ErrNotPermitted)
		}
		return fmt.Errorf("writing ICMP message to %s: %w", ip, err)
	}

	receivedData, err := receiveEchoReply(conn, id, i.buffer, ipVersion, i.logger)
	if err != nil {
		if errors.Is(err, net.ErrClosed) && ctx.Err() != nil {
			return fmt.Errorf("%w from %s", ErrTimedOut, ip)
		}
		return fmt.Errorf("receiving ICMP echo reply from %s: %w", ip, err)
	}

	sentData := message.Body.(*icmp.Echo).Data //nolint:forcetypeassert
	if !bytes.Equal(receivedData, sentData) {
		return fmt.Errorf("%w: sent %x to %s and received %x", ErrICMPEchoDataMismatch, sentData, ip, receivedData)
	}

	return nil
}

func buildMessageToSend(ipVersion string, size uint, randomSource io.Reader) (id int, message *icmp.Message) {
	const uint16Bytes = 2
	idBytes := make([]byte, uint16Bytes)
	_, _ = randomSource.Read(idBytes)
	id = int(binary.BigEndian.Uint16(idBytes))

	var icmpType icmp.Type
	switch ipVersion {
	case "v4":
		icmpType = ipv4.ICMPTypeEcho
	case "v6":
		icmpType = ipv6.ICMPTypeEchoRequest
	default:
		panic(fmt.Sprintf("IP version %q not supported", ipVersion))
	}
	messageBodyData := make([]byte, size)
	_, _ = randomSource.Read(messageBodyData)

	// See https://www.iana.org/assignments/icmp-parameters/icmp-parameters.xhtml#icmp-parameters-types
	message = &icmp.Message{
		Type:     icmpType, // echo request
		Code:     0,        // no code
		Checksum: 0,        // calculated at encoding (ipv4) or sending (ipv6)
		Body: &icmp.Echo{
			ID:   id,
			Seq:  0, // only one packet
			Data: messageBodyData,
		},
	}
	return id, message
}

func receiveEchoReply(conn net.PacketConn, id int, buffer []byte, ipVersion string, logger Logger,
) (data []byte, err error) {
	var icmpProtocol int
	const (
		icmpv4Protocol = 1
		icmpv6Protocol = 58
	)
	switch ipVersion {
	case "v4":
		icmpProtocol = icmpv4Protocol
	case "v6":
		icmpProtocol = icmpv6Protocol
	default:
		panic(fmt.Sprintf("unknown IP version: %s", ipVersion))
	}

	for {
		// Note we need to read the whole packet in one call to ReadFrom, so the buffer
		// must be large enough to read the entire reply packet. See:
		// https://groups.google.com/g/golang-nuts/c/5dy2Q4nPs08/m/KmuSQAGEtG4J
		bytesRead, returnAddr, err := conn.ReadFrom(buffer)
		if err != nil {
			return nil, fmt.Errorf("reading from ICMP connection: %w", err)
		}
		packetBytes := buffer[:bytesRead]

		// Parse the ICMP message
		message, err := icmp.ParseMessage(icmpProtocol, packetBytes)
		if err != nil {
			return nil, fmt.Errorf("parsing message: %w", err)
		}

		switch body := message.Body.(type) {
		case *icmp.Echo:
			if id != body.ID {
				logger.Warnf("ignoring ICMP echo reply mismatching expected id %d "+
					"(id: %d, type: %d, code: %d, length: %d, return address %s)",
					id, body.ID, message.Type, message.Code, len(packetBytes), returnAddr)
				continue // not the ID we are looking for
			}
			return body.Data, nil
		case *icmp.DstUnreach:
			logger.Debugf("ignoring ICMP destination unreachable message (type: 3, code: %d, return address %s, expected-id %d)",
				message.Code, returnAddr, id)
			// See https://github.com/qdm12/gluetun/pull/2923#issuecomment-3377532249
			// on why we ignore this message. If it is actually unreachable, the timeout on waiting for
			// the echo reply will do instead of returning an error error.
			continue
		case *icmp.TimeExceeded:
			logger.Debugf("ignoring ICMP time exceeded message (type: 11, code: %d, return address %s, expected-id %d)",
				message.Code, returnAddr, id)
			continue
		default:
			return nil, fmt.Errorf("%w: %T (type %d, code %d, return address %s, expected-id %d)",
				ErrICMPBodyUnsupported, body, message.Type, message.Code, returnAddr, id)
		}
	}
}