// Package ping is an ICMP ping library seeking to emulate the unix "ping"
// command.
//
// Here is a very simple example that sends & receives 3 packets:
//
//	pinger, err := ping.NewPinger("www.google.com")
//	if err != nil {
//		panic(err)
//	}
//
//	pinger.Count = 3
//	pinger.Run() // blocks until finished
//	stats := pinger.Statistics() // get send/receive/rtt stats
//
// Here is an example that emulates the unix ping command:
//
//	pinger, err := ping.NewPinger("www.google.com")
//	if err != nil {
//		fmt.Printf("ERROR: %s\n", err.Error())
//		return
//	}
//
//	pinger.OnRecv = func(pkt *ping.Packet) {
//		fmt.Printf("%d bytes from %s: icmp_seq=%d time=%v\n",
//			pkt.Nbytes, pkt.IPAddr, pkt.Seq, pkt.Rtt)
//	}
//	pinger.OnFinish = func(stats *ping.Statistics) {
//		fmt.Printf("\n--- %s ping statistics ---\n", stats.Addr)
//		fmt.Printf("%d packets transmitted, %d packets received, %v%% packet loss\n",
//			stats.PacketsSent, stats.PacketsRecv, stats.PacketLoss)
//		fmt.Printf("round-trip min/avg/max/stddev = %v/%v/%v/%v\n",
//			stats.MinRtt, stats.AvgRtt, stats.MaxRtt, stats.StdDevRtt)
//	}
//
//	fmt.Printf("PING %s (%s):\n", pinger.Addr(), pinger.IPAddr())
//	pinger.Run()
//
// It sends ICMP packet(s) and waits for a response. If it receives a response,
// it calls the "receive" callback. When it's finished, it calls the "finish"
// callback.
//
// For a full ping example, see "cmd/ping/ping.go".
//
package ping

