repeater/internal/station/hostapd/backend.go

package hostapd

import (
	"bufio"
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/nemunaire/repeater/internal/station/backend"
)

// Backend implements StationBackend using hostapd_cli
type Backend struct {
	interfaceName string
	hostapdCLI    string // Path to hostapd_cli executable

	stations  map[string]*HostapdStation // Key: MAC address
	callbacks backend.EventCallbacks

	mu       sync.RWMutex
	running  bool
	stopCh   chan struct{}
	stopOnce sync.Once

	// IP / hostname correlation - populated by periodic DHCP lease correlation
	ipByMAC       map[string]string // MAC -> IP mapping
	hostnameByMAC map[string]string // MAC -> hostname mapping

	correlator *DHCPCorrelator
}

// NewBackend creates a new hostapd backend
func NewBackend() *Backend {
	return &Backend{
		stations:      make(map[string]*HostapdStation),
		ipByMAC:       make(map[string]string),
		hostnameByMAC: make(map[string]string),
		hostapdCLI:    "hostapd_cli",
		stopCh:        make(chan struct{}),
	}
}

// Initialize initializes the hostapd backend
func (b *Backend) Initialize(config backend.BackendConfig) error {
	b.mu.Lock()
	defer b.mu.Unlock()

	b.interfaceName = config.InterfaceName
	if b.interfaceName == "" {
		b.interfaceName = "wlan1" // Default AP interface
	}

	// Check if hostapd_cli is available
	if _, err := exec.LookPath(b.hostapdCLI); err != nil {
		return fmt.Errorf("hostapd_cli not found in PATH: %w", err)
	}

	// Verify we can communicate with hostapd
	if err := b.runCommand("ping"); err != nil {
		return fmt.Errorf("failed to communicate with hostapd on interface %s: %w", b.interfaceName, err)
	}

	log.Printf("Hostapd backend initialized for interface %s", b.interfaceName)

	// Load initial station list
	if err := b.loadStations(); err != nil {
		log.Printf("Warning: Failed to load initial stations: %v", err)
	}

	// Hostapd only knows MACs. Pair every station with its IP/hostname by
	// running a background correlator that polls the DHCP lease file and
	// the ARP table. Without this, the device list shows MACs only.
	b.correlator = NewDHCPCorrelator(b, config.DHCPLeasesPath, config.ARPTablePath)
	b.correlator.Start()

	return nil
}

// Close cleans up backend resources
func (b *Backend) Close() error {
	if b.correlator != nil {
		b.correlator.Stop()
	}
	b.StopEventMonitoring()
	return nil
}

// runCommand executes a hostapd_cli command and returns the output
func (b *Backend) runCommand(args ...string) error {
	cmdArgs := []string{"-i", b.interfaceName}
	cmdArgs = append(cmdArgs, args...)
	cmd := exec.Command(b.hostapdCLI, cmdArgs...)
	return cmd.Run()
}

// runCommandOutput executes a hostapd_cli command and returns the output
func (b *Backend) runCommandOutput(args ...string) (string, error) {
	cmdArgs := []string{"-i", b.interfaceName}
	cmdArgs = append(cmdArgs, args...)
	cmd := exec.Command(b.hostapdCLI, cmdArgs...)
	out, err := cmd.Output()
	if err != nil {
		return "", err
	}
	return string(out), nil
}

// GetStations returns all connected stations
func (b *Backend) GetStations() ([]backend.Station, error) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	stations := make([]backend.Station, 0, len(b.stations))
	for mac, hs := range b.stations {
		station := b.convertStation(mac, hs)
		stations = append(stations, station)
	}

	return stations, nil
}

// StartEventMonitoring starts monitoring for station events via polling
func (b *Backend) StartEventMonitoring(callbacks backend.EventCallbacks) error {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.running {
		return nil
	}

	b.callbacks = callbacks
	b.running = true

	// Start polling goroutine
	go b.pollStations()

	log.Printf("Hostapd event monitoring started (polling mode)")
	return nil
}

// StopEventMonitoring stops event monitoring. Safe to call multiple times —
// without sync.Once, racing callers could each pass the running guard before
// either closed the channel, panicking on the second close.
func (b *Backend) StopEventMonitoring() {
	b.mu.Lock()
	if !b.running {
		b.mu.Unlock()
		return
	}
	b.running = false
	b.mu.Unlock()

	b.stopOnce.Do(func() { close(b.stopCh) })
}

// SupportsRealTimeEvents returns false (hostapd_cli uses polling)
func (b *Backend) SupportsRealTimeEvents() bool {
	return false
}

// pollStations periodically polls for station changes
func (b *Backend) pollStations() {
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-b.stopCh:
			return
		case <-ticker.C:
			if err := b.checkStationChanges(); err != nil {
				log.Printf("Error polling stations: %v", err)
			}
		}
	}
}

