package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"math/bits"
	"net"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const observationShardCount = 64

type observation struct {
	topic         string
	seq           int64
	senderMachine string
	senderIndex   int
	scenario      string
	sentAt        int64
	bits          []uint64
}

type observationShard struct {
	mu       sync.RWMutex
	observed map[string]*observation
	history  map[string]bool
}

type sequenceState struct {
	mu   sync.Mutex
	last map[string]int64
}

type deliveryRecorder struct {
	machine      string
	localUsers   int
	latency      *DurationHistogram
	localLatency *DurationHistogram
	crossLatency *DurationHistogram
	disconnects  atomic.Int64
	received     atomic.Int64
	runID        atomic.Value
	historyMode  atomic.Bool

	shards                []observationShard
	sequences             []sequenceState
	duplicates            atomic.Int64
	outOfOrder            atomic.Int64
	malformed             atomic.Int64
	crossClockSkewSamples atomic.Int64

	orderingMu sync.Mutex
	ordering   []OrderingViolation

	disconnectMu      sync.Mutex
	disconnectReasons map[string]int64
	disconnectSamples []DisconnectSample
}

func newDeliveryRecorder(runID, machine string, localUsers int) *deliveryRecorder {
	shards := make([]observationShard, observationShardCount)
	for i := range shards {
		shards[i].observed = make(map[string]*observation)
		shards[i].history = make(map[string]bool)
	}
	sequences := make([]sequenceState, localUsers)
	for i := range sequences {
		sequences[i].last = make(map[string]int64)
	}
	recorder := &deliveryRecorder{
		machine: machine, localUsers: localUsers, latency: NewDurationHistogram(),
		localLatency: NewDurationHistogram(), crossLatency: NewDurationHistogram(),
		shards: shards, sequences: sequences, disconnectReasons: make(map[string]int64),
	}
	recorder.runID.Store(runID)
	return recorder
}

func (r *deliveryRecorder) setRunID(runID string) { r.runID.Store(runID) }

func (r *deliveryRecorder) resetDisconnects() {
	r.disconnectMu.Lock()
	r.disconnects.Store(0)
	r.disconnectReasons = make(map[string]int64)
	r.disconnectSamples = nil
	r.disconnectMu.Unlock()
}

func (r *deliveryRecorder) onData(account AccountRecord, data dataPacket) {
	runID, _ := r.runID.Load().(string)
	mid, _ := data.Head["x-im-client-mid"].(string)
	if mid == "" || !strings.HasPrefix(mid, runID+"-") {
		return
	}
	var content struct {
		RunID         string `json:"run_id"`
		SenderMachine string `json:"sender_machine"`
		SenderIndex   int    `json:"sender_index"`
		Scenario      string `json:"scenario"`
		SentAt        int64  `json:"sent_at_unix_nano"`
	}
	if json.Unmarshal(data.Content, &content) != nil || content.RunID != runID {
		r.malformed.Add(1)
		return
	}
	if r.historyMode.Load() {
		shard := r.shard(mid)
		shard.mu.Lock()
		shard.history[mid] = true
		shard.mu.Unlock()
		return
	}
	index := account.Index - 1
	if index < 0 || index >= r.localUsers {
		return
	}
	shard := r.shard(mid)
	shard.mu.Lock()
	obs := shard.observed[mid]
	if obs == nil {
		obs = &observation{topic: data.Topic, seq: data.Seq, senderMachine: content.SenderMachine,
			senderIndex: content.SenderIndex, scenario: content.Scenario, sentAt: content.SentAt,
			bits: make([]uint64, (r.localUsers+63)/64)}
		shard.observed[mid] = obs
	}
	word, bit := index/64, uint(index%64)
	mask := uint64(1) << bit
	if obs.bits[word]&mask != 0 {
		shard.mu.Unlock()
		r.duplicates.Add(1)
		return
	}
	obs.bits[word] |= mask
	shard.mu.Unlock()
	r.received.Add(1)
	sequence := &r.sequences[index]
	sequence.mu.Lock()
	previous := sequence.last[data.Topic]
	violation := previous > 0 && data.Seq < previous
	if data.Seq > previous {
		sequence.last[data.Topic] = data.Seq
	}
	sequence.mu.Unlock()
	if violation {
		r.outOfOrder.Add(1)
		r.orderingMu.Lock()
		if len(r.ordering) < 100 {
			r.ordering = append(r.ordering, OrderingViolation{
				ReceiverMachine: r.machine,
				ReceiverIndex:   account.Index,
				TopicHash:       topicHash(data.Topic),
				PreviousSeq:     previous,
				CurrentSeq:      data.Seq,
				ClientMid:       mid,
			})
		}
		r.orderingMu.Unlock()
	}
	if content.SentAt > 0 {
		delay := time.Since(time.Unix(0, content.SentAt))
		r.latency.Observe(delay)
		senderMachine := content.SenderMachine
		if senderMachine == "" {
			senderMachine = senderMachineFromClientMid(runID, mid)
		}
		if senderMachine == r.machine {
			r.localLatency.Observe(delay)
		} else if delay < 0 {
			r.crossClockSkewSamples.Add(1)
		} else {
			r.crossLatency.Observe(delay)
		}
	}
}

