package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"log"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

var (
	target      = flag.String("server", "127.0.0.1:6379", "Valkey server address")
	concurrency = flag.Int("concurrency", 640, "Number of concurrent connections")
	pipeline    = flag.Int("pipeline", 100, "Pipeline size")
	payloadSize = flag.Int("size", 1024, "Data size in bytes")
	ops         uint64
)

func main() {
	flag.Parse()
	fmt.Printf("Starting WRITE-heavy repro against %s\n", *target)
	fmt.Printf("Concurrency: %d | Pipeline: %d | Payload: %d bytes\n", *concurrency, *pipeline, *payloadSize)
	fmt.Println("Workload: SET key value")

	var wg sync.WaitGroup
	wg.Add(*concurrency)

	// Start status reporter
	go func() {
		ticker := time.NewTicker(1 * time.Second)
		defer ticker.Stop()
		for range ticker.C {
			count := atomic.SwapUint64(&ops, 0)
			fmt.Printf("OPS: %d\n", count)
		}
	}()

	// Launch Load
	for i := 0; i < *concurrency; i++ {
		go func(id int) {
			defer wg.Done()
			worker(id)
		}(i)
	}

	wg.Wait()
}

func worker(id int) {
	conn, err := net.DialTimeout("tcp", *target, 5*time.Second)
	if err != nil {
		log.Printf("Worker %d failed to connect: %v", id, err)
		return
	}
	defer conn.Close()

	if tcpConn, ok := conn.(*net.TCPConn); ok {
		tcpConn.SetNoDelay(true)
	}

	// 1. Pre-build the pipeline buffer for SET commands
	// We use a unique key per worker to avoid lock contention masking the IO bug
	// Key: repro:<id>
	key := fmt.Sprintf("repro:%d", id)
	value := bytes.Repeat([]byte("x"), *payloadSize)

	// RESP: *3\r\n$3\r\nSET\r\n$<klen>\r\n<key>\r\n$<vlen>\r\n<val>\r\n
	cmdHeader := fmt.Sprintf("*3\r\n$3\r\nSET\r\n$%d\r\n%s\r\n$%d\r\n", len(key), key, len(value))
	// Combine header + body + CRLF
	fullCmd := append([]byte(cmdHeader), value...)
	fullCmd = append(fullCmd, []byte("\r\n")...)

	var buf bytes.Buffer
	for i := 0; i < *pipeline; i++ {
		buf.Write(fullCmd)
	}
	payload := buf.Bytes()

	reader := bufio.NewReader(conn)

	for {
		// Strict Deadline
		conn.SetDeadline(time.Now().Add(2 * time.Second))

		// Write the massive pipeline
		if _, err := conn.Write(payload); err != nil {
			log.Printf("Worker %d write failed: %v", id, err)
			return
		}

		// Read exactly 'pipeline' responses (Expecting +OK)
		for i := 0; i < *pipeline; i++ {
			line, err := reader.ReadBytes('\n')
			if err != nil {
				log.Printf("Worker %d HANG/ERROR on reply %d/%d: %v", id, i+1, *pipeline, err)
				return
			}

			// Verify we got "+OK"
			// (Trimming space to handle \r\n)
			if len(line) < 1 || line[0] != '+' {
				log.Printf("Worker %d unexpected response: %s", id, string(line))
				return
			}
		}
		atomic.AddUint64(&ops, uint64(*pipeline))
	}
}