package bcap

import "time"

type trackerTCPState struct {
	firstSeen   time.Time
	lastSeen    time.Time
	packetCount uint64
	byteCount   uint64

	seq          uint32
	ack          uint32
	window       uint16
	payload      int
	finState     bool
	synState     bool
	isFirst      bool
	state        uint8
	segments     [trackedTCPSegments]tcpSegmentRange
	segmentCount int
	segmentNext  int
}

func (t *Tracker) observeTCP(packet Packet, flow FlowRef) (HintSet, error) {
	tcp := packet.Transport.TCP
	if tcp == nil {
		return HintSet{}, NewParseError(ErrTypeTransport, "TCP", "missing tcp facts", nil)
	}

	forwardKey := flow.Forward.StableString()
	reverseKey := flow.Reverse.StableString()
	if forwardKey == "" {
		forwardKey = flow.Stable
	}
	if reverseKey == "" {
		reverseKey = stableFlowKey(flow.Reverse)
	}

	t.mu.Lock()
	defer t.mu.Unlock()

	lastState, exists := t.tcpStates[forwardKey]
	if !exists {
		lastState = trackerTCPState{
			firstSeen: packet.Meta.Timestamp,
			lastSeen:  packet.Meta.Timestamp,
			isFirst:   true,
		}
	}
	lastReverse := t.tcpStates[reverseKey]

	payloadLen := tcp.Payload
	seqEnd := tcpSeqAdvanceFacts(tcp.Seq, tcp.SYN, tcp.FIN, payloadLen)
	state := StateUnknown
	hintOpts := tcpHintOptions{}
	connectionClosed := false

	if tcp.RST {
		state = StateTcpRst
		connectionClosed = true
		delete(t.tcpStates, forwardKey)
		delete(t.tcpStates, reverseKey)
	} else if tcp.SYN && !tcp.ACK {
		state = StateTcpConnect1
	} else if tcp.SYN && tcp.ACK {
		state = StateTcpConnect2
	} else if tcp.ACK && !tcp.FIN {
		state, hintOpts, connectionClosed = classifyTrackerAckNoFIN(packet, tcp, seqEnd, lastState, lastReverse)
		if connectionClosed {
			delete(t.tcpStates, forwardKey)
			delete(t.tcpStates, reverseKey)
		}
	} else if tcp.ACK && tcp.FIN {
		state = classifyTrackerAckFIN(tcp, lastState, lastReverse)
	}

	if !connectionClosed {
		next := trackerTCPState{
			firstSeen:    firstSeenOrNow(lastState.firstSeen, packet.Meta.Timestamp),
			lastSeen:     packet.Meta.Timestamp,
			packetCount:  lastState.packetCount + 1,
			byteCount:    lastState.byteCount + packetWireLength(packet),
			seq:          tcp.Seq,
			ack:          tcp.Ack,
			window:       tcp.Window,
			payload:      payloadLen,
			finState:     tcp.FIN,
			synState:     tcp.SYN,
			isFirst:      false,
			state:        state,
			segments:     lastState.segments,
			segmentCount: lastState.segmentCount,
			segmentNext:  lastState.segmentNext,
		}
		next.rememberSegment(tcp.Seq, seqEnd)
		t.tcpStates[forwardKey] = next
	}

	return newTCPHints(packet, state, hintOpts), nil
}

func (f FlowKey) StableString() string {
	return stableFlowKey(f)
}

func firstSeenOrNow(firstSeen, fallback time.Time) time.Time {
	if firstSeen.IsZero() {
		return fallback
	}
	return firstSeen
}

func packetWireLength(packet Packet) uint64 {
	switch {
	case packet.Meta.Length > 0:
		return uint64(packet.Meta.Length)
	case packet.Meta.CaptureLength > 0:
		return uint64(packet.Meta.CaptureLength)
	case packet.Raw.Packet != nil:
		return uint64(len(packet.Raw.Packet.Data()))
	default:
		return 0
	}
}

func tcpSeqAdvanceFacts(seq uint32, syn, fin bool, payloadLen int) uint32 {
	span := payloadLen
	if syn {
		span++
	}
	if fin {
		span++
	}
	return tcpSeqAdd(seq, uint32(span))
}

func classifyTrackerAckNoFIN(
	packet Packet,
	tcp *TCPFacts,
	seqEnd uint32,
	lastState, lastReverse trackerTCPState,
) (uint8, tcpHintOptions, bool) {
	if isTrackerHandshakeAck(lastState, lastReverse, tcp.Ack) {
		return StateTcpConnect3, tcpHintOptions{}, false
	}
	if tcp.CWR {
		return StateTcpCwr, tcpHintOptions{}, false
	}
	if tcp.ECE {
		return StateTcpEce, tcpHintOptions{}, false
	}
	if state, opts, ok := classifyTrackerTCPKeepalive(packet, tcp, lastState, lastReverse); ok {
		return state, opts, false
	}
	if lastState.hasSeenSegment(tcp.Seq, seqEnd) {
		return StateTcpRetransmit, tcpHintOptions{}, false
	}
	if lastReverse.finState && lastState.finState {
		return StateTcpDisconnect4, tcpHintOptions{}, true
	}
	if lastReverse.finState && tcpSeqAdd(lastReverse.seq, 1) == tcp.Ack {
		return StateTcpDisconnect2, tcpHintOptions{}, false
	}
	return StateTcpAckOk, tcpHintOptions{}, false
}