func senderMachineFromClientMid(runID, mid string) string {
	prefix := strings.TrimPrefix(mid, runID+"-")
	parts := strings.SplitN(prefix, "-", 3)
	if len(parts) < 2 {
		return ""
	}
	return parts[1]
}

func (r *deliveryRecorder) metricCounts() (received, duplicates, outOfOrder, disconnects int64) {
	return r.received.Load(), r.duplicates.Load(), r.outOfOrder.Load(), r.disconnects.Load()
}

func (r *deliveryRecorder) beginHistory() { r.historyMode.Store(true) }

func (r *deliveryRecorder) historyFound(mid string) bool {
	shard := r.shard(mid)
	shard.mu.RLock()
	found := shard.history[mid]
	shard.mu.RUnlock()
	return found
}

func (r *deliveryRecorder) historyFoundAll(messageIDs []string) bool {
	for _, mid := range messageIDs {
		if !r.historyFound(mid) {
			return false
		}
	}
	return true
}

func (r *deliveryRecorder) unackedObservedMessages(ledger []AcceptedMessage, scenario string) []AcceptedMessage {
	accepted := make(map[string]bool, len(ledger))
	for _, message := range ledger {
		accepted[message.ClientMid] = true
	}
	runID, _ := r.runID.Load().(string)
	messages := make([]AcceptedMessage, 0)
	for i := range r.shards {
		shard := &r.shards[i]
		shard.mu.RLock()
		for mid, obs := range shard.observed {
			if accepted[mid] || obs.seq <= 0 || obs.topic == "" {
				continue
			}
			kind := "p2p"
			if strings.HasPrefix(obs.topic, "grp") {
				kind = "group"
			}
			observedScenario := obs.scenario
			if observedScenario == "" {
				observedScenario = scenario
			}
			message := AcceptedMessage{RunID: runID, Scenario: observedScenario, Kind: kind,
				ClientMid: mid, Topic: obs.topic, Seq: obs.seq, SenderMachine: obs.senderMachine,
				SenderIndex: obs.senderIndex, PeerMachine: r.machine, PeerIndex: firstObservedIndex(obs.bits)}
			if obs.sentAt > 0 {
				message.SentAt = time.Unix(0, obs.sentAt).UTC()
			}
			messages = append(messages, message)
		}
		shard.mu.RUnlock()
	}
	sort.Slice(messages, func(i, j int) bool { return messages[i].ClientMid < messages[j].ClientMid })
	return messages
}

func firstObservedIndex(words []uint64) int {
	for wordIndex, word := range words {
		if word != 0 {
			return wordIndex*64 + bits.TrailingZeros64(word) + 1
		}
	}
	return 0
}

func (r *deliveryRecorder) onDisconnect(account AccountRecord, err error) {
	category := disconnectCategory(err)
	r.disconnectMu.Lock()
	r.disconnects.Add(1)
	r.disconnectReasons[category]++
	if len(r.disconnectSamples) < 100 {
		message := ""
		if err != nil {
			message = err.Error()
			if len(message) > 300 {
				message = message[:300]
			}
		}
		r.disconnectSamples = append(r.disconnectSamples, DisconnectSample{
			ReceiverMachine: r.machine, ReceiverIndex: account.Index, Category: category,
			Error: message, At: time.Now().UTC(),
		})
	}
	r.disconnectMu.Unlock()
}

func disconnectCategory(err error) string {
	if err == nil {
		return "unknown"
	}
	if errors.Is(err, io.EOF) {
		return "eof"
	}
	if errors.Is(err, net.ErrClosed) {
		return "closed_connection"
	}
	var netErr net.Error
	if errors.As(err, &netErr) && netErr.Timeout() {
		return "timeout"
	}
	message := strings.ToLower(err.Error())
	switch {
	case strings.Contains(message, "close 1006") || strings.Contains(message, "abnormal closure"):
		return "abnormal_close"
	case strings.Contains(message, "close 1000") || strings.Contains(message, "normal closure"):
		return "normal_close"
	case strings.Contains(message, "connection reset"):
		return "connection_reset"
	case strings.Contains(message, "broken pipe"):
		return "broken_pipe"
	case strings.Contains(message, "timeout") || strings.Contains(message, "deadline exceeded"):
		return "timeout"
	case strings.Contains(message, "use of closed network connection"):
		return "closed_connection"
	default:
		return "other"
	}
}

