// Enrollment backend for the interop spike.
//
// Provisions Mosquitto Dynamic Security, then answers enroll/request/<clientid> with an X25519
// key agreement (signed with Ed25519), creating the FN broker client with the derived password.
// Every step is logged. Passwords are never logged, only a short SHA-256 fingerprint.
package main

import (
	"bytes"
	"crypto/ecdh"
	"crypto/ed25519"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
)

// Dynamic Security is controlled by publishing JSON commands to $CONTROL/dynamic-security/v1 and
// reading the answer on .../response. hkdfInfoLabel is mixed into the password derivation so a
// secret derived for this purpose cannot be reused for anything else (the FN or SN is appended).
const (
	controlTopic  = "$CONTROL/dynamic-security/v1"
	controlResp   = "$CONTROL/dynamic-security/v1/response"
	hkdfInfoLabel = "robot-market/mqtt-pw/v1|"
)

// stepDelay is the pause between steps (STEP_DELAY seconds) so a human can follow the logs.
var stepDelay time.Duration

var logFile *os.File // optional copy of every line for the web viewer (LOG_FILE)

// logf prints one timestamped line to stdout and, when LOG_FILE is set, to that file too.
// The web viewer reads the file; "docker compose logs" shows stdout.
func logf(format string, a ...any) {
	line := fmt.Sprintf("%s [backend] %s\n", time.Now().Format("15:04:05.000"), fmt.Sprintf(format, a...))
	fmt.Print(line)
	if logFile != nil {
		logFile.WriteString(line)
	}
}

// step prints a header line; the web page turns lines in this exact format into step bands.
func step(n int, title string) {
	logf("---- step %d: %s ----", n, title)
}

func pause() { time.Sleep(stepDelay) }

func env(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}

func envInt(key string, def int) int {
	if v, err := strconv.Atoi(os.Getenv(key)); err == nil {
		return v
	}
	return def
}

func fatal(format string, a ...any) {
	logf("FATAL: "+format, a...)
	os.Exit(1)
}

// fingerprint is a short SHA-256 of the password. Logging this instead of the password lets the
// backend and the boards prove they derived the same value without ever exposing it.
func fingerprint(s string) string {
	h := sha256.Sum256([]byte(s))
	return hex.EncodeToString(h[:8])
}

// control wraps the broker admin connection. do() sends one batch of commands and blocks until the
// matching response arrives (one request at a time, guarded by mu).
// control sends Dynamic Security commands and waits for the matching response.
type control struct {
	client mqtt.Client
	mu     sync.Mutex
	resp   chan []byte
}

func (c *control) do(cmds ...map[string]any) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	body, _ := json.Marshal(map[string]any{"commands": cmds})
	select { // drop any stale response
	case <-c.resp:
	default:
	}
	if t := c.client.Publish(controlTopic, 1, false, body); t.Wait() && t.Error() != nil {
		return t.Error()
	}
	select {
	case raw := <-c.resp:
		var r struct {
			Responses []struct {
				Command string `json:"command"`
				Error   string `json:"error"`
			} `json:"responses"`
		}
		if err := json.Unmarshal(raw, &r); err != nil {
			return fmt.Errorf("bad control response: %w", err)
		}
		var failed []string
		for _, x := range r.Responses {
			if x.Error != "" {
				failed = append(failed, x.Command+": "+x.Error)
			}
		}
		if len(failed) > 0 {
			return fmt.Errorf("%s", strings.Join(failed, "; "))
		}
		return nil
	case <-time.After(5 * time.Second):
		return fmt.Errorf("timeout waiting for dynamic security response")
	}
}

// acl builds an addRoleACL command: allow `typ` (publishClientSend, subscribePattern, ...) on `topic`.
// %c in a topic means the client's Client ID, %u its username (Mosquitto 2.1 or newer).
func acl(role, typ, topic string) map[string]any {
	return map[string]any{"command": "addRoleACL", "rolename": role, "acltype": typ, "topic": topic, "priority": 0, "allow": true}
}