import (
	"fmt"
	"math/rand"
	"net"
	"sync"
	"syscall"
	"time"

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

var (
	myid = rand.Intn(65535)
)

// NewMultiPinger returns a new MultiPinger struct pointer
func NewMultiPinger(addr []string) (*MultiPinger, error) {
	var ipaddr []*net.IPAddr
	for _, v := range addr {
		if tmpipaddr, err := net.ResolveIPAddr("ip", v); err == nil {
			ipaddr = append(ipaddr, tmpipaddr)
		} else {
			return nil, err
		}
	}

	return &MultiPinger{
		ipaddr:   ipaddr,
		addr:     addr,
		Interval: time.Second,
		Timeout:  time.Second * 100000,
		Count:    -1,

		network: "udp",
		ipv4:    true,
		size:    timeSliceLength,

		done: make(chan bool),
	}, nil
}

// MultiPinger represents ICMP packet sender/receiver
type MultiPinger struct {
	// Interval is the wait time between each packet send. Default is 1s.
	Interval time.Duration

	// Timeout specifies a timeout before ping exits, regardless of how many
	// packets have been received.
	Timeout time.Duration

	// Count tells MultiMultiPinger to stop after sending (and receiving) Count echo
	// packets. If this option is not specified, MultiPinger will operate until
	// interrupted.
	Count int

	// Debug runs in debug mode
	Debug bool

	// Number of packets sent
	PacketsSent int

	// Number of packets received
	PacketsRecv int

	// OnRecv is called when MultiPinger receives and processes a packet
	OnRecv func(*Packet)

	// OnFinish is called when MultiPinger exits
	OnFinish func()

	// stop chan bool
	done chan bool

	ipaddr []*net.IPAddr
	addr   []string

	ipv4     bool
	source   string
	size     int
	sequence int
	network  string
}

// SetIPAddr sets the ip address of the target host.
func (p *MultiPinger) AddIPAddr(ipaddr *net.IPAddr) {
	p.ipaddr = append(p.ipaddr, ipaddr)
	//p.addr = append(p.addr, ipaddr.String())
}

// IPAddr returns the ip address of the target host.
func (p *MultiPinger) IPAddr() []*net.IPAddr {
	return p.ipaddr
}

// SetAddr resolves and sets the ip address of the target host, addr can be a
// DNS name like "www.google.com" or IP like "127.0.0.1".
func (p *MultiPinger) SetAddr(addr string) error {
	ipaddr, err := net.ResolveIPAddr("ip", addr)
	if err != nil {
		return err
	}

	p.AddIPAddr(ipaddr)
	p.addr = append(p.addr, addr)

	return nil
}

//Replace addrs
func (p *MultiPinger) ReplaceAddrs(addr []string) error {
	var ipaddrs []*net.IPAddr
	for _, v := range addr {
		if ipaddr, err := net.ResolveIPAddr("ip", v); err != nil {
			return err
		} else {
			ipaddrs = append(ipaddrs, ipaddr)
		}
	}

	p.ipaddr = ipaddrs
	p.addr = addr

	return nil
}

// Addr returns the string ip address of the target host.
func (p *MultiPinger) Addr() []string {
	return p.addr
}

// SetPrivileged sets the type of ping MultiPinger will send.
// false means MultiPinger will send an "unprivileged" UDP ping.
// true means MultiPinger will send a "privileged" raw ICMP ping.
// NOTE: setting to true requires that it be run with super-user privileges.
func (p *MultiPinger) SetPrivileged(privileged bool) {
	if privileged {
		p.network = "ip"
	} else {
		p.network = "udp"
	}
}

// Privileged returns whether MultiPinger is running in privileged mode.
func (p *MultiPinger) Privileged() bool {
	return p.network == "ip"
}

// Run runs the MultiPinger. This is a blocking function that will exit when it's
// done. If Count or Interval are not specified, it will run continuously until
// it is interrupted.
func (p *MultiPinger) Run() {
	p.run()
}

func (p *MultiPinger) sendPingPack(conn *icmp.PacketConn) {
	for _, v := range p.ipaddr {
		err := p.sendICMP(conn, v)
		if err != nil {
			fmt.Println(err.Error())
		}
	}
}

func (p *MultiPinger) run() {
	var conn *icmp.PacketConn
	if p.ipv4 {
		if conn = p.listen(ipv4Proto[p.network], p.source); conn == nil {
			return
		}
	} else {
		if conn = p.listen(ipv6Proto[p.network], p.source); conn == nil {
			return
		}
	}
	defer conn.Close()
	defer p.finish()

	var wg sync.WaitGroup
	recv := make(chan *packet, 5)
	wg.Add(1)
	go p.recvICMP(conn, recv, &wg)

	p.sendPingPack(conn)

	timeout := time.NewTicker(p.Timeout)
	interval := time.NewTicker(p.Interval)

	for {
		select {
		case <-p.done:
			wg.Wait()
			return
		case <-timeout.C:
			close(p.done)
			wg.Wait()
			return
		case <-interval.C:
			p.sendPingPack(conn)
		case r := <-recv:
			err := p.processPacket(r)
			if err != nil {
				fmt.Println("FATAL: ", err.Error())
			}
		default:
			if p.Count > 0 && p.PacketsRecv >= p.Count {
				close(p.done)
				wg.Wait()
				return
			}
		}
	}
}

func (p *MultiPinger) finish() {
	handler := p.OnFinish
	if handler != nil {
		handler()
	}
}

func (p *MultiPinger) recvICMP(
	conn *icmp.PacketConn,
	recv chan<- *packet,
	wg *sync.WaitGroup,
) {
	defer wg.Done()
	for {
		select {
		case <-p.done:
			return
		default:
			bytes := make([]byte, 512)
			conn.SetReadDeadline(time.Now().Add(time.Millisecond * 100))
			n, addr, err := conn.ReadFrom(bytes)
			if err != nil {
				if neterr, ok := err.(*net.OpError); ok {
					if neterr.Timeout() {
						// Read timeout
						continue
					} else {
						close(p.done)
						return
					}
				}
			}

			recv <- &packet{bytes: bytes, nbytes: n, addr: addr}
		}
	}
}

func (p *MultiPinger) processPacket(recv *packet) error {
	var bytes []byte
	var proto int
	if p.ipv4 {
		if p.network == "ip" {
			bytes = ipv4Payload(recv.bytes)
		} else {
			bytes = recv.bytes
		}
		proto = protocolICMP
	} else {
		bytes = recv.bytes
		proto = protocolIPv6ICMP
	}

	var m *icmp.Message
	var err error
	if m, err = icmp.ParseMessage(proto, bytes[:recv.nbytes]); err != nil {
		return fmt.Errorf("Error parsing icmp message")
	}

	if m.Type != ipv4.ICMPTypeEchoReply && m.Type != ipv6.ICMPTypeEchoReply {
		// Not an echo reply, ignore it
		return nil
	}

	//fmt.Printf("Recv pack: %+v %+v\n", m, m.Body)

	outPkt := &Packet{
		Nbytes: recv.nbytes,
		IPAddr: fmt.Sprintln(recv.addr),
	}

	switch pkt := m.Body.(type) {
	case *icmp.Echo:
		if len(pkt.Data) >= 2 {
			if pkt.Data[len(pkt.Data)-2] == magicA && pkt.Data[len(pkt.Data)-1] == magicB {
				outPkt.Rtt = time.Since(bytesToTime(pkt.Data[:timeSliceLength]))
				outPkt.Seq = pkt.Seq
				outPkt.Id = pkt.ID
				p.PacketsRecv += 1
			} else {
				return errors.New("Magic not match!!")
			}
		} else {
			return errors.New("Recv pkt len not match!!")
		}
	default:
		// Very bad, not sure how this can happen
		return fmt.Errorf("Error, invalid ICMP echo reply. Body type: %T, %s",
			pkt, pkt)
	}

	handler := p.OnRecv
	if handler != nil {
		handler(outPkt)
	}

	return nil
}

func (p *MultiPinger) sendICMP(conn *icmp.PacketConn, ipaddr *net.IPAddr) error {
	var typ icmp.Type
	if p.ipv4 {
		typ = ipv4.ICMPTypeEcho
	} else {
		typ = ipv6.ICMPTypeEchoRequest
	}

	var dst net.Addr = ipaddr
	if p.network == "udp" {
		dst = &net.UDPAddr{IP: ipaddr.IP, Zone: ipaddr.Zone}
	}

	t := timeToBytes(time.Now())
	if p.size-timeSliceLength != 0 {
		t = append(t, byteSliceOfSize(p.size-timeSliceLength)...)
	}
	t = append(t, magicA, magicB)
	bytes, err := (&icmp.Message{
		Type: typ, Code: 0,
		Body: &icmp.Echo{
			ID:   myid,
			Seq:  p.sequence,
			Data: t,
		},
	}).Marshal(nil)

	if err != nil {
		return err
	}

	for {
		if _, err := conn.WriteTo(bytes, dst); err != nil {
			if neterr, ok := err.(*net.OpError); ok {
				if neterr.Err == syscall.ENOBUFS {
					continue
				}
			}
		}
		p.PacketsSent += 1
		p.sequence += 1
		break
	}
	return nil
}

func (p *MultiPinger) listen(netProto string, source string) *icmp.PacketConn {
	conn, err := icmp.ListenPacket(netProto, source)
	if err != nil {
		fmt.Printf("Error listening for ICMP packets: %s\n", err.Error())
		close(p.done)
		return nil
	}
	return conn
}
