package peer

import (
	"net/netip"
	"time"
)

type sentProbe struct {
	SentAt time.Time
	Addr   netip.AddrPort
}

type stateClientRelayed2 struct {
	*peerData
	lastSeen time.Time
	syn      packetSyn
	probes   map[uint64]sentProbe
}

func enterStateClientRelayed2(data *peerData) peerState {
	data.staged.Relay = false
	data.staged.Direct = false
	data.staged.DirectAddr = netip.AddrPort{}
	data.publish(data.staged)

	state := &stateClientRelayed2{
		peerData: data,
		lastSeen: time.Now(),
		syn: packetSyn{
			TraceID:       newTraceID(),
			SharedKey:     data.staged.DataCipher.Key(),
			Direct:        false,
			PossibleAddrs: data.pubAddrs.Get(),
		},
		probes: map[uint64]sentProbe{},
	}

	state.Send(state.staged, state.syn)

	data.pingTimer.Reset(pingInterval)

	state.logf("==> ClientRelayed")
	return state
}

func (s *stateClientRelayed2) logf(str string, args ...any) {
	s.peerData.logf("CLNT | "+str, args...)
}

func (s *stateClientRelayed2) OnMsg(raw any) peerState {
	switch msg := raw.(type) {
	case peerUpdateMsg:
		return initPeerState(s.peerData, msg.Peer)
	case controlMsg[packetAck]:
		return s.onAck(msg)
	case controlMsg[packetProbe]:
		return s.onProbe(msg)
	case controlMsg[packetLocalDiscovery]:
		return s.onLocalDiscovery(msg)
	case pingTimerMsg:
		return s.onPingTimer()
	default:
		s.logf("Ignoring message: %v", raw)
		return s
	}
}

func (s *stateClientRelayed2) onAck(msg controlMsg[packetAck]) peerState {
	if msg.Packet.TraceID != s.syn.TraceID {
		return s
	}

	s.lastSeen = time.Now()

	if !s.staged.Up {
		s.staged.Up = true
		s.publish(s.staged)
		s.logf("Got ACK.")
	}

	s.pubAddrs.Store(msg.Packet.ToAddr)

	for _, addr := range msg.Packet.PossibleAddrs {
		if !addr.IsValid() {
			break
		}
		s.sendProbeTo(addr)
	}

	s.cleanProbes()

	return s
}

func (s *stateClientRelayed2) onPingTimer() peerState {
	if time.Since(s.lastSeen) > timeoutInterval {
		if s.staged.Up {
			s.logf("Timeout.")
		}
		return initPeerState(s.peerData, s.peer)
	}

	s.Send(s.staged, s.syn)
	return s
}

func (s *stateClientRelayed2) onProbe(msg controlMsg[packetProbe]) peerState {
	s.cleanProbes()

	sent, ok := s.probes[msg.Packet.TraceID]
	if !ok {
		return s
	}

	s.logf("Successful probe.")
	return enterStateClientDirect2(s.peerData, sent.Addr)
}

func (s *stateClientRelayed2) onLocalDiscovery(msg controlMsg[packetLocalDiscovery]) peerState {
	// The source port will be the multicast port, so we'll have to
	// construct the correct address using the peer's listed port.
	addr := netip.AddrPortFrom(msg.SrcAddr.Addr(), s.peer.Port)
	s.sendProbeTo(addr)
	return s
}

func (s *stateClientRelayed2) cleanProbes() {
	for key, sent := range s.probes {
		if time.Since(sent.SentAt) > pingInterval {
			delete(s.probes, key)
		}
	}
}

func (s *stateClientRelayed2) sendProbeTo(addr netip.AddrPort) {
	probe := packetProbe{TraceID: newTraceID()}
	s.probes[probe.TraceID] = sentProbe{
		SentAt: time.Now(),
		Addr:   addr,
	}
	s.SendTo(probe, addr)
}