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")
	messageSize := flag.Int("message-size", 1024, "Size of the full application message in bytes (min 36)")
	chunkSize := flag.Int("chunk-size", 0, "Chunk size. If > 0, splits message-size into multiple chunks (min 36)")
	flag.Parse()

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

	if *messageSize < 36 {
		*messageSize = 36 // 4 len + 8 seq + 8 ts + 8 chunk_idx + 8 tot_chunks
	}

	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 (Message Size: %d bytes, Chunk Size: %d)...\n", *duration, *messageSize, *chunkSize)

	end := time.Now().Add(time.Duration(*duration) * time.Second)
	
	var msgId uint64 = 0
	totalBytes := 0
	totalChunksSent := 0
	
	actualChunkSize := *messageSize
	if *chunkSize > 36 && *chunkSize <= *messageSize {
		actualChunkSize = *chunkSize
	}

	totalChunks := uint64((*messageSize + actualChunkSize - 1) / actualChunkSize)
	
	for time.Now().Before(end) {
		msgId++
		sendTime := uint64(time.Now().UnixNano())
		
		bytesLeft := *messageSize
		
		for chunkIdx := uint64(0); chunkIdx < totalChunks; chunkIdx++ {
			currentLength := actualChunkSize
			if bytesLeft < actualChunkSize {
				currentLength = bytesLeft
			}
			if currentLength < 36 {
				currentLength = 36
			}
			
			payload := make([]byte, currentLength)
			binary.BigEndian.PutUint32(payload[0:4], uint32(currentLength))
			binary.BigEndian.PutUint64(payload[4:12], msgId)
			binary.BigEndian.PutUint64(payload[12:20], sendTime)
			binary.BigEndian.PutUint64(payload[20:28], chunkIdx)
			binary.BigEndian.PutUint64(payload[28:36], totalChunks)

			n, err := stream.Write(payload)
			if err != nil {
				log.Fatal("Stream write error:", err)
			}
			totalBytes += n
			bytesLeft -= currentLength
			totalChunksSent++
		}
	}
	
	fmt.Printf("Finished sending %d messages (%d chunks), %d bytes (%.2f MB)\n", msgId, totalChunksSent, totalBytes, float64(totalBytes)/1024/1024)
}