// provision configures the broker once at start-up. After it runs:
//   - the service client (backend-svc) may read enrollment requests, answer them, and talk to FN boards;
//   - anonymous clients may only publish enroll/request/<their own client id> and listen on
//     enroll/response/<their own client id>;
//   - FN boards may only publish board/<FN>/hello and listen on board/<FN>/cmd;
//   - everything not explicitly allowed is denied (default ACL access = deny).
func provision(c *control, svcUser, svcPassword string) {
	step(2, "provision Dynamic Security (roles, anonymous group, deny-by-default)")
	logf("creating role backend-service and a separate service client %q (the admin client only manages the broker; changing a connected client's roles disconnects it, so it must not modify itself)", svcUser)
	err := c.do(
		map[string]any{"command": "createRole", "rolename": "backend-service"},
		acl("backend-service", "subscribePattern", "enroll/request/+"),
		acl("backend-service", "publishClientReceive", "enroll/request/+"),
		acl("backend-service", "publishClientSend", "enroll/response/+"),
		acl("backend-service", "subscribePattern", "board/+/hello"),
		acl("backend-service", "publishClientReceive", "board/+/hello"),
		acl("backend-service", "publishClientSend", "board/+/cmd"),
		map[string]any{"command": "createClient", "username": svcUser, "password": svcPassword, "textname": "backend service",
			"roles": []map[string]any{{"rolename": "backend-service"}}},
	)
	if err != nil {
		fatal("backend-service role: %v", err)
	}
	pause()

	logf("creating role enroll (anonymous): publish enroll/request/%%c, subscribe+receive enroll/response/%%c only")
	err = c.do(
		map[string]any{"command": "createRole", "rolename": "enroll"},
		acl("enroll", "publishClientSend", "enroll/request/%c"),
		acl("enroll", "subscribePattern", "enroll/response/%c"),
		acl("enroll", "publishClientReceive", "enroll/response/%c"),
		acl("enroll", "unsubscribePattern", "enroll/response/%c"),
		map[string]any{"command": "createGroup", "groupname": "enrollment"},
		map[string]any{"command": "addGroupRole", "groupname": "enrollment", "rolename": "enroll"},
		map[string]any{"command": "setAnonymousGroup", "groupname": "enrollment"},
	)
	if err != nil {
		fatal("enroll role / anonymous group: %v", err)
	}
	pause()

	logf("creating role board-bootstrap (FN boards): publish board/%%u/hello, subscribe+receive board/%%u/cmd only")
	err = c.do(
		map[string]any{"command": "createRole", "rolename": "board-bootstrap"},
		acl("board-bootstrap", "publishClientSend", "board/%u/hello"),
		acl("board-bootstrap", "subscribePattern", "board/%u/cmd"),
		acl("board-bootstrap", "publishClientReceive", "board/%u/cmd"),
		acl("board-bootstrap", "unsubscribePattern", "board/%u/cmd"),
	)
	if err != nil {
		fatal("board-bootstrap role: %v", err)
	}
	pause()

	logf("setting default ACL access to deny for publishClientSend/Receive, subscribe, unsubscribe")
	err = c.do(map[string]any{"command": "setDefaultACLAccess", "acls": []map[string]any{
		{"acltype": "publishClientSend", "allow": false},
		{"acltype": "publishClientReceive", "allow": false},
		{"acltype": "subscribe", "allow": false},
		{"acltype": "unsubscribe", "allow": false},
	}})
	if err != nil {
		fatal("default ACL: %v", err)
	}
	logf("provisioning done")
}

// server holds the enrollment state. replies caches the answer per Client ID: an empty string means
// "being processed", a non-empty one is the finished answer that can be re-sent if the board retries.
type server struct {
	ctl     *control
	pub     mqtt.Client
	signKey ed25519.PrivateKey

	mu      sync.Mutex
	nextFN  int
	replies map[string]string // clientid -> cached response (public values only)
}

// parseFields reads the simple "key=value" lines used by the request and response messages.
func parseFields(payload []byte) map[string]string {
	out := map[string]string{}
	for _, line := range strings.Split(string(payload), "\n") {
		if k, v, ok := strings.Cut(strings.TrimSpace(line), "="); ok {
			out[k] = v
		}
	}
	return out
}

// validClientID accepts only 32 hex characters (16 random bytes), the format boards generate.
func validClientID(cid string) bool {
	if len(cid) != 32 {
		return false
	}
	_, err := hex.DecodeString(cid)
	return err == nil
}