// checkStationChanges checks for station connect/disconnect events
func (b *Backend) checkStationChanges() error {
	// Get current stations from hostapd
	currentStations, err := b.fetchStations()
	if err != nil {
		return err
	}

	now := time.Now()

	b.mu.Lock()
	defer b.mu.Unlock()

	// Build a map of current MACs
	currentMACs := make(map[string]bool)
	for mac := range currentStations {
		currentMACs[mac] = true
	}

	// Reconcile current poll with cached state
	for mac, station := range currentStations {
		if existing, exists := b.stations[mac]; exists {
			// Refresh signal/byte counters in place — without this, GetStations
			// would keep returning the values from the very first poll.
			existing.Signal = station.Signal
			existing.RxBytes = station.RxBytes
			existing.TxBytes = station.TxBytes
			continue
		}
		// New station connected
		station.FirstSeen = now
		b.stations[mac] = station
		if cb := b.callbacks.OnStationConnected; cb != nil {
			st := b.convertStation(mac, station)
			backend.SafeGo("OnStationConnected", func() { cb(st) })
		}
		log.Printf("Station connected: %s", mac)
	}

	// Check for removed stations
	for mac := range b.stations {
		if !currentMACs[mac] {
			// Station disconnected
			delete(b.stations, mac)
			delete(b.ipByMAC, mac)
			if cb := b.callbacks.OnStationDisconnected; cb != nil {
				macCopy := mac
				backend.SafeGo("OnStationDisconnected", func() { cb(macCopy) })
			}
			log.Printf("Station disconnected: %s", mac)
		}
	}

	return nil
}

// loadStations loads the initial list of stations from hostapd
func (b *Backend) loadStations() error {
	stations, err := b.fetchStations()
	if err != nil {
		return err
	}

	now := time.Now()
	for _, st := range stations {
		// We can't tell when these stations actually connected; treat the
		// daemon-start time as a lower bound rather than leaving it zero.
		st.FirstSeen = now
	}
	b.stations = stations
	log.Printf("Loaded %d initial stations from hostapd", len(b.stations))
	return nil
}

// fetchStations fetches all stations using hostapd_cli all_sta command
func (b *Backend) fetchStations() (map[string]*HostapdStation, error) {
	output, err := b.runCommandOutput("all_sta")
	if err != nil {
		return nil, fmt.Errorf("failed to get stations: %w", err)
	}

	return b.parseAllStaOutput(output), nil
}

// parseAllStaOutput parses the output of "hostapd_cli all_sta"
func (b *Backend) parseAllStaOutput(output string) map[string]*HostapdStation {
	stations := make(map[string]*HostapdStation)
	scanner := bufio.NewScanner(bytes.NewBufferString(output))

	var currentMAC string
	var currentStation *HostapdStation

	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" {
			continue
		}

		// Check if this is a MAC address line (starts the station block)
		if !strings.Contains(line, "=") && len(line) == 17 && strings.Count(line, ":") == 5 {
			// Save previous station if exists
			if currentMAC != "" && currentStation != nil {
				stations[currentMAC] = currentStation
			}
			// Start new station
			currentMAC = strings.ToLower(line)
			currentStation = &HostapdStation{}
			continue
		}

		// Parse key=value pairs
		if currentStation != nil && strings.Contains(line, "=") {
			parts := strings.SplitN(line, "=", 2)
			if len(parts) != 2 {
				continue
			}
			key := strings.TrimSpace(parts[0])
			value := strings.TrimSpace(parts[1])

			switch key {
			case "signal":
				if v, err := strconv.Atoi(value); err == nil {
					currentStation.Signal = int32(v)
				}
			case "rx_bytes":
				if v, err := strconv.ParseUint(value, 10, 64); err == nil {
					currentStation.RxBytes = v
				}
			case "tx_bytes":
				if v, err := strconv.ParseUint(value, 10, 64); err == nil {
					currentStation.TxBytes = v
				}
			}
		}
	}

	// Save last station
	if currentMAC != "" && currentStation != nil {
		stations[currentMAC] = currentStation
	}

	return stations
}

// convertStation converts HostapdStation to backend.Station
func (b *Backend) convertStation(mac string, hs *HostapdStation) backend.Station {
	ip := b.ipByMAC[mac]
	hostname := b.hostnameByMAC[mac]

	connectedAt := hs.FirstSeen
	if connectedAt.IsZero() {
		connectedAt = time.Now()
	}

	return backend.Station{
		MAC:         mac,
		IP:          ip,
		Hostname:    hostname,
		Type:        backend.GuessDeviceType(hostname, mac),
		Vendor:      backend.LookupVendor(mac),
		Signal:      hs.Signal,
		RxBytes:     hs.RxBytes,
		TxBytes:     hs.TxBytes,
		ConnectedAt: connectedAt,
	}
}

// UpdateLeaseInfo updates MAC -> IP and MAC -> hostname mappings from an
// external source (e.g., DHCP lease file). Called periodically by the
// correlator. An empty hostname leaves the existing one in place — DHCP
// leases sometimes drop the hostname between renewals.
func (b *Backend) UpdateLeaseInfo(macToIP, macToHostname map[string]string) {
	b.mu.Lock()
	defer b.mu.Unlock()

	updated := make(map[string]bool)

	for mac, ip := range macToIP {
		if oldIP, exists := b.ipByMAC[mac]; !exists || oldIP != ip {
			updated[mac] = true
		}
		b.ipByMAC[mac] = ip
	}

	for mac, hostname := range macToHostname {
		if hostname == "" {
			continue
		}
		if oldName, exists := b.hostnameByMAC[mac]; !exists || oldName != hostname {
			updated[mac] = true
		}
		b.hostnameByMAC[mac] = hostname
	}

	// Trigger update callbacks for stations whose info changed
	for mac := range updated {
		if station, exists := b.stations[mac]; exists {
			if cb := b.callbacks.OnStationUpdated; cb != nil {
				st := b.convertStation(mac, station)
				backend.SafeGo("OnStationUpdated", func() { cb(st) })
			}
		}
	}
}