func classifyTrackerAckFIN(tcp *TCPFacts, lastState, lastReverse trackerTCPState) uint8 {
	if !lastReverse.finState {
		return StateTcpDisconnect1
	}
	if lastReverse.finState && tcpSeqAdd(lastReverse.seq, 1) == tcp.Ack &&
		lastState.ack == tcp.Ack && lastState.seq == tcp.Seq {
		return StateTcpDisconnect3
	}
	return StateTcpDisconnect23
}

func isTrackerHandshakeAck(lastState, lastReverse trackerTCPState, ack uint32) bool {
	if lastReverse.state != StateTcpConnect2 {
		return false
	}
	if tcpSeqAdd(lastReverse.seq, 1) != ack {
		return false
	}
	return lastState.state == StateTcpConnect1 && lastState.synState
}

func classifyTrackerTCPKeepalive(packet Packet, tcp *TCPFacts, lastState, lastReverse trackerTCPState) (uint8, tcpHintOptions, bool) {
	if isTrackerRepeatedKeepaliveProbe(packet, tcp, lastState, lastReverse) {
		return StateTcpKeepalive, tcpHintOptions{}, true
	}
	if lastReverse.matchesKeepaliveResponse(packet, tcp) {
		return StateTcpKeepalive, tcpHintOptions{keepaliveResponse: true}, true
	}
	if tcp.Seq == tcpSeqPrev(lastReverse.ack) &&
		tcp.Seq == tcpSeqPrev(tcpSeqAdd(lastState.seq, uint32(lastState.payload))) {
		return StateTcpKeepalive, tcpHintOptions{}, true
	}
	return StateUnknown, tcpHintOptions{}, false
}

func isTrackerRepeatedKeepaliveProbe(packet Packet, tcp *TCPFacts, lastState, lastReverse trackerTCPState) bool {
	if lastState.isFirst {
		return false
	}
	if tcp == nil || !tcp.ACK || tcp.SYN || tcp.FIN || tcp.RST || tcp.ECE || tcp.CWR {
		return false
	}
	if tcp.Payload != 0 {
		return false
	}
	if lastState.synState || lastState.finState || lastState.payload != 0 {
		return false
	}
	switch lastState.state {
	case StateTcpAckOk, StateTcpKeepalive, StateTcpConnect3, StateTcpDisconnect2:
	default:
		return false
	}
	if tcp.Seq != lastState.seq || tcp.Ack != lastState.ack || tcp.Window != lastState.window {
		return false
	}
	if packet.Meta.Timestamp.Sub(lastState.lastSeen) < repeatedKeepaliveMinInterval {
		return false
	}
	if !lastReverse.lastSeen.IsZero() && lastReverse.lastSeen.After(lastState.lastSeen) {
		return false
	}
	return true
}

func (s trackerTCPState) matchesKeepaliveResponse(packet Packet, tcp *TCPFacts) bool {
	if tcp == nil || s.state != StateTcpKeepalive {
		return false
	}
	if !tcp.ACK || tcp.SYN || tcp.FIN || tcp.RST || tcp.ECE || tcp.CWR {
		return false
	}
	if tcp.Payload != 0 {
		return false
	}
	if packet.Meta.Timestamp.Sub(s.lastSeen) > keepaliveResponseMaxDelay {
		return false
	}
	return tcp.Seq == s.ack && tcp.Ack == tcpSeqAdd(s.seq, 1)
}

func (s trackerTCPState) hasSeenSegment(seq, end uint32) bool {
	if !tcpSeqLess(seq, end) {
		return false
	}
	for i := 0; i < s.segmentCount; i++ {
		seg := s.segments[i]
		if !tcpSeqLess(seg.seq, seg.end) {
			continue
		}
		if !tcpSeqLess(seq, seg.seq) && tcpSeqLEQ(end, seg.end) {
			return true
		}
	}
	return false
}

func (s *trackerTCPState) rememberSegment(seq, end uint32) {
	if !tcpSeqLess(seq, end) {
		return
	}
	if s.hasSeenSegment(seq, end) {
		return
	}
	if s.segmentCount < trackedTCPSegments {
		s.segments[s.segmentCount] = tcpSegmentRange{seq: seq, end: end}
		s.segmentCount++
		return
	}
	s.segments[s.segmentNext] = tcpSegmentRange{seq: seq, end: end}
	s.segmentNext = (s.segmentNext + 1) % trackedTCPSegments
}