package main

import (
	"bufio"
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/csv"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"math"
	"net/http"
	"net/url"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
)

type account struct {
	Username string
	Password string
	Token    string
}

type client struct {
	id      int
	account account
	wsURL   string
	conn    *websocket.Conn
	sendMu  sync.Mutex
	stats   *probeStats
}

type probeStats struct {
	expected atomic.Int64
	received atomic.Int64
	acks     atomic.Int64
	errors   atomic.Int64

	mu        sync.Mutex
	latencies []time.Duration
	seen      map[string]struct{}
}

func main() {
	var wsList string
	var apiKey string
	var accountsPath string
	var topic string
	var subscribers int
	var publishers int
	var rate int
	var duration time.Duration
	var ramp time.Duration
	var insecure bool

	flag.StringVar(&wsList, "ws", "ws://localhost:6060/v0/channels", "Comma-separated websocket endpoints, one or more pods/LBs.")
	flag.StringVar(&apiKey, "apikey", "", "Tinode API key appended to websocket URL.")
	flag.StringVar(&accountsPath, "accounts", "loadtest/users.csv", "CSV: username,password[,token].")
	flag.StringVar(&topic, "topic", "", "Group topic to subscribe and publish to, e.g. grpABC.")
	flag.IntVar(&subscribers, "subscribers", 100, "Number of websocket subscribers.")
	flag.IntVar(&publishers, "publishers", 1, "Number of publishers selected from subscribers.")
	flag.IntVar(&rate, "rate", 10, "Total publish rate, messages per second.")
	flag.DurationVar(&duration, "duration", time.Minute, "Publish duration.")
	flag.DurationVar(&ramp, "ramp", 10*time.Second, "Subscriber connection ramp duration.")
	flag.BoolVar(&insecure, "insecure", false, "Skip TLS verification for wss endpoints.")
	flag.Parse()

	if topic == "" || subscribers <= 0 || publishers <= 0 || rate <= 0 || duration <= 0 {
		log.Fatal("topic, subscribers, publishers, rate, and duration are required and must be positive")
	}
	endpoints := parseEndpoints(wsList, apiKey)
	accounts, err := readAccounts(accountsPath)
	if err != nil {
		log.Fatal(err)
	}
	if len(accounts) < subscribers {
		log.Fatalf("need at least %d accounts, got %d", subscribers, len(accounts))
	}
	if publishers > subscribers {
		publishers = subscribers
	}

	stats := &probeStats{seen: make(map[string]struct{})}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	clients := make([]*client, subscribers)
	var wg sync.WaitGroup
	for i := 0; i < subscribers; i++ {
		c := &client{id: i, account: accounts[i], wsURL: endpoints[i%len(endpoints)], stats: stats}
		clients[i] = c
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := c.connectAndSubscribe(ctx, topic, insecure); err != nil {
				stats.errors.Add(1)
				log.Printf("client %d connect failed: %v", c.id, err)
			}
		}()
		if ramp > 0 {
			time.Sleep(ramp / time.Duration(subscribers))
		}
	}
	wg.Wait()
	active := 0
	for _, c := range clients {
		if c.conn != nil {
			active++
		}
	}
	if active == 0 {
		log.Fatal("no active websocket clients")
	}
	log.Printf("subscribed clients=%d publishers=%d endpoints=%d topic=%s", active, publishers, len(endpoints), topic)

	publishDone := make(chan struct{})
	go publishLoop(ctx, clients[:publishers], topic, rate, duration, stats, publishDone)
	<-publishDone
	log.Printf("publish finished; waiting 10s for fanout drain")
	time.Sleep(10 * time.Second)
	cancel()
	for _, c := range clients {
		if c.conn != nil {
			_ = c.conn.Close()
		}
	}
	printSummary(stats, active)
}

