ha-n95-local-control/cmd/n95bridge/main.go

package main

import (
	"context"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"git.i3omb.com/gronod/ha-n95-local-control/internal/config"
	"git.i3omb.com/gronod/ha-n95-local-control/internal/ctl"
	"git.i3omb.com/gronod/ha-n95-local-control/internal/ha"
	"git.i3omb.com/gronod/ha-n95-local-control/internal/httpx"
	"git.i3omb.com/gronod/ha-n95-local-control/internal/mqttbridge"
	"git.i3omb.com/gronod/ha-n95-local-control/internal/robot"
	"git.i3omb.com/gronod/ha-n95-local-control/internal/session"
	"git.i3omb.com/gronod/ha-n95-local-control/internal/xmpp"
)

var (
	// shutdownBridge is assigned by phase 04.
	shutdownBridge func(context.Context) error
	// shutdownXMPP is assigned by phase 02.
	shutdownXMPP func(context.Context) error
	// shutdownTimeout is the graceful shutdown budget.
	shutdownTimeout = 5 * time.Second
	// logOutput is the destination for slog TextHandler.
	logOutput io.Writer = os.Stderr
)

func main() {
	if len(os.Args) == 2 && os.Args[1] == "healthcheck" {
		os.Exit(healthcheck())
	}

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	if err := run(ctx, os.Environ()); err != nil {
		os.Exit(1)
	}
}

func run(ctx context.Context, environ []string) error {
	cfg, err := config.Load(environ)
	if err != nil {
		slog.Error("invalid configuration", "err", err)
		return err
	}

	slog.SetDefault(slog.New(slog.NewTextHandler(logOutput, &slog.HandlerOptions{
		Level: cfg.LogLevel,
	})))
	slog.Info("starting with config", "config", cfg)

	groupCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

	registry := session.NewRegistry()
	bus := session.NewBus()
	
	// Create late-bound sinks to break the circular dependency.
	var bridge *mqttbridge.Bridge
	diagSink := func(d session.Diagnostic) {
		if bridge != nil {
			bridge.DiagnosticSink()(d)
		}
	}
	traceSink := func(serial string, tr ctl.Trace) {
		if bridge != nil {
			bridge.TraceSink()(serial, tr)
		}
	}

	fleet := robot.NewFleet(groupCtx, cfg.ControllerJID, nil, traceSink, diagSink, nil)
	var errBridge error
	bridge, errBridge = mqttbridge.New(groupCtx, cfg, fleet.Submit)
	if errBridge != nil {
		return errBridge
	}

	robot.ParseSchedules = robot.ScheduleParserHook
	robot.ParseLifespan = robot.LifespanParserHook
	bridge.SendCommand = ha.SendCommandHook(cfg)

	fleet.SetRepublishFactory(bridge.Republisher)
	fleet.SetSendCommandFunc(bridge.DispatchSendCommand)
	bus.Register(fleet)
	bus.Register(bridge)
	if shutdownBridge == nil {
		shutdownBridge = bridge.Shutdown
	}
	xmppServer := xmpp.NewServer(cfg, registry, bus, nil, diagSink)
	if shutdownXMPP == nil {
		shutdownXMPP = xmppServer.Shutdown
	}

	// Robot-facing listeners. A failure here stops the process.
	const robotListeners = 3
	errCh := make(chan error, robotListeners)
	go func() { errCh <- httpx.ServeLookup(groupCtx, cfg) }()
	go func() { errCh <- httpx.ServeFirmware(groupCtx, cfg) }()
	go func() { errCh <- xmppServer.Serve(groupCtx) }()

	// Health is optional. Occupied 8080 (OTBR and other HA add-ons) must
	// not tear down lookup/XMPP.
	go func() {
		if err := httpx.ServeHealth(groupCtx, cfg); err != nil && groupCtx.Err() == nil {
			slog.Error("health listener failed; robot listeners continue", "err", err)
		}
	}()

	var runErr error
	received := 0
	select {
	case runErr = <-errCh:
		received++
		cancel()
	case <-ctx.Done():
	}

	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout)
	defer shutdownCancel()
	if shutdownBridge != nil {
		_ = shutdownBridge(shutdownCtx)
	}
	if shutdownXMPP != nil {
		_ = shutdownXMPP(shutdownCtx)
	}

	cancel()
	for received < robotListeners {
		if err := <-errCh; err != nil && runErr == nil {
			runErr = err
		}
		received++
	}

	if shutdownCtx.Err() != nil {
		slog.Info("shutdown budget expired")
		return nil
	}

	if runErr != nil {
		slog.Error("listener failed", "err", runErr)
	}
	return runErr
}

func healthcheck() int {
	cfg, err := config.Load(os.Environ())
	if err != nil {
		slog.Error("invalid configuration", "err", err)
		return 1
	}

	url := fmt.Sprintf("http://127.0.0.1:%d/healthz", cfg.HealthPort)
	client := &http.Client{Timeout: 2 * time.Second}
	resp, err := client.Get(url)
	if err != nil {
		slog.Error("healthcheck request failed", "err", err)
		return 1
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		slog.Error("healthcheck body read failed", "err", err)
		return 1
	}

	if resp.StatusCode != http.StatusOK || string(body) != "ok\n" {
		slog.Error("healthcheck failed", "status", resp.StatusCode, "body", string(body))
		return 1
	}
	return 0
}