package main

// LocalSend discovery, protocol section 3.
//
// Two halves, and both are needed for the phone to show this computer in its
// device list:
//
//	Listening  Join 224.0.0.167:53317 on every multicast-capable interface.
//	           When another device announces itself, answer it — over HTTP, by
//	           POSTing our own descriptor to its /api/localsend/v2/register.
//	           That is what puts us in ITS list.
//	Announcing Send our own announcement to the group so devices that were
//	           already running answer us. The reference implementation sends a
//	           burst of three, because a single UDP datagram is easily lost.
//
// Only IPv4 is implemented. The current LocalSend app additionally announces on
// an IPv6 group (ff12::fd3a:e420) as an extension on top of v2.2; that is
// listed as not implemented in the README rather than half-done here.

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"sort"
	"sync"
	"time"
)

// announceDelays matches the reference implementation's burst: a device that
// has just joined the network may not be listening yet when the first
// datagram goes out.
var announceDelays = []time.Duration{
	100 * time.Millisecond,
	500 * time.Millisecond,
	2000 * time.Millisecond,
}

// reannounceInterval keeps a long-running receiver visible to phones that
// wake up hours after it started. The protocol does not require it; the cost
// is one small datagram per interface per minute.
const reannounceInterval = 60 * time.Second

// maxDatagram is the read buffer. Announcements are a few hundred bytes;
// anything larger is not one.
const maxDatagram = 64 << 10

// peer is another LocalSend-speaking device we have heard from.
type peer struct {
	Alias       string
	Version     string
	DeviceModel string
	DeviceType  string
	Fingerprint string
	Address     string // IP the datagram or request came from
	Port        int
	Protocol    string
	Download    bool
	LastSeen    time.Time
}

// peerFromInfo turns a device descriptor received over HTTP into a peer,
// using the address the request came from rather than anything it claimed.
func peerFromInfo(info deviceInfo, address string) peer {
	port := info.Port
	if port == 0 {
		port = defaultLocalSendPort
	}
	return peer{
		Alias:       info.Alias,
		Version:     info.Version,
		DeviceModel: info.DeviceModel,
		DeviceType:  info.DeviceType,
		Fingerprint: info.Fingerprint,
		Address:     address,
		Port:        port,
		Protocol:    info.Protocol,
		Download:    info.Download,
		LastSeen:    time.Now(),
	}
}

func (p peer) baseURL() string {
	scheme := p.Protocol
	if scheme != "https" {
		scheme = "http"
	}
	return fmt.Sprintf("%s://%s", scheme, net.JoinHostPort(p.Address, fmt.Sprint(p.Port)))
}

// discovery runs the multicast half of the protocol for one receiver.
type discovery struct {
	rc *receiver

	// Port the group is joined on. Normally the same as the HTTP port.
	Port int

	// Announce says whether we shout as well as listen. `snapbeam peers`
	// announces too, so that silent devices answer it.
	Announce bool

	// OnPeer is called every time a device is heard from.
	OnPeer func(peer)

	mu       sync.Mutex
	listen   []*net.UDPConn
	send     []*net.UDPConn
	peers    map[string]peer
	stopped  bool
	stopOnce sync.Once
	wg       sync.WaitGroup
	client   *http.Client

	// done is closed by stop, which is what lets every wait in this file be
	// interrupted rather than run to completion.
	done chan struct{}
}

func newDiscovery(rc *receiver, port int) *discovery {
	return &discovery{
		rc:     rc,
		Port:   port,
		peers:  map[string]peer{},
		client: &http.Client{Timeout: 4 * time.Second},
		done:   make(chan struct{}),
	}
}

// groupAddr is the multicast socket address discovery uses.
func (d *discovery) groupAddr() *net.UDPAddr {
	return &net.UDPAddr{IP: multicastGroup, Port: d.Port}
}

// start binds the sockets and begins listening. It returns the interfaces it
// managed to join, and an error only when it could not join anything at all.
func (d *discovery) start() ([]string, error) {
	group := d.groupAddr()
	ifaces := multicastInterfaces()

	var joined []string
	var firstErr error
	for i := range ifaces {
		ifc := ifaces[i]
		conn, err := net.ListenMulticastUDP("udp4", &ifc, group)
		if err != nil {
			if firstErr == nil {
				firstErr = fmt.Errorf("%s: %w", ifc.Name, err)
			}
			continue
		}
		_ = conn.SetReadBuffer(maxDatagram)
		d.listen = append(d.listen, conn)
		joined = append(joined, ifc.Name)

		// A separate socket for sending, bound to this interface's own
		// address so the datagram leaves by the right door. The standard
		// library exposes no IP_MULTICAST_IF, and binding the source address
		// achieves the same routing decision.
		if src := firstIPv4Of(ifc); src != nil {
			if s, err := net.DialUDP("udp4", &net.UDPAddr{IP: src}, group); err == nil {
				d.send = append(d.send, s)
			}
		}
	}

	if len(d.listen) == 0 {
		if firstErr != nil {
			return nil, fmt.Errorf("cannot join the LocalSend multicast group %s: %w", group, firstErr)
		}
		return nil, fmt.Errorf("cannot join the LocalSend multicast group %s: no multicast-capable interface is up", group)
	}
	// A last-resort sender for the case where no interface offered a usable
	// source address; the kernel picks the route.
	if len(d.send) == 0 {
		if s, err := net.DialUDP("udp4", nil, group); err == nil {
			d.send = append(d.send, s)
		}
	}

	for _, c := range d.listen {
		d.wg.Add(1)
		go d.readLoop(c)
	}
	if d.Announce {
		d.wg.Add(1)
		go d.announceLoop()
	}
	return joined, nil
}