func (c *client) connectAndSubscribe(ctx context.Context, topic string, insecure bool) error {
	dialer := websocket.DefaultDialer
	if insecure {
		dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
	}
	conn, _, err := dialer.DialContext(ctx, c.wsURL, http.Header{})
	if err != nil {
		return err
	}
	c.conn = conn
	if err = c.sendAndWaitCtrl("hi", map[string]any{"hi": map[string]any{"id": c.msgID("hi"), "ver": "0.22.8", "ua": "realtime-fanout-probe/1.0"}}); err != nil {
		return err
	}
	if c.account.Token != "" {
		if err = c.sendAndWaitCtrl("login", map[string]any{"login": map[string]any{"id": c.msgID("login"), "scheme": "token", "secret": c.account.Token}}); err != nil {
			return err
		}
	} else {
		secret := base64.StdEncoding.EncodeToString([]byte(c.account.Username + ":" + c.account.Password))
		if err = c.sendAndWaitCtrl("login", map[string]any{"login": map[string]any{"id": c.msgID("login"), "scheme": "basic", "secret": secret}}); err != nil {
			return err
		}
	}
	if err = c.sendAndWaitCtrl("sub", map[string]any{"sub": map[string]any{"id": c.msgID("sub"), "topic": topic, "get": map[string]any{"what": "desc"}}}); err != nil {
		return err
	}
	go c.readLoop(ctx)
	return nil
}

func (c *client) sendAndWaitCtrl(name string, msg map[string]any) error {
	c.sendMu.Lock()
	if err := c.conn.WriteJSON(msg); err != nil {
		c.sendMu.Unlock()
		return err
	}
	c.sendMu.Unlock()
	deadline := time.Now().Add(15 * time.Second)
	for time.Now().Before(deadline) {
		var raw map[string]json.RawMessage
		if err := c.conn.ReadJSON(&raw); err != nil {
			return err
		}
		if ctrlRaw, ok := raw["ctrl"]; ok {
			var ctrl struct {
				Code int    `json:"code"`
				Text string `json:"text"`
			}
			_ = json.Unmarshal(ctrlRaw, &ctrl)
			if ctrl.Code >= 200 && ctrl.Code < 300 {
				return nil
			}
			return fmt.Errorf("%s ctrl code=%d text=%s", name, ctrl.Code, ctrl.Text)
		}
	}
	return fmt.Errorf("%s ctrl timeout", name)
}

func (c *client) readLoop(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		var raw map[string]json.RawMessage
		if err := c.conn.ReadJSON(&raw); err != nil {
			return
		}
		if _, ok := raw["ctrl"]; ok {
			c.stats.acks.Add(1)
			continue
		}
		dataRaw, ok := raw["data"]
		if !ok {
			continue
		}
		var data struct {
			Content json.RawMessage `json:"content"`
		}
		if err := json.Unmarshal(dataRaw, &data); err != nil || len(data.Content) == 0 {
			continue
		}
		var content struct {
			ProbeID string `json:"probe_id"`
			SentAt  int64  `json:"sent_at_unix_nano"`
		}
		if err := json.Unmarshal(data.Content, &content); err != nil || content.ProbeID == "" || content.SentAt <= 0 {
			continue
		}
		c.stats.record(content.ProbeID, time.Since(time.Unix(0, content.SentAt)))
	}
}

func publishLoop(ctx context.Context, publishers []*client, topic string, rate int, duration time.Duration, stats *probeStats, done chan<- struct{}) {
	defer close(done)
	interval := time.Second / time.Duration(rate)
	if interval <= 0 {
		interval = time.Nanosecond
	}
	ticker := time.NewTicker(interval)
	defer ticker.Stop()
	deadline := time.Now().Add(duration)
	seq := int64(0)
	for time.Now().Before(deadline) {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
		}
		n := atomic.AddInt64(&seq, 1)
		publisher := publishers[int(n)%len(publishers)]
		probeID := fmt.Sprintf("probe-%d-%d", time.Now().UnixNano(), n)
		msg := map[string]any{"pub": map[string]any{
			"id":    publisher.msgID("pub-" + strconv.FormatInt(n, 10)),
			"topic": topic,
			"head":  map[string]any{"mime": "application/json"},
			"content": map[string]any{
				"probe_id":          probeID,
				"sent_at_unix_nano": time.Now().UnixNano(),
				"seq":               n,
			},
		}}
		publisher.sendMu.Lock()
		err := publisher.conn.WriteJSON(msg)
		publisher.sendMu.Unlock()
		if err != nil {
			stats.errors.Add(1)
			log.Printf("publish failed client=%d seq=%d: %v", publisher.id, n, err)
			continue
		}
		stats.expected.Add(1)
	}
}

