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 enterStateClient(data *peerData) peerState {
	ip, ipValid := netip.AddrFromSlice(data.peer.PublicIP)

	data.staged.Relay = data.peer.Relay && ipValid
	data.staged.Direct = ipValid
	data.staged.DirectAddr = netip.AddrPortFrom(ip, data.peer.Port)
	data.publish(data.staged)

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

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

	data.pingTimer.Reset(pingInterval)

	state.logf("==> Client")
	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]:
		s.onAck(msg)
	case controlMsg[packetProbe]:
		return s.onProbe(msg)
	case controlMsg[packetLocalDiscovery]:
		s.onLocalDiscovery(msg)
	case pingTimerMsg:
		return s.onPingTimer()
	default:
		s.logf("Ignoring message: %v", raw)
	}
	return s
}

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

	s.lastSeen = time.Now()

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

	if s.staged.Direct {
		s.pubAddrs.Store(msg.Packet.ToAddr)
		return
	}

	// Relayed below.

	s.cleanProbes()

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

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 {
	if s.staged.Direct {
		return s
	}

	s.cleanProbes()

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

	s.staged.Direct = true
	s.staged.DirectAddr = sent.Addr
	s.publish(s.staged)

	s.syn.TraceID = newTraceID()
	s.syn.Direct = true
	s.Send(s.staged, s.syn)

	s.logf("Successful probe.")
	return s
}

func (s *stateClientRelayed2) onLocalDiscovery(msg controlMsg[packetLocalDiscovery]) {
	if s.staged.Direct {
		return
	}

	// 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)
}

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.logf("Probing %v...", addr)
	s.SendTo(probe, addr)
}