// Command wsclient is a standalone WebSocket load/behaviour test client for the
// wstester dashboard. Build it, or download a prebuilt binary from the
// dashboard, then point it at any WebSocket URL:
//
//	wsclient -url wss://app-xxxx.zerops.app/ws -label my-laptop -interval 1s -size 512
//
// Flags let you drive the different scenarios the dashboard visualises:
//   - steady traffic:      -interval 500ms -size 1024
//   - idle / no-data test:  -idle           (connect and send nothing)
//   - fixed run then quit:  -duration 30s
package main

import (
	"context"
	"crypto/tls"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/gorilla/websocket"
)

func main() {
	var (
		rawURL   = flag.String("url", "", "WebSocket URL to connect to, e.g. wss://host/ws (required)")
		label    = flag.String("label", "", "client label shown on the dashboard")
		interval = flag.Duration("interval", time.Second, "how often to send a message (e.g. 500ms, 2s)")
		size     = flag.Int("size", 256, "payload size in bytes per message")
		duration = flag.Duration("duration", 0, "stop after this long (0 = run until interrupted)")
		idle     = flag.Bool("idle", false, "connect but send no data (test server idle/no-data timeouts)")
		recv     = flag.Bool("recv", true, "read and count messages sent by the server")
		insecure = flag.Bool("insecure", false, "skip TLS certificate verification")
		forceV4  = flag.Bool("4", false, "force IPv4: resolve and connect over IPv4 only")
		forceV6  = flag.Bool("6", false, "force IPv6: resolve and connect over IPv6 only")
	)
	flag.Parse()

	if *rawURL == "" {
		fmt.Fprintln(os.Stderr, "error: -url is required")
		flag.Usage()
		os.Exit(2)
	}
	if *forceV4 && *forceV6 {
		log.Fatalf("cannot use -4 and -6 together")
	}

	// network selects the address family used when dialing the resolved host:
	// "tcp" = dual-stack (default), "tcp4" = IPv4 only, "tcp6" = IPv6 only.
	network := "tcp"
	if *forceV4 {
		network = "tcp4"
	} else if *forceV6 {
		network = "tcp6"
	}

	u, err := url.Parse(*rawURL)
	if err != nil {
		log.Fatalf("invalid url: %v", err)
	}
	// Attach the label as a query param the server understands.
	if *label != "" {
		q := u.Query()
		q.Set("label", *label)
		u.RawQuery = q.Encode()
	}

	// Copy the library defaults (Proxy, HandshakeTimeout) and layer our options on.
	d := *websocket.DefaultDialer
	dialer := &d
	if *insecure {
		dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
	}
	if network != "tcp" {
		nd := &net.Dialer{Timeout: 45 * time.Second}
		// Forcing the network on the resolving dialer makes Go's resolver pick
		// only A (tcp4) or AAAA (tcp6) records for the URL host.
		dialer.NetDialContext = func(ctx context.Context, _, addr string) (net.Conn, error) {
			return nd.DialContext(ctx, network, addr)
		}
	}

	log.Printf("connecting to %s (%s)", u.String(), network)
	conn, resp, err := dialer.Dial(u.String(), nil)
	if err != nil {
		if resp != nil {
			log.Fatalf("dial failed (%s): %v", resp.Status, err)
		}
		log.Fatalf("dial failed: %v", err)
	}
	defer conn.Close()
	log.Printf("connected to %s", conn.UnderlyingConn().RemoteAddr())

	var rxBytes, txBytes, rxMsgs, txMsgs int64

	done := make(chan struct{})

	// Reader: count inbound messages (server-generated traffic, pings handled
	// automatically by gorilla, welcome frame, etc.).
	go func() {
		defer close(done)
		for {
			mt, data, err := conn.ReadMessage()
			if err != nil {
				log.Printf("read: %v", err)
				return
			}
			if *recv {
				atomic.AddInt64(&rxBytes, int64(len(data)))
				atomic.AddInt64(&rxMsgs, 1)
			}
			if mt == websocket.TextMessage && strings.Contains(string(data), "\"type\":\"welcome\"") {
				log.Printf("server welcome: %s", string(data))
			}
		}
	}()

	// Stats printer.
	stats := time.NewTicker(time.Second)
	defer stats.Stop()

	// Writer: send periodic payloads unless in idle mode.
	var sendTick <-chan time.Time
	if !*idle {
		t := time.NewTicker(*interval)
		defer t.Stop()
		sendTick = t.C
	} else {
		log.Printf("idle mode: connected, sending no data")
	}

	payload := make([]byte, *size)
	for i := range payload {
		payload[i] = byte('a' + rand.Intn(26))
	}

	var deadline <-chan time.Time
	if *duration > 0 {
		deadline = time.After(*duration)
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)

	start := time.Now()
	for {
		select {
		case <-done:
			log.Printf("connection closed by server")
			printFinal(start, &rxBytes, &txBytes, &rxMsgs, &txMsgs)
			return
		case <-sig:
			log.Printf("interrupted, closing")
			_ = conn.WriteControl(websocket.CloseMessage,
				websocket.FormatCloseMessage(websocket.CloseNormalClosure, "client exit"),
				time.Now().Add(time.Second))
			printFinal(start, &rxBytes, &txBytes, &rxMsgs, &txMsgs)
			return
		case <-deadline:
			log.Printf("duration reached, closing")
			_ = conn.WriteControl(websocket.CloseMessage,
				websocket.FormatCloseMessage(websocket.CloseNormalClosure, "duration reached"),
				time.Now().Add(time.Second))
			printFinal(start, &rxBytes, &txBytes, &rxMsgs, &txMsgs)
			return
		case <-sendTick:
			_ = conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
			if err := conn.WriteMessage(websocket.TextMessage, payload); err != nil {
				log.Printf("write: %v", err)
				printFinal(start, &rxBytes, &txBytes, &rxMsgs, &txMsgs)
				return
			}
			atomic.AddInt64(&txBytes, int64(len(payload)))
			atomic.AddInt64(&txMsgs, 1)
		case <-stats.C:
			log.Printf("tx: %d msgs / %s | rx: %d msgs / %s",
				atomic.LoadInt64(&txMsgs), humanBytes(atomic.LoadInt64(&txBytes)),
				atomic.LoadInt64(&rxMsgs), humanBytes(atomic.LoadInt64(&rxBytes)))
		}
	}
}

func printFinal(start time.Time, rxBytes, txBytes, rxMsgs, txMsgs *int64) {
	elapsed := time.Since(start)
	log.Printf("=== summary (%.1fs) ===", elapsed.Seconds())
	log.Printf("sent:     %d msgs, %s", atomic.LoadInt64(txMsgs), humanBytes(atomic.LoadInt64(txBytes)))
	log.Printf("received: %d msgs, %s", atomic.LoadInt64(rxMsgs), humanBytes(atomic.LoadInt64(rxBytes)))
}

func humanBytes(n int64) string {
	const unit = 1024
	if n < unit {
		return fmt.Sprintf("%d B", n)
	}
	div, exp := int64(unit), 0
	for m := n / unit; m >= unit; m /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.1f %cB", float64(n)/float64(div), "KMGT"[exp])
}