func (s *probeStats) record(probeID string, latency time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.received.Add(1)
	s.seen[probeID] = struct{}{}
	s.latencies = append(s.latencies, latency)
}

func printSummary(stats *probeStats, activeSubscribers int) {
	stats.mu.Lock()
	latencies := append([]time.Duration(nil), stats.latencies...)
	unique := len(stats.seen)
	stats.mu.Unlock()
	sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] })
	published := stats.expected.Load()
	expectedDeliveries := published * int64(activeSubscribers)
	received := stats.received.Load()
	lost := expectedDeliveries - received
	lossRate := 0.0
	if expectedDeliveries > 0 {
		lossRate = float64(lost) / float64(expectedDeliveries) * 100
	}
	fmt.Println("==== realtime fanout probe summary ====")
	fmt.Printf("published=%d active_subscribers=%d expected_deliveries=%d received=%d lost=%d loss_rate=%.4f%% unique_messages_seen=%d errors=%d ctrl_acks=%d\n",
		published, activeSubscribers, expectedDeliveries, received, lost, lossRate, unique, stats.errors.Load(), stats.acks.Load())
	fmt.Printf("latency p50=%s p95=%s p99=%s max=%s samples=%d\n",
		percentile(latencies, 50), percentile(latencies, 95), percentile(latencies, 99), maxDuration(latencies), len(latencies))
}

func percentile(values []time.Duration, p int) time.Duration {
	if len(values) == 0 {
		return 0
	}
	idx := int(math.Ceil(float64(p)/100*float64(len(values)))) - 1
	if idx < 0 {
		idx = 0
	}
	if idx >= len(values) {
		idx = len(values) - 1
	}
	return values[idx]
}

func maxDuration(values []time.Duration) time.Duration {
	if len(values) == 0 {
		return 0
	}
	return values[len(values)-1]
}

func (c *client) msgID(suffix string) string {
	return fmt.Sprintf("rtp-%d-%s", c.id, suffix)
}

func parseEndpoints(raw string, apiKey string) []string {
	parts := strings.Split(raw, ",")
	out := make([]string, 0, len(parts))
	for _, part := range parts {
		endpoint := strings.TrimSpace(part)
		if endpoint == "" {
			continue
		}
		u, err := url.Parse(endpoint)
		if err != nil {
			log.Fatalf("bad ws endpoint %q: %v", endpoint, err)
		}
		q := u.Query()
		if apiKey != "" && q.Get("apikey") == "" {
			q.Set("apikey", apiKey)
		}
		u.RawQuery = q.Encode()
		out = append(out, u.String())
	}
	if len(out) == 0 {
		log.Fatal("no websocket endpoints")
	}
	return out
}

func readAccounts(path string) ([]account, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()
	reader := csv.NewReader(bufio.NewReader(file))
	reader.FieldsPerRecord = -1
	rows, err := reader.ReadAll()
	if err != nil {
		return nil, err
	}
	accounts := make([]account, 0, len(rows))
	for _, row := range rows {
		if len(row) < 2 || strings.HasPrefix(strings.TrimSpace(row[0]), "#") {
			continue
		}
		acc := account{Username: strings.TrimSpace(row[0]), Password: strings.TrimSpace(row[1])}
		if len(row) > 2 {
			acc.Token = strings.TrimSpace(row[2])
		}
		accounts = append(accounts, acc)
	}
	return accounts, nil
}