package mtprotoedge

import (
	"context"
	"fmt"
	"os"
	"runtime"
	"sort"
	"strconv"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"go.uber.org/zap"

	"github.com/iamxvbaba/td/proto"
	"github.com/iamxvbaba/td/tg"
)

// Connection-layer fan-out / churn load test: focused on SessionManager lock
// contention, without going through real sockets / encryption.
//
// The constructed Conn deliberately does not startOutbound: after pushToUser takes
// a locked snapshot of byUser, it calls c.Send for each conn outside the lock, and
// at that point c.outbound==nil immediately returns ErrConnClosed (see
// outbound.go), so the measurement concentrates on "the locked section + dispatch
// overhead", i.e. the global lock hot spot that sharding aims to eliminate. Each
// connection uses only struct memory (no 1024-capacity outbound channel, no
// goroutine), so it can be registered at the 200k scale.
//
// Usage:
//
//	go test ./internal/mtprotoedge/ -run '^$' -bench BenchmarkSessionManager -benchmem -cpu 1,4,8
//	go test ./internal/mtprotoedge/ -run '^$' -bench BenchmarkSessionManagerPushConcurrent -mutexprofile mu.out
//	TELESRV_LOAD_CONNS=200000 go test ./internal/mtprotoedge/ -run TestSessionManagerFanoutThroughput -v -timeout 300s

func benchConn(sessionID int64, authKeyID [8]byte, userID int64) *Conn {
	c := &Conn{sessionID: sessionID, authKeyID: authKeyID}
	if userID != 0 {
		c.userID.Store(userID)
		c.userIDResolved.Store(true)
	}
	c.receivesUpdates.Store(true) // take fanout's "collect conns→Send outside lock" branch, not pending staging
	return c
}

func authKeyIDFromInt(v uint64) [8]byte {
	var id [8]byte
	for i := 0; i < 8; i++ {
		id[i] = byte(v >> (8 * i))
	}
	return id
}

// seedSessions registers conns connections, binding connsPerUser connections per
// user (simulating multiple devices). It returns the number of registered users
// (deduplicated, ordered range [1, userCount]).
func seedSessions(sm *SessionManager, conns, connsPerUser int) (userCount int) {
	if connsPerUser < 1 {
		connsPerUser = 1
	}
	for i := 0; i < conns; i++ {
		userID := int64(i/connsPerUser) + 1
		sm.Register(benchConn(int64(i)+1, authKeyIDFromInt(uint64(i)+1), userID))
	}
	return (conns + connsPerUser - 1) / connsPerUser
}

// BenchmarkSessionManagerPushConcurrent simulates the real hot spot at 200k online:
// many goroutines concurrently pushToUser to different users, all contending for
// the same global lock. -mutexprofile will surface SessionManager.mu at the top.
func BenchmarkSessionManagerPushConcurrent(b *testing.B) {
	const conns = 200_000
	const connsPerUser = 2
	sm := NewSessionManager(zap.NewNop())
	userCount := seedSessions(sm, conns, connsPerUser)
	msg := &tg.UpdatesTooLong{}
	ctx := context.Background()

	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		var n uint64
		for pb.Next() {
			n++
			userID := int64(n%uint64(userCount)) + 1
			_, _ = sm.PushToUser(ctx, userID, proto.MessageFromServer, msg)
		}
	})
}

// BenchmarkSessionManagerRegisterChurn measures the lock cost of connection
// setup/teardown: concurrent Register+Unregister. 200k online means continuous
// connect/disconnect churn, each contending for the global write lock.
func BenchmarkSessionManagerRegisterChurn(b *testing.B) {
	sm := NewSessionManager(zap.NewNop())
	var seq atomic.Uint64

	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			id := seq.Add(1)
			c := benchConn(int64(id), authKeyIDFromInt(id), int64(id))
			sm.Register(c)
			sm.Unregister(c)
		}
	})
}

