package main

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

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

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

	mu                    sync.Mutex
	observed              map[string]*observation
	lastSeq               []map[string]int64
	duplicates            int64
	outOfOrder            int64
	ordering              []OrderingViolation
	malformed             int64
	crossClockSkewSamples int64
	history               map[string]bool
}

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

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

func (r *deliveryRecorder) resetDisconnects() { r.disconnects.Store(0) }

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.mu.Lock()
		r.malformed++
		r.mu.Unlock()
		return
	}
	if r.historyMode.Load() {
		r.mu.Lock()
		r.history[mid] = true
		r.mu.Unlock()
		return
	}
	index := account.Index - 1
	if index < 0 || index >= r.localUsers {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	obs := r.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)}
		r.observed[mid] = obs
	}
	word, bit := index/64, uint(index%64)
	mask := uint64(1) << bit
	if obs.bits[word]&mask != 0 {
		r.duplicates++
		return
	}
	obs.bits[word] |= mask
	r.received.Add(1)
	if previous := r.lastSeq[index][data.Topic]; previous > 0 && data.Seq < previous {
		r.outOfOrder++
		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,
			})
		}
	}
	if data.Seq > r.lastSeq[index][data.Topic] {
		r.lastSeq[index][data.Topic] = data.Seq
	}
	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++
		} 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) {
	r.mu.Lock()
	duplicates = r.duplicates
	outOfOrder = r.outOfOrder
	r.mu.Unlock()
	return r.received.Load(), duplicates, outOfOrder, r.disconnects.Load()
}

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

func (r *deliveryRecorder) historyFound(mid string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.history[mid]
}

func (r *deliveryRecorder) historyFoundAll(messageIDs []string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, mid := range messageIDs {
		if !r.history[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)
	r.mu.Lock()
	defer r.mu.Unlock()
	messages := make([]AcceptedMessage, 0)
	for mid, obs := range r.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)
	}
	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(_ AccountRecord, _ error) {
	r.disconnects.Add(1)
}

func (r *deliveryRecorder) summarize(base AgentSummary, ledger []AcceptedMessage) AgentSummary {
	base.DeliveryLatency = r.latency.Snapshot()
	base.DeliveryLatencyLocal = r.localLatency.Snapshot()
	base.DeliveryLatencyCross = r.crossLatency.Snapshot()
	base.Disconnects = r.disconnects.Load()
	runID, _ := r.runID.Load().(string)
	accepted := make(map[string]AcceptedMessage, len(ledger))
	const missingLimit = 100000

	r.mu.Lock()
	defer r.mu.Unlock()
	base.Duplicates = r.duplicates
	base.OutOfOrder = r.outOfOrder
	base.CrossClockSkewSamples = r.crossClockSkewSamples
	base.OrderingViolations = append([]OrderingViolation(nil), r.ordering...)
	if r.malformed > 0 {
		base.Errors = append(base.Errors, fmt.Sprintf("malformed run messages=%d", r.malformed))
	}
	for _, message := range ledger {
		if message.RunID != runID {
			continue
		}
		accepted[message.ClientMid] = message
		obs := r.observed[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 mid, obs := range r.observed {
		if _, ok := accepted[mid]; !ok {
			base.UnexpectedMessages += bitsetCount(obs.bits)
		}
	}
	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
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	var count int64
	for _, message := range ledger {
		if message.RunID != runID || !stored[message.ClientMid] {
			continue
		}
		obs := r.observed[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 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}
}