// Package traceroute provides functions for executing a tracroute to a remote
// host.
package traceroute

import (
	"errors"
	"fmt"
	"net"
	"os"
	"time"

	"github.com/vishvananda/netlink"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
)

const DEFAULT_PORT = 33434
const DEFAULT_MAX_HOPS = 30
const DEFAULT_FIRST_HOP = 1
const DEFAULT_TIMEOUT_MS = 1000
const DEFAULT_RETRIES = 3
const DEFAULT_PACKET_SIZE = 60

func socketAddr(ip net.IP) (IPAddr [4]byte, err error) {
	routes, err := netlink.RouteGet(ip)

	if err != nil {
		return
	}

	copy(IPAddr[:], routes[0].Src.To4())
	err = errors.New("You do not appear to be connected to the Internet")
	return
}

// Given a host name convert it to a 4 byte IP address.
func destAddr(dest string) (destAddr [4]byte, IPAddr net.IPAddr, 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())
	IPAddr.IP = ipAddr.IP
	return
}

// TracrouteOptions type
type TracerouteOptions struct {
	port       int
	maxHops    int
	firstHop   int
	timeoutMs  int
	retries    int
	packetSize int
}

func (options *TracerouteOptions) Port() int {
	if options.port == 0 {
		options.port = DEFAULT_PORT
	}
	return options.port
}

func (options *TracerouteOptions) SetPort(port int) {
	options.port = port
}

func (options *TracerouteOptions) MaxHops() int {
	if options.maxHops == 0 {
		options.maxHops = DEFAULT_MAX_HOPS
	}
	return options.maxHops
}

func (options *TracerouteOptions) SetMaxHops(maxHops int) {
	options.maxHops = maxHops
}

func (options *TracerouteOptions) FirstHop() int {
	if options.firstHop == 0 {
		options.firstHop = DEFAULT_FIRST_HOP
	}
	return options.firstHop
}

func (options *TracerouteOptions) SetFirstHop(firstHop int) {
	options.firstHop = firstHop
}

func (options *TracerouteOptions) TimeoutMs() int {
	if options.timeoutMs == 0 {
		options.timeoutMs = DEFAULT_TIMEOUT_MS
	}
	return options.timeoutMs
}

func (options *TracerouteOptions) SetTimeoutMs(timeoutMs int) {
	options.timeoutMs = timeoutMs
}

func (options *TracerouteOptions) Retries() int {
	if options.retries == 0 {
		options.retries = DEFAULT_RETRIES
	}
	return options.retries
}

func (options *TracerouteOptions) SetRetries(retries int) {
	options.retries = retries
}

func (options *TracerouteOptions) PacketSize() int {
	if options.packetSize == 0 {
		options.packetSize = DEFAULT_PACKET_SIZE
	}
	return options.packetSize
}

func (options *TracerouteOptions) SetPacketSize(packetSize int) {
	options.packetSize = packetSize
}

// TracerouteHop type
type TracerouteHop struct {
	Success     bool
	Address     [4]byte
	Host        string
	N           int
	ElapsedTime time.Duration
	TTL         int
}

func (hop *TracerouteHop) AddressString() string {
	return fmt.Sprintf("%v.%v.%v.%v", hop.Address[0], hop.Address[1], hop.Address[2], hop.Address[3])
}

func (hop *TracerouteHop) HostOrAddressString() string {
	hostOrAddr := hop.AddressString()
	if hop.Host != "" {
		hostOrAddr = hop.Host
	}
	return hostOrAddr
}

// TracerouteResult type
type TracerouteResult struct {
	DestinationAddress [4]byte
	Hops               []TracerouteHop
}

func notify(hop TracerouteHop, channels []chan TracerouteHop) {
	for _, c := range channels {
		c <- hop
	}
}

func closeNotify(channels []chan TracerouteHop) {
	for _, c := range channels {
		close(c)
	}
}

func Traceroute(dest string, options *TracerouteOptions, c ...chan TracerouteHop) (result TracerouteResult, err error) {
	result.Hops = []TracerouteHop{}
	destAddrBytes, destIPAddr, err := destAddr(dest)
	result.DestinationAddress = destAddrBytes
	if err != nil {
		return
	}

	timeoutDuration, err := time.ParseDuration(fmt.Sprintf("%vs", options.TimeoutMs()))

	if err != nil {
		return
	}

	ttl := options.FirstHop()
	retry := 0

	socketAddr, err := socketAddr(destIPAddr.IP)

	connection, err := net.ListenPacket("ip4:1", fmt.Sprintf("%v.%v.%v.%v", socketAddr[0], socketAddr[1], socketAddr[2], socketAddr[3]))
	if err != nil {
		return
	}

	defer connection.Close()
	packet := ipv4.NewPacketConn(connection)

	if err := packet.SetControlMessage(ipv4.FlagTTL|ipv4.FlagSrc|ipv4.FlagDst|ipv4.FlagInterface, true); err != nil {
		return TracerouteResult{}, err
	}

	icmpMessage := icmp.Message{
		Type: ipv4.ICMPTypeEcho, Code: 0,
		Body: &icmp.Echo{
			ID:   os.Getpid() & 0xffff,
			Data: []byte("HELLO-R-U-THERE"),
		},
	}

	rb := make([]byte, 1500)

	for {
		icmpMessage.Body.(*icmp.Echo).Seq = ttl

		wb, err := icmpMessage.Marshal(nil)

		if err != nil {
			return TracerouteResult{}, err
		}

		if err := packet.SetTTL(ttl); err != nil {
			return TracerouteResult{}, err
		}

		start := time.Now()
		if _, err := packet.WriteTo(wb, nil, &destIPAddr); err != nil {
			return TracerouteResult{}, err
		}

		if err := packet.SetReadDeadline(start.Add(timeoutDuration)); err != nil {
			return TracerouteResult{}, err
		}

		var n int
		var srcAddr net.Addr

		n, _, srcAddr, err = packet.ReadFrom(rb)

		if err != nil {
			if err, ok := err.(net.Error); ok && err.Timeout() {
				// means timeout here
				notify(TracerouteHop{Success: false, TTL: ttl}, c)
				retry += 1
				if retry > options.Retries() {
					ttl += 1
					retry = 0
				}
				continue
			}
			return TracerouteResult{}, err
		}

		elapsed := time.Since(start)
		rm, err := icmp.ParseMessage(1, rb[:n])

		if err != nil {
			return TracerouteResult{}, err
		}

		srcAddrBytes, _, _ := destAddr(srcAddr.String())

		hop := TracerouteHop{
			Success:     true,
			Address:     srcAddrBytes,
			N:           n,
			ElapsedTime: elapsed,
			TTL:         ttl,
		}

		currHost, err := net.LookupAddr(srcAddr.String())
		if err == nil {
			hop.Host = currHost[0]
		}

		if rm.Type == ipv4.ICMPTypeEchoReply || rm.Type == ipv4.ICMPTypeTimeExceeded {
			notify(hop, c)
			result.Hops = append(result.Hops, hop)
		}

		ttl += 1
		retry = 0

		if ttl > options.MaxHops() || rm.Type == ipv4.ICMPTypeEchoReply {
			closeNotify(c)
			return result, nil
		}

		time.Sleep(time.Millisecond * 500)
	}
}