// BenchmarkSessionManagerPushFanoutWidth measures the fanout breadth cost of a
// single push: one user bound to many connections, and a single PushToUser must
// traverse all of them while holding the lock. A real private-chat user has few
// devices (2-4), so this is an upper-bound reference.
func BenchmarkSessionManagerPushFanoutWidth(b *testing.B) {
	for _, width := range []int{1, 4, 16, 64} {
		b.Run(fmt.Sprintf("width=%d", width), func(b *testing.B) {
			sm := NewSessionManager(zap.NewNop())
			const userID = 1
			for i := 0; i < width; i++ {
				sm.Register(benchConn(int64(i)+1, authKeyIDFromInt(uint64(i)+1), userID))
			}
			msg := &tg.UpdatesTooLong{}
			ctx := context.Background()
			b.ResetTimer()
			for i := 0; i < b.N; i++ {
				_, _ = sm.PushToUser(ctx, userID, proto.MessageFromServer, msg)
			}
		})
	}
}

// TestSessionManagerFanoutThroughput is a data-driven throughput test: after
// registering N connections, P goroutines continuously push concurrently, measuring
// actual push throughput and p99 under the global lock. Small smoke scale by
// default; set TELESRV_LOAD_CONNS to scale up to 200k.
func TestSessionManagerFanoutThroughput(t *testing.T) {
	conns := envIntDefault("TELESRV_LOAD_CONNS", 20_000)
	connsPerUser := envIntDefault("TELESRV_LOAD_CONNS_PER_USER", 2)
	workers := envIntDefault("TELESRV_LOAD_PUSH_WORKERS", 0) // 0 → GOMAXPROCS
	duration := time.Duration(envIntDefault("TELESRV_LOAD_SECONDS", 3)) * time.Second
	if workers <= 0 {
		workers = runtime.GOMAXPROCS(0)
	}

	sm := NewSessionManager(zap.NewNop())
	t0 := time.Now()
	userCount := seedSessions(sm, conns, connsPerUser)
	seedWall := time.Since(t0)
	if got := sm.Online(); got != conns {
		t.Fatalf("online = %d, want %d", got, conns)
	}

	msg := &tg.UpdatesTooLong{}
	ctx := context.Background()
	var ops atomic.Int64
	perWorkerLat := make([][]time.Duration, workers)

	deadline := time.Now().Add(duration)
	var wg sync.WaitGroup
	start := time.Now()
	for w := 0; w < workers; w++ {
		wg.Add(1)
		go func(w int) {
			defer wg.Done()
			lat := make([]time.Duration, 0, 1<<16)
			var n uint64
			for time.Now().Before(deadline) {
				// Check the clock once every batch of 256, reducing time.Now's share.
				for j := 0; j < 256; j++ {
					n++
					userID := int64(n%uint64(userCount)) + 1
					s := time.Now()
					_, _ = sm.PushToUser(ctx, userID, proto.MessageFromServer, msg)
					lat = append(lat, time.Since(s))
				}
				ops.Add(256)
			}
			perWorkerLat[w] = lat
		}(w)
	}
	wg.Wait()
	wall := time.Since(start)

	all := make([]time.Duration, 0, ops.Load())
	for _, l := range perWorkerLat {
		all = append(all, l...)
	}
	sortDurations(all)
	total := ops.Load()
	thr := float64(total) / wall.Seconds()

	t.Logf("==== session_manager fan-out throughput ====")
	t.Logf("config:  conns=%d connsPerUser=%d users=%d pushWorkers=%d dur=%s seed=%s",
		conns, connsPerUser, userCount, workers, duration, seedWall.Round(time.Millisecond))
	t.Logf("push:    %d ops in %s -> %.0f push/s", total, wall.Round(time.Millisecond), thr)
	t.Logf("push.lat p50=%s p90=%s p99=%s max=%s",
		pct(all, 50), pct(all, 90), pct(all, 99), pct(all, 100))
	t.Logf("=============================================")
}

func pct(sorted []time.Duration, p int) time.Duration {
	if len(sorted) == 0 {
		return 0
	}
	idx := (p*len(sorted))/100 - 1
	if idx < 0 {
		idx = 0
	}
	if idx >= len(sorted) {
		idx = len(sorted) - 1
	}
	return sorted[idx]
}

func sortDurations(d []time.Duration) {
	sort.Slice(d, func(i, j int) bool { return d[i] < d[j] })
}

func envIntDefault(key string, def int) int {
	if v := os.Getenv(key); v != "" {
		if n, err := strconv.Atoi(v); err == nil {
			return n
		}
	}
	return def
}