package main

import (
	"context"
	"crypto/tls"
	"encoding/binary"
	"flag"
	"fmt"
	"log"
	"strings"
	"time"

	quic "github.com/AeonDave/mp-quic-go"
	"mpquic-exp/client/scheduler"
)

func main() {
	schedName := flag.String("scheduler", "roundrobin", "Scheduler to use")
	listScheds := flag.Bool("list-schedulers", false, "List available schedulers")
	addr := flag.String("addr", "127.0.0.1:4242", "Server address")
	duration := flag.Int("duration", 10, "Duration in seconds")
	payloadSize := flag.Int("payload-size", 1024, "Size of the payload in bytes (min 20)")
	flag.Parse()

	if *listScheds {
		fmt.Println("Available schedulers:")
		for _, name := range scheduler.List() {
			fmt.Println(" -", name)
		}
		return
	}

	if *payloadSize < 20 {
		*payloadSize = 20 // 4 bytes length, 8 bytes seq, 8 bytes timestamp
	}

	tlsConf := &tls.Config{
		InsecureSkipVerify: true,
		NextProtos:         []string{"mpquic-exp"},
	}

	sched, err := scheduler.Get(*schedName)
	if err != nil {
		log.Fatalf("Error: %v. Available: %s", err, strings.Join(scheduler.List(), ", "))
	}

	quicConfig := &quic.Config{
		MaxPaths:            4,
		MultipathController: quic.NewDefaultMultipathController(sched),
	}

	fmt.Printf("Connecting to %s using %s scheduler...\n", *addr, *schedName)
	conn, err := quic.DialAddr(context.Background(), *addr, tlsConf, quicConfig)
	if err != nil {
		log.Fatal(err)
	}

	stream, err := conn.OpenStreamSync(context.Background())
	if err != nil {
		log.Fatal(err)
	}

	fmt.Printf("Stream opened, sending data for %d seconds (Payload Size: %d bytes)...\n", *duration, *payloadSize)

	end := time.Now().Add(time.Duration(*duration) * time.Second)
	payload := make([]byte, *payloadSize)
	
	// Frame structure:
	// [0:4] uint32 Total Length
	// [4:12] uint64 Sequence Number
	// [12:20] uint64 Send Timestamp (nanoseconds)
	// [20:] Padding (dummy data)
	binary.BigEndian.PutUint32(payload[0:4], uint32(*payloadSize))

	var seqNum uint64 = 0
	totalBytes := 0
	
	// Target sending rate: we don't want to lock the CPU entirely in a busy loop.
	// We yield slightly to allow the network stack to process.
	// But to measure max throughput we just send as fast as stream.Write allows.
	for time.Now().Before(end) {
		seqNum++
		sendTime := uint64(time.Now().UnixNano())
		
		binary.BigEndian.PutUint64(payload[4:12], seqNum)
		binary.BigEndian.PutUint64(payload[12:20], sendTime)
		
		n, err := stream.Write(payload)
		if err != nil {
			log.Fatal("Stream write error:", err)
		}
		totalBytes += n
	}
	
	fmt.Printf("Finished sending %d packets, %d bytes (%.2f MB)\n", seqNum, totalBytes, float64(totalBytes)/1024/1024)
}