// handleRequest is the whole enrollment: validate the request, reserve the Client ID, assign an FN,
// run the X25519 key agreement, derive the password, create the broker account, sign the answer and
// publish it. The password is created on both sides independently; it never travels.
func (s *server) handleRequest(cid string, payload []byte) {
	if !validClientID(cid) {
		logf("REJECT request: client id %q is not 32 hex chars", cid)
		return
	}
	f := parseFields(payload)
	boardPub, err1 := hex.DecodeString(f["B"])
	nonce, err2 := hex.DecodeString(f["n1"])
	if err1 != nil || err2 != nil || len(boardPub) != 32 || len(nonce) != 16 {
		logf("REJECT request from %s: bad B/n1 fields", cid)
		return
	}
	mac := f["mac"]

	s.mu.Lock()
	if cached, ok := s.replies[cid]; ok {
		s.mu.Unlock()
		if cached == "" {
			logf("duplicate request from %s while the first one is still being processed: ignored", cid)
			return
		}
		logf("duplicate request from %s: re-sending the cached response (no new FN, no new password)", cid)
		s.pub.Publish("enroll/response/"+cid, 1, false, cached).Wait()
		return
	}
	s.replies[cid] = "" // reserve: marks this Client ID as being processed
	s.nextFN++
	fn := fmt.Sprintf("FN-%06d", s.nextFN)
	s.mu.Unlock()
	done := false
	defer func() {
		if !done { // failed: release the reservation so the board can retry
			s.mu.Lock()
			delete(s.replies, cid)
			s.mu.Unlock()
		}
	}()

	step(4, "enrollment request from "+cid)
	logf("board public key B=%s nonce n1=%s mac=%s (mac is only a label, not proof)", f["B"], f["n1"], mac)
	pause()

	logf("assigned %s; generating ephemeral X25519 key pair (s, S)", fn)
	// The backend's own throw-away key pair (s, S). Combined with the board's public key B it gives
	// the shared secret; the private half s is discarded when this function returns.
	eph, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		logf("ERROR ephemeral key: %v", err)
		return
	}
	bPub, err := ecdh.X25519().NewPublicKey(boardPub)
	if err != nil {
		logf("REJECT %s: invalid board public key: %v", cid, err)
		return
	}
	shared, err := eph.ECDH(bPub)
	if err != nil {
		logf("REJECT %s: ECDH failed: %v", cid, err)
		return
	}
	logf("S=%x", eph.PublicKey().Bytes())
	pause()

	// Turn the shared secret into the broker password. salt = the board's nonce, info = label + FN,
	// so the result is unique to this board, this request and this identity.
	pwBytes, err := hkdf.Key(sha256.New, shared, nonce, hkdfInfoLabel+fn, 32)
	if err != nil {
		logf("ERROR hkdf: %v", err)
		return
	}
	pw := hex.EncodeToString(pwBytes)
	logf("derived password with HKDF-SHA256(shared, salt=n1, info=%q); fingerprint=%s (password itself is never logged or sent)", hkdfInfoLabel+fn, fingerprint(pw))
	pause()

	logf("creating broker client username=%s clientid=%s role=board-bootstrap via Dynamic Security", fn, fn)
	if err := s.ctl.do(map[string]any{
		"command": "createClient", "username": fn, "password": pw, "clientid": fn,
		"textname": "board " + mac, "roles": []map[string]any{{"rolename": "board-bootstrap"}},
	}); err != nil {
		logf("ERROR createClient %s: %v", fn, err)
		return
	}
	logf("client %s created", fn)
	pause()

	// Sign everything the board needs to trust: both public keys, its nonce, its Client ID and the FN.
	// A reply for another request, another board or another FN will not verify on the board.
	msg := bytes.Join([][]byte{eph.PublicKey().Bytes(), boardPub, nonce, []byte(cid), []byte(fn)}, nil)
	sig := ed25519.Sign(s.signKey, msg)
	logf("signed S||B||n1||clientid||FN with the backend Ed25519 key; sig=%x", sig)
	reply := fmt.Sprintf("fn=%s\nS=%x\nsig=%x\n", fn, eph.PublicKey().Bytes(), sig)

	s.mu.Lock()
	s.replies[cid] = reply
	s.mu.Unlock()
	done = true
	if t := s.pub.Publish("enroll/response/"+cid, 1, false, reply); t.Wait() && t.Error() != nil {
		logf("ERROR publishing response: %v", t.Error())
		return
	}
	logf("response published on enroll/response/%s (fn, S, signature only; not retained)", cid)
}

