package peer

import (
	"encoding/binary"
	"fmt"
	"log"
	"net"
	"net/netip"

	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
)

const (
	controlPort    = uint16(4561)
	controlMsgLen  = 7 // 1 type + 4 IPv4 + 2 port
	msgYourEndpt   = uint8(1)
	msgMyEndpt     = uint8(2)
)

type ControlServer struct {
	localPeerIP byte
	network     []byte
	conn        *net.UDPConn
	hp          *HolePunch
	netName     string
}

func NewControlServer(g Globals, hp *HolePunch, netName string) (*ControlServer, error) {
	vpnIP := netip.AddrFrom4([4]byte{
		g.Network[0], g.Network[1], g.Network[2], g.LocalPeerIP,
	})
	listenAddr := net.UDPAddrFromAddrPort(netip.AddrPortFrom(vpnIP, controlPort))
	conn, err := net.ListenUDP("udp4", listenAddr)
	if err != nil {
		return nil, fmt.Errorf("listen on control port: %w", err)
	}
	return &ControlServer{
		localPeerIP: g.LocalPeerIP,
		network:     g.Network,
		conn:        conn,
		hp:          hp,
		netName:     netName,
	}, nil
}

func (cs *ControlServer) Run() {
	buf := make([]byte, 64)
	for {
		n, src, err := cs.conn.ReadFromUDPAddrPort(buf)
		if err != nil {
			log.Printf("[Control] read: %v", err)
			continue
		}
		if n < controlMsgLen {
			continue
		}
		cs.handle(buf[:n], src)
	}
}

func (cs *ControlServer) handle(msg []byte, src netip.AddrPort) {
	msgType, ep, ok := decodeControlMsg(msg)
	if !ok {
		return
	}
	switch msgType {
	case msgYourEndpt:
		cs.onYourEndpoint(ep)
	case msgMyEndpt:
		cs.onMyEndpoint(src, ep)
	}
}

// onYourEndpoint is called when the relay tells us our external WG endpoint.
// We broadcast MsgMyEndpoint to all known peers so they can attempt direct
// connections to us.
func (cs *ControlServer) onYourEndpoint(ourEndpoint netip.AddrPort) {
	log.Printf("[Control] external endpoint: %v", ourEndpoint)

	state, err := loadNetworkState(cs.netName)
	if err != nil {
		log.Printf("[Control] load state: %v", err)
		return
	}

	for _, p := range state.Peers {
		if p == nil || p.PeerIP == cs.localPeerIP || len(p.WGPubKey) != wgtypes.KeyLen {
			continue
		}
		peerVPNIP := netip.AddrFrom4([4]byte{
			cs.network[0], cs.network[1], cs.network[2], p.PeerIP,
		})
		cs.sendMsg(msgMyEndpt, peerVPNIP, ourEndpoint)
	}
}

// onMyEndpoint is called when a peer tells us their external WG endpoint.
// We start a probe-before-commit attempt to reach them directly.
func (cs *ControlServer) onMyEndpoint(src netip.AddrPort, theirEndpoint netip.AddrPort) {
	peerIPByte := src.Addr().As4()[3]

	state, err := loadNetworkState(cs.netName)
	if err != nil {
		log.Printf("[Control] load state: %v", err)
		return
	}

	peer := state.Peers[peerIPByte]
	if peer == nil || len(peer.WGPubKey) != wgtypes.KeyLen {
		return
	}

	pubKey, err := wgtypes.NewKey(peer.WGPubKey)
	if err != nil {
		return
	}

	cs.hp.OnEndpointLearned(peerIPByte, pubKey, theirEndpoint, false)
}

// SendYourEndpoint sends MsgYourEndpoint to a peer, informing them of their
// own external WG endpoint. Used by the relay's endpoint reporter (Phase 6).
func (cs *ControlServer) SendYourEndpoint(peerVPNIP netip.Addr, theirEndpoint netip.AddrPort) {
	cs.sendMsg(msgYourEndpt, peerVPNIP, theirEndpoint)
}

func (cs *ControlServer) sendMsg(msgType uint8, peerVPNIP netip.Addr, ep netip.AddrPort) {
	dst := net.UDPAddrFromAddrPort(netip.AddrPortFrom(peerVPNIP, controlPort))
	if _, err := cs.conn.WriteTo(encodeControlMsg(msgType, ep), dst); err != nil {
		log.Printf("[Control] send to %v: %v", peerVPNIP, err)
	}
}

func encodeControlMsg(msgType uint8, ep netip.AddrPort) []byte {
	msg := make([]byte, controlMsgLen)
	msg[0] = msgType
	a4 := ep.Addr().Unmap().As4()
	copy(msg[1:5], a4[:])
	binary.BigEndian.PutUint16(msg[5:7], ep.Port())
	return msg
}

func decodeControlMsg(msg []byte) (msgType uint8, ep netip.AddrPort, ok bool) {
	if len(msg) < controlMsgLen {
		return 0, netip.AddrPort{}, false
	}
	ip := netip.AddrFrom4([4]byte(msg[1:5]))
	port := binary.BigEndian.Uint16(msg[5:7])
	return msg[0], netip.AddrPortFrom(ip, port), true
}