package main

import (
	"fmt"
	"io"
	"math/rand"
	"net/http"
	"strconv"
	"time"
)

const (
	randomPoolSize = 4 * 1024 * 1024 // 4MB pool of random data, read-only after init
	uploadBufSize  = 32 * 1024        // per-request local buffer
)

var randomPool []byte

func init() {
	rng := rand.New(rand.NewSource(time.Now().UnixNano()))
	randomPool = make([]byte, randomPoolSize)
	rng.Read(randomPool)
}

func main() {
	http.HandleFunc("/download", downloadHandler)
	http.HandleFunc("/upload", uploadHandler)

	port := ":8080"
	fmt.Printf("Speedtest server starting on http://localhost%s\n", port)
	fmt.Printf("Download endpoint: GET /download?bytes=<size>\n")
	fmt.Printf("Upload endpoint: POST /upload\n")

	if err := http.ListenAndServe(port, nil); err != nil {
		fmt.Printf("Server error: %v\n", err)
	}
}

func downloadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	bytesParam := r.URL.Query().Get("bytes")
	if bytesParam == "" {
		http.Error(w, "Missing 'bytes' parameter", http.StatusBadRequest)
		return
	}

	numBytes, err := strconv.ParseInt(bytesParam, 10, 64)
	if err != nil || numBytes <= 0 {
		http.Error(w, "Invalid 'bytes' parameter", http.StatusBadRequest)
		return
	}

	w.Header().Set("Content-Type", "application/octet-stream")
	w.Header().Set("Content-Length", strconv.FormatInt(numBytes, 10))

	poolLen := int64(len(randomPool))
	offset := int64(0)
	remaining := numBytes

	for remaining > 0 {
		start := offset % poolLen
		end := start + remaining
		if end > poolLen {
			end = poolLen
		}
		n, writeErr := w.Write(randomPool[start:end])
		if writeErr != nil {
			return
		}
		written := int64(n)
		remaining -= written
		offset += written
	}
}

func uploadHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	defer r.Body.Close()

	buf := make([]byte, uploadBufSize)
	for {
		_, err := r.Body.Read(buf)
		if err == io.EOF {
			break
		}
		if err != nil {
			http.Error(w, "Error reading body", http.StatusInternalServerError)
			return
		}
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Upload complete"))
}