// The package `m` contains models shared between the hub and peer programs.
package m

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/netip"

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

type PeerInitArgs struct {
	WGPubKey   []byte
	SignPubKey []byte
}

type PeerInitResp struct {
	PeerIP       byte
	Network      []byte
	LocalDomain  string
	NetworkState NetworkState
}

// Peer is the network membership record for a single peer, exchanged between
// the hub and peers. Addr4/Addr6 are the peer's public endpoint addresses (zero
// if it has none); Port is its WireGuard listen port, meaningful even for a
// non-public peer (it is the peer's own bind/beacon port).
type Peer struct {
	PeerIP     byte
	Name       string
	Addr4      netip.Addr // zero if none
	Addr6      netip.Addr // zero if none
	Port       uint16
	Relay      bool
	WGPubKey   wgtypes.Key
	SignPubKey [32]byte
}

// IsPublic reports whether the peer advertises at least one reachable endpoint.
func (p Peer) IsPublic() bool {
	return p.Addr4.IsValid() || p.Addr6.IsValid()
}

// Endpoint4 returns the IPv4 endpoint (addr+port), or the zero AddrPort if the
// peer has no IPv4 address.
func (p Peer) Endpoint4() netip.AddrPort {
	if !p.Addr4.IsValid() {
		return netip.AddrPort{}
	}
	return netip.AddrPortFrom(p.Addr4, p.Port)
}

// Endpoint6 returns the IPv6 endpoint (addr+port), or the zero AddrPort if the
// peer has no IPv6 address.
func (p Peer) Endpoint6() netip.AddrPort {
	if !p.Addr6.IsValid() {
		return netip.AddrPort{}
	}
	return netip.AddrPortFrom(p.Addr6, p.Port)
}

// PreferredEndpoint returns the IPv4 endpoint if present, else IPv6.
func (p Peer) PreferredEndpoint() netip.AddrPort {
	if ep := p.Endpoint4(); ep.IsValid() {
		return ep
	}
	return p.Endpoint6()
}

// peerJSON is the wire representation. netip.Addr fields round-trip as text
// strings automatically; only the fixed-size key arrays need base64 (otherwise
// encoding/json would emit them as arrays of numbers).
type peerJSON struct {
	PeerIP     byte
	Name       string
	Addr4      netip.Addr
	Addr6      netip.Addr
	Port       uint16
	Relay      bool
	WGPubKey   string
	SignPubKey string
}

func (p Peer) MarshalJSON() ([]byte, error) {
	return json.Marshal(peerJSON{
		PeerIP:     p.PeerIP,
		Name:       p.Name,
		Addr4:      p.Addr4,
		Addr6:      p.Addr6,
		Port:       p.Port,
		Relay:      p.Relay,
		WGPubKey:   base64.StdEncoding.EncodeToString(p.WGPubKey[:]),
		SignPubKey: base64.StdEncoding.EncodeToString(p.SignPubKey[:]),
	})
}

func (p *Peer) UnmarshalJSON(data []byte) error {
	var j peerJSON
	if err := json.Unmarshal(data, &j); err != nil {
		return err
	}
	wg, err := base64.StdEncoding.DecodeString(j.WGPubKey)
	if err != nil {
		return fmt.Errorf("decode WGPubKey: %w", err)
	}
	key, err := wgtypes.NewKey(wg)
	if err != nil {
		return fmt.Errorf("invalid WGPubKey: %w", err)
	}
	sign, err := base64.StdEncoding.DecodeString(j.SignPubKey)
	if err != nil {
		return fmt.Errorf("decode SignPubKey: %w", err)
	}
	if len(sign) != 32 {
		return fmt.Errorf("invalid SignPubKey length: %d", len(sign))
	}
	*p = Peer{
		PeerIP:     j.PeerIP,
		Name:       j.Name,
		Addr4:      j.Addr4,
		Addr6:      j.Addr6,
		Port:       j.Port,
		Relay:      j.Relay,
		WGPubKey:   key,
		SignPubKey: [32]byte(sign),
	}
	return nil
}

type NetworkState struct {
	Peers []Peer
}