package traceroute

import (
	"fmt"
	"net"
	"syscall"
	"time"
)

const PORT = 33434
const MAX_HOPS = 30
const TIMEOUT_MS = 500
const RETRIES = 3
const PACKET_SIZE = 52

type Hop struct {
	Success     bool
	Address     [4]byte
	ElapsedTime time.Duration
	TTL         int
}

func Traceroute(dest string, c chan<- Hop) error {

	ttl := 1
	retry := 0
	tv := syscall.NsecToTimeval(1000 * 1000 * TIMEOUT_MS)

	for {
		start := time.Now()

		recvSock, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_RAW, syscall.IPPROTO_ICMP)
		if err != nil {
			return err
		}

		sendSock, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_DGRAM, syscall.IPPROTO_UDP)
		if err != nil {
			return err
		}
		syscall.SetsockoptInt(sendSock, 0x0, syscall.IP_TTL, ttl)
		syscall.SetsockoptTimeval(recvSock, syscall.SOL_SOCKET, syscall.SO_RCVTIMEO, &tv)

		defer syscall.Close(recvSock)
		defer syscall.Close(sendSock)

		sockaddr, err := socketAddr()
		if err != nil {
			return err
		}
		syscall.Bind(recvSock, &syscall.SockaddrInet4{Port: PORT, Addr: sockaddr})

		destaddr, err := destAddr(dest)
		if err != nil {
			return err
		}
		syscall.Sendto(sendSock, []byte{0x0}, 0, &syscall.SockaddrInet4{Port: PORT, Addr: destaddr})

		var p = make([]byte, PACKET_SIZE)
		_, from, err := syscall.Recvfrom(recvSock, p, 0)
		elapsed := time.Since(start)

		if err == nil {
			currAddr := from.(*syscall.SockaddrInet4).Addr

			hop := Hop{Success: true, Address: currAddr, ElapsedTime: elapsed, TTL: ttl}
			c <- hop

			ttl += 1
			retry = 0

			if ttl > MAX_HOPS || currAddr == destaddr {
				close(c)
				return nil
			}

		} else {
			retry += 1
			if retry > RETRIES {
				hop := Hop{Success: false, TTL: ttl}
				c <- hop

				ttl += 1
				retry = 0
			}
			if ttl > MAX_HOPS {
				close(c)
				return nil
			}
		}

	}
}

func socketAddr() (addr [4]byte, err error) {
	addrs, err := net.InterfaceAddrs()
	if err != nil {
		return
	}

	for _, a := range addrs {
		if ipnet, ok := a.(*net.IPNet); ok && !ipnet.IP.IsLoopback() {
			if len(ipnet.IP.To4()) == net.IPv4len {
				copy(addr[:], ipnet.IP.To4())
				return
			}
		}
	}
	return
}

func destAddr(dest string) (destAddr [4]byte, err error) {
	addrs, err := net.LookupHost(dest)
	if err != nil {
		return
	}
	addr := addrs[0]

	ipAddr, err := net.ResolveIPAddr("ip", addr)
	if err != nil {
		return
	}
	copy(destAddr[:], ipAddr.IP.To4())
	return
}

func AddressToString(addr [4]byte) string {
	return fmt.Sprintf("%v.%v.%v.%v", addr[0], addr[1], addr[2], addr[3])
}
