package scan

import (
	"fmt"
	"net"
	"sync/atomic"
	"time"

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

type Pinger struct {
	conn *icmp.PacketConn
	seq  atomic.Uint32
}

func NewPinger() (*Pinger, error) {
	conn, err := icmp.ListenPacket("udp4", "0.0.0.0")
	if err != nil {
		return nil, fmt.Errorf("icmp listen: %w", err)
	}
	return &Pinger{conn: conn}, nil
}

func (p *Pinger) Close() {
	if p.conn != nil {
		p.conn.Close()
	}
}

func (p *Pinger) Ping(ip string, timeout time.Duration) bool {
	dst, err := net.ResolveIPAddr("ip4", ip)
	if err != nil {
		return false
	}

	var target net.Addr
	switch p.conn.LocalAddr().(type) {
	case *net.UDPAddr:
		target = &net.UDPAddr{IP: dst.IP}
	case *net.IPAddr:
		target = &net.IPAddr{IP: dst.IP}
	default:
		target = dst
	}

	seq := uint16(p.seq.Add(1))
	msg := icmp.Message{
		Type: ipv4.ICMPTypeEcho,
		Code: 0,
		Body: &icmp.Echo{
			ID:   0,
			Seq:  int(seq),
			Data: []byte("PING"),
		},
	}
	buf, err := msg.Marshal(nil)
	if err != nil {
		return false
	}

	if _, err := p.conn.WriteTo(buf, target); err != nil {
		return false
	}

	deadline := time.Now().Add(timeout)
	rb := make([]byte, 1500)
	for {
		remaining := time.Until(deadline)
		if remaining <= 0 {
			return false
		}
		_ = p.conn.SetReadDeadline(time.Now().Add(remaining))
		n, peer, err := p.conn.ReadFrom(rb)
		if err != nil {
			return false
		}
		rm, err := icmp.ParseMessage(ipv4.ICMPTypeEchoReply.Protocol(), rb[:n])
		if err != nil {
			continue
		}
		if rm.Type != ipv4.ICMPTypeEchoReply {
			continue
		}
		var peerIP net.IP
		switch a := peer.(type) {
		case *net.UDPAddr:
			peerIP = a.IP
		case *net.IPAddr:
			peerIP = a.IP
		default:
			peerIP = net.ParseIP(peer.String())
		}
		if peerIP != nil && peerIP.Equal(dst.IP) {
			return true
		}
	}
}

func Ping(ip string, timeout time.Duration) (bool, error) {
	p, err := NewPinger()
	if err != nil {
		return false, err
	}
	defer p.Close()
	return p.Ping(ip, timeout), nil
}