// handleHello is called when an FN board publishes board/<FN>/hello. Receiving it proves the board
// logged in with the derived password and that its ACL allows publishing there; the backend answers
// on board/<FN>/cmd, which also proves the board may receive there.
func (s *server) handleHello(topic string, payload []byte) {
	parts := strings.Split(topic, "/")
	if len(parts) != 3 {
		return
	}
	fn := parts[1]
	step(9, "hello from "+fn)
	logf("received %q on %s (this proves the board logged in with the derived password and its ACL allows board/%%u/hello)", string(payload), topic)
	if t := s.pub.Publish("board/"+fn+"/cmd", 1, false, "ack-hello"); t.Wait() && t.Error() != nil {
		logf("ERROR publishing ack: %v", t.Error())
		return
	}
	logf("sent ack-hello on board/%s/cmd", fn)
}

// main: wait (so logs can be followed) -> connect as admin -> provision the broker -> connect the
// service client -> subscribe -> serve requests forever.
func main() {
	if path := os.Getenv("LOG_FILE"); path != "" {
		logFile, _ = os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC|os.O_APPEND, 0o644)
	}
	stepDelay = time.Duration(envInt("STEP_DELAY", 2)) * time.Second

	seed, err := hex.DecodeString(os.Getenv("SIGNING_SEED_HEX"))
	if err != nil || len(seed) != ed25519.SeedSize {
		fatal("SIGNING_SEED_HEX must be 32 bytes of hex")
	}
	signKey := ed25519.NewKeyFromSeed(seed)

	start := envInt("START_DELAY", 15)
	logf("backend starting; waiting %d s so the logs can be followed (set START_DELAY to change)", start)
	for left := start; left > 0; left -= 5 {
		logf("starting in %d s ...", left)
		time.Sleep(time.Duration(min(5, left)) * time.Second)
	}

	step(1, "connect to the broker as the Dynamic Security admin")
	ctl := &control{resp: make(chan []byte, 4)}
	opts := mqtt.NewClientOptions().
		AddBroker(env("BROKER_URL", "tcp://broker:1883")).
		SetClientID("backend-service").
		SetUsername(env("ADMIN_USER", "backend")).
		SetPassword(os.Getenv("ADMIN_PASSWORD")).
		SetConnectRetry(true).
		SetConnectRetryInterval(2 * time.Second).
		SetAutoReconnect(true).
		SetOrderMatters(false) // handlers may block waiting for control responses
	client := mqtt.NewClient(opts)
	ctl.client = client
	client.Connect().Wait()
	logf("connected as user %q", env("ADMIN_USER", "backend"))

	if t := client.Subscribe(controlResp, 1, func(_ mqtt.Client, m mqtt.Message) {
		select {
		case ctl.resp <- m.Payload():
		default:
		}
	}); t.Wait() && t.Error() != nil {
		fatal("subscribe control response: %v", t.Error())
	}
	pause()

	svcUser, svcPassword := "backend-svc", env("SERVICE_PASSWORD", "demo-service-secret")
	provision(ctl, svcUser, svcPassword)
	pause()

	step(3, "connect the service client and subscribe to enrollment requests and board hellos")
	svc := mqtt.NewClient(mqtt.NewClientOptions().
		AddBroker(env("BROKER_URL", "tcp://broker:1883")).
		SetClientID("backend-svc").
		SetUsername(svcUser).
		SetPassword(svcPassword).
		SetConnectRetry(true).
		SetConnectRetryInterval(2 * time.Second).
		SetAutoReconnect(true).
		SetOrderMatters(false))
	svc.Connect().Wait()
	logf("service client connected as %q", svcUser)
	s := &server{ctl: ctl, pub: svc, signKey: signKey, replies: map[string]string{}}
	svc.Subscribe("enroll/request/+", 1, func(_ mqtt.Client, m mqtt.Message) {
		cid := strings.TrimPrefix(m.Topic(), "enroll/request/")
		go s.handleRequest(cid, m.Payload())
	}).Wait()
	svc.Subscribe("board/+/hello", 1, func(_ mqtt.Client, m mqtt.Message) {
		go s.handleHello(m.Topic(), m.Payload())
	}).Wait()
	logf("ready: waiting for boards (signing public key=%x)", signKey.Public().(ed25519.PublicKey))

	select {}
}