func firstIPv4Of(ifc net.Interface) net.IP {
	addrs, err := ifc.Addrs()
	if err != nil {
		return nil
	}
	for _, a := range addrs {
		ip, _, err := net.ParseCIDR(a.String())
		if err != nil {
			continue
		}
		if ip4 := ip.To4(); ip4 != nil {
			return ip4
		}
	}
	return nil
}

func (d *discovery) stop() {
	d.stopOnce.Do(func() {
		d.mu.Lock()
		d.stopped = true
		d.mu.Unlock()
		close(d.done)
		for _, c := range d.listen {
			_ = c.Close()
		}
		for _, c := range d.send {
			_ = c.Close()
		}
		d.wg.Wait()
	})
}

func (d *discovery) isStopped() bool {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.stopped
}

// announceBytes is the exact datagram SnapBeam puts on the wire.
func (d *discovery) announceBytes() ([]byte, error) {
	msg := announcement{deviceInfo: d.rc.info(), Announce: true}
	return json.Marshal(msg)
}

func (d *discovery) announceOnce() {
	payload, err := d.announceBytes()
	if err != nil {
		return
	}
	group := d.groupAddr()
	for _, s := range d.send {
		// The socket is connected, so a plain Write goes to the group.
		if _, err := s.Write(payload); err != nil {
			// Fall back to an explicit destination; some stacks refuse a
			// connected multicast socket.
			_, _ = s.WriteToUDP(payload, group)
		}
	}
}

// announceLoop sends the opening burst and then keeps a long-running receiver
// visible. Every wait is on a select with the stop channel: a plain sleep or a
// bare ticker would hold stop() for up to a full re-announce interval, which
// is exactly how a program ends up taking a minute to answer Ctrl+C.
func (d *discovery) announceLoop() {
	defer d.wg.Done()
	wait := func(dur time.Duration) bool {
		t := time.NewTimer(dur)
		defer t.Stop()
		select {
		case <-d.done:
			return false
		case <-t.C:
			return true
		}
	}
	for _, delay := range announceDelays {
		if !wait(delay) {
			return
		}
		d.announceOnce()
	}
	for wait(reannounceInterval) {
		d.announceOnce()
	}
}

func (d *discovery) readLoop(c *net.UDPConn) {
	defer d.wg.Done()
	buf := make([]byte, maxDatagram)
	for {
		n, src, err := c.ReadFromUDP(buf)
		if err != nil {
			return // socket closed, or gave up
		}
		d.handleDatagram(buf[:n], src)
	}
}

// handleDatagram parses one announcement and, when it is somebody else's
// announcement, answers it over HTTP.
func (d *discovery) handleDatagram(data []byte, src *net.UDPAddr) {
	var msg announcement
	// A message with no "announce" key decodes to false, which is not what the
	// protocol means: the reference implementation omits it on reply and the
	// document shows it present on announcements. Decode into a map first to
	// tell "absent" from "false".
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(data, &raw); err != nil {
		return
	}
	if err := json.Unmarshal(data, &msg); err != nil {
		return
	}
	if msg.Alias == "" || msg.Fingerprint == "" {
		return
	}
	if msg.Fingerprint == d.rc.Fingerprint {
		return // our own datagram, looped back
	}

	isAnnouncement := msg.Announce
	if _, present := raw["announce"]; !present {
		if _, presentV1 := raw["announcement"]; presentV1 {
			// Protocol v1 spelled the flag "announcement".
			var v1 struct {
				Announcement bool `json:"announcement"`
			}
			_ = json.Unmarshal(data, &v1)
			isAnnouncement = v1.Announcement
		} else {
			isAnnouncement = true
		}
	}

	port := msg.Port
	if port == 0 {
		port = defaultLocalSendPort
	}
	p := peer{
		Alias:       msg.Alias,
		Version:     msg.Version,
		DeviceModel: msg.DeviceModel,
		DeviceType:  msg.DeviceType,
		Fingerprint: msg.Fingerprint,
		Address:     src.IP.String(),
		Port:        port,
		Protocol:    msg.Protocol,
		Download:    msg.Download,
		LastSeen:    time.Now(),
	}
	d.mu.Lock()
	_, known := d.peers[p.Fingerprint]
	d.peers[p.Fingerprint] = p
	d.mu.Unlock()

	if d.OnPeer != nil && !known {
		d.OnPeer(p)
	}

	// Section 3.1: a response is only triggered when the message is an
	// announcement. Answering a response would put two receivers in a loop.
	if isAnnouncement {
		go d.registerWith(p)
	}
}

// registerWith is the HTTP half of the discovery handshake: we tell the peer
// who we are, which is what makes SnapBeam appear in its device list.
func (d *discovery) registerWith(p peer) {
	body, err := json.Marshal(d.rc.info())
	if err != nil {
		return
	}
	url := p.baseURL() + "/api/localsend/v2/register"
	req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return
	}
	req.Header.Set("Content-Type", "application/json")
	resp, err := d.client.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	_, _ = readAllLimited(resp.Body, 64<<10)
}

// record adds a peer learned by some route other than multicast, notably an
// HTTP registration answering our own announcement.
func (d *discovery) record(p peer) bool {
	if p.Fingerprint == "" || p.Fingerprint == d.rc.Fingerprint {
		return false
	}
	d.mu.Lock()
	_, known := d.peers[p.Fingerprint]
	d.peers[p.Fingerprint] = p
	d.mu.Unlock()
	return !known
}

// knownPeers returns everything heard from so far, most recent first.
func (d *discovery) knownPeers() []peer {
	d.mu.Lock()
	out := make([]peer, 0, len(d.peers))
	for _, p := range d.peers {
		out = append(out, p)
	}
	d.mu.Unlock()
	sort.SliceStable(out, func(i, j int) bool { return out[i].LastSeen.After(out[j].LastSeen) })
	return out
}