package main

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"encoding/csv"
	"encoding/pem"
	"flag"
	"fmt"
	"io"
	"log"
	"math/big"
	"os"
	"strconv"
	"sync"
	"time"

	quic "github.com/AeonDave/mp-quic-go"
)

func generateTLSConfig() *tls.Config {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		panic(err)
	}
	template := x509.Certificate{SerialNumber: big.NewInt(1)}
	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
	if err != nil {
		panic(err)
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(key)})
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})

	tlsCert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		panic(err)
	}
	return &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
		NextProtos:   []string{"mpquic-exp"},
	}
}

func main() {
	addr := flag.String("addr", "0.0.0.0:4242", "Address to listen on")
	maxPaths := flag.Int("maxpaths", 4, "Maximum number of paths")
	outCSV := flag.String("output", "app_metrics.csv", "Output CSV for application-level metrics")
	flag.Parse()

	listener, err := quic.ListenAddr(*addr, generateTLSConfig(), &quic.Config{
		MaxPaths: *maxPaths,
	})
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("Server listening on %s with max %d paths\n", *addr, *maxPaths)
	fmt.Printf("Logging app-level metrics to %s\n", *outCSV)

	// Open CSV
	f, err := os.OpenFile(*outCSV, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
	if err != nil {
		log.Fatal("Failed to open CSV:", err)
	}
	defer f.Close()

	writer := csv.NewWriter(f)
	writer.Write([]string{"message_id", "chunk_index", "total_chunks", "drone_send_time", "ground_recv_time", "latency_ns", "bytes_received"})
	writer.Flush()

	// Mutex to protect CSV writer if multiple streams are used
	var mu sync.Mutex

	for {
		conn, err := listener.Accept(context.Background())
		if err != nil {
			log.Fatal(err)
		}
		fmt.Println("Client connected:", conn.RemoteAddr())

		go func(conn *quic.Conn) {
			for {
				stream, err := conn.AcceptStream(context.Background())
				if err != nil {
					fmt.Println("Session error:", err)
					return
				}
				go func(stream *quic.Stream) {
					header := make([]byte, 4)
					for {
						// Read 4-byte length prefix
						_, err := io.ReadFull(stream, header)
						if err != nil {
							fmt.Printf("Stream closed\n")
							return
						}
						
						length := binary.BigEndian.Uint32(header)
						if length < 36 {
							fmt.Printf("Warning: Invalid frame length %d\n", length)
							continue
						}
						
						// Read rest of the payload
						payload := make([]byte, length-4)
						_, err = io.ReadFull(stream, payload)
						if err != nil {
							fmt.Printf("Stream read error: %v\n", err)
							return
						}
						
						recvTime := time.Now().UnixNano()
						
						msgId := binary.BigEndian.Uint64(payload[0:8])
						sendTime := binary.BigEndian.Uint64(payload[8:16])
						chunkIdx := binary.BigEndian.Uint64(payload[16:24])
						totalChunks := binary.BigEndian.Uint64(payload[24:32])
						latency := recvTime - int64(sendTime)
						
						mu.Lock()
						writer.Write([]string{
							strconv.FormatUint(msgId, 10),
							strconv.FormatUint(chunkIdx, 10),
							strconv.FormatUint(totalChunks, 10),
							strconv.FormatUint(sendTime, 10),
							strconv.FormatInt(recvTime, 10),
							strconv.FormatInt(latency, 10),
							strconv.FormatUint(uint64(length), 10),
						})
						// Periodically flush? 
						if msgId % 1000 == 0 {
							writer.Flush()
						}
						mu.Unlock()
					}
				}(stream)
			}
		}(conn)
	}
}