func (r *deliveryRecorder) summarize(base AgentSummary, ledger []AcceptedMessage) AgentSummary {
	base.DeliveryLatency = r.latency.Snapshot()
	base.DeliveryLatencyLocal = r.localLatency.Snapshot()
	base.DeliveryLatencyCross = r.crossLatency.Snapshot()
	r.disconnectMu.Lock()
	base.Disconnects = r.disconnects.Load()
	base.DisconnectReasons = cloneCountMap(r.disconnectReasons)
	base.DisconnectSamples = append([]DisconnectSample(nil), r.disconnectSamples...)
	r.disconnectMu.Unlock()
	runID, _ := r.runID.Load().(string)
	accepted := make(map[string]AcceptedMessage, len(ledger))
	const missingLimit = 100000

	base.Duplicates = r.duplicates.Load()
	base.OutOfOrder = r.outOfOrder.Load()
	base.CrossClockSkewSamples = r.crossClockSkewSamples.Load()
	r.orderingMu.Lock()
	base.OrderingViolations = append([]OrderingViolation(nil), r.ordering...)
	r.orderingMu.Unlock()
	if malformed := r.malformed.Load(); malformed > 0 {
		base.Errors = append(base.Errors, fmt.Sprintf("malformed run messages=%d", malformed))
	}
	for _, message := range ledger {
		if message.RunID != runID {
			continue
		}
		accepted[message.ClientMid] = message
		obs := r.observationSnapshot(message.ClientMid)
		switch message.Kind {
		case "group":
			base.ExpectedGroup += int64(r.localUsers)
			for index := 1; index <= r.localUsers; index++ {
				if observationHas(obs, index) {
					base.ReceivedGroup++
				} else if len(base.Missing) < missingLimit {
					base.Missing = append(base.Missing, missingFor(message, index))
				}
			}
		case "p2p":
			if message.PeerMachine == r.machine && message.PeerIndex > 0 {
				base.ExpectedPeer++
				if observationHas(obs, message.PeerIndex) {
					base.ReceivedPeer++
				} else if len(base.Missing) < missingLimit {
					base.Missing = append(base.Missing, missingFor(message, message.PeerIndex))
				}
			}
			if message.SenderMachine == r.machine && message.SenderIndex > 0 {
				base.ExpectedEcho++
				if observationHas(obs, message.SenderIndex) {
					base.ReceivedEcho++
				}
			}
		}
	}
	for i := range r.shards {
		shard := &r.shards[i]
		shard.mu.RLock()
		for mid, obs := range shard.observed {
			if _, ok := accepted[mid]; !ok {
				base.UnexpectedMessages += bitsetCount(obs.bits)
			}
		}
		shard.mu.RUnlock()
	}
	base.RealtimeLost = (base.ExpectedPeer - base.ReceivedPeer) + (base.ExpectedGroup - base.ReceivedGroup)
	if len(base.Missing) == missingLimit {
		base.Errors = append(base.Errors, "missing delivery details truncated at 100000")
	}
	return base
}

func (r *deliveryRecorder) countRealtimeLostStored(ledger []AcceptedMessage, storedMessageIDs []string) int64 {
	runID, _ := r.runID.Load().(string)
	stored := make(map[string]bool, len(storedMessageIDs))
	for _, mid := range storedMessageIDs {
		stored[mid] = true
	}
	var count int64
	for _, message := range ledger {
		if message.RunID != runID || !stored[message.ClientMid] {
			continue
		}
		obs := r.observationSnapshot(message.ClientMid)
		switch message.Kind {
		case "group":
			for index := 1; index <= r.localUsers; index++ {
				if !observationHas(obs, index) {
					count++
				}
			}
		case "p2p":
			if message.PeerMachine == r.machine && message.PeerIndex > 0 && !observationHas(obs, message.PeerIndex) {
				count++
			}
		}
	}
	return count
}

func (r *deliveryRecorder) shard(mid string) *observationShard {
	var hash uint32 = 2166136261
	for i := 0; i < len(mid); i++ {
		hash ^= uint32(mid[i])
		hash *= 16777619
	}
	return &r.shards[hash%uint32(len(r.shards))]
}

func (r *deliveryRecorder) observationSnapshot(mid string) *observation {
	shard := r.shard(mid)
	shard.mu.RLock()
	obs := shard.observed[mid]
	if obs != nil {
		copyOfObservation := *obs
		copyOfObservation.bits = append([]uint64(nil), obs.bits...)
		obs = &copyOfObservation
	}
	shard.mu.RUnlock()
	return obs
}

func cloneCountMap(source map[string]int64) map[string]int64 {
	if len(source) == 0 {
		return nil
	}
	clone := make(map[string]int64, len(source))
	for key, value := range source {
		clone[key] = value
	}
	return clone
}

func observationHas(obs *observation, receiverIndex int) bool {
	if obs == nil || receiverIndex <= 0 {
		return false
	}
	index := receiverIndex - 1
	word, bit := index/64, uint(index%64)
	return word < len(obs.bits) && obs.bits[word]&(uint64(1)<<bit) != 0
}

func bitsetCount(words []uint64) int64 {
	var count int64
	for _, word := range words {
		count += int64(bits.OnesCount64(word))
	}
	return count
}

func missingFor(message AcceptedMessage, receiver int) MissingDelivery {
	return MissingDelivery{ClientMid: message.ClientMid, Kind: message.Kind, TopicHash: topicHash(message.Topic),
		Seq: message.Seq, ReceiverIndex: receiver}
}