package main

import (
	"encoding/binary"
	"fmt"
	"gopkg.in/alecthomas/kingpin.v2"
	"log"
	"net"
	"os"
	"time"
)

var (
	app = kingpin.New("tcpdelay", "Measure TCP delays")

	server        = app.Command("server", "Run as server")
	listenAddress = server.Arg("address", "Address:port to listen on").String()

	client        = app.Command("client", "Request measurement")
	serverAddress = client.Arg("address", "Address:port to connect to").String()
	relativeTimes = client.Flag("relative", "Print times relative to start").Bool()
)

func tcpdelayServer() {
	ln, err := net.Listen("tcp", *listenAddress)
	if err != nil {
		log.Fatal("Cannot listen on ", *listenAddress, err)
	}
	log.Println("Listening on ", ln.Addr())
	for {
		conn, err := ln.Accept()
		log.Println("accepted connection")
		if err != nil {
			log.Fatal("Cannot accept on ", listenAddress, err)
		}
		o := 0
		t0 := time.Now()
		d, _ := time.ParseDuration("1s")
		te := t0.Add(d)
		buffer := make([]byte, 1024)
		for t := time.Now(); t.Before(te); t = time.Now() {
			copy(buffer[0:8], "tcpdelay")
			binary.BigEndian.PutUint64(buffer[8:], uint64(o))
			binary.BigEndian.PutUint64(buffer[16:], uint64(t.UnixNano()))
			n, err := conn.Write(buffer)
			if err != nil {
				log.Fatal("Cannot write on ", listenAddress, err)
			}
			if n != len(buffer) {
				log.Fatal("Wrote ", n, " bytes intead of ", len(buffer))
			}
			o += n
		}
		conn.Close()
		log.Println("closed connection")
	}
}

func tcpdelayClient() {
	t0 := time.Now().UnixNano()
	conn, err := net.Dial("tcp", *serverAddress)
	if err != nil {
		log.Fatal("Cannot connect to ", *serverAddress, err)
	}
	o := 0
	buffer := make([]byte, 1024)
	for {
		n := 0
		for n < len(buffer) {
			nr, err := conn.Read(buffer[n:])
			if err != nil {
				log.Fatal("Cannot read from ", *serverAddress, err.Error())
			}
			if nr <= 0 {
				log.Fatal("Read ", n, " bytes intead of ", len(buffer))
			}
			n += nr
		}
		tr := time.Now().UnixNano()
		if string(buffer[0:8]) != "tcpdelay" {
			log.Fatal("Unexpected magic ", buffer[0:8])
		}
		oSent := int(binary.BigEndian.Uint64(buffer[8:]))
		if oSent != o {
			log.Fatal("Unexpected offset ", oSent, " instead of ", o)
		}
		ts := binary.BigEndian.Uint64(buffer[16:])
		if *relativeTimes {
			fmt.Println(o, int64(ts)-t0, tr-t0)
		} else {
			fmt.Println(o, ts, tr)
		}
		o += n
	}

}

func main() {
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {
	case server.FullCommand():
		tcpdelayServer()
	case client.FullCommand():
		tcpdelayClient()
	}
}