// Package mqttclient wraps paho for device-control-service's two-way MQTT
// use: subscribing to telemetry/ack topics for liveness and reported state,
// and publishing commands to devices.
package mqttclient

import (
	"fmt"
	"log/slog"
	"sync"
	"time"

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

// Handler receives one message's topic, raw payload, and arrival time.
type Handler func(topic string, payload []byte, receivedAt time.Time)

type Config struct {
	BrokerURL string // e.g. tcp://mosquitto:1883
	ClientID  string
}

type subscription struct {
	topic   string
	qos     byte
	handler Handler
}

// Client wraps a paho client, replaying every registered subscription on
// each (re)connect so a dropped connection resubscribes automatically.
type Client struct {
	client mqtt.Client
	logger *slog.Logger

	mu   sync.Mutex
	subs []subscription
}

func Connect(cfg Config, logger *slog.Logger) (*Client, error) {
	c := &Client{logger: logger}

	opts := mqtt.NewClientOptions().
		AddBroker(cfg.BrokerURL).
		SetClientID(cfg.ClientID).
		SetAutoReconnect(true).
		SetConnectRetry(true).
		SetOnConnectHandler(func(mc mqtt.Client) {
			c.resubscribeAll(mc)
		}).
		SetConnectionLostHandler(func(_ mqtt.Client, err error) {
			logger.Warn("mqtt connection lost", "error", err)
		})

	mqttClient := mqtt.NewClient(opts)
	token := mqttClient.Connect()
	token.Wait()
	if err := token.Error(); err != nil {
		return nil, fmt.Errorf("connect to mqtt broker %s: %w", cfg.BrokerURL, err)
	}
	c.client = mqttClient
	return c, nil
}

// Subscribe registers a handler for topic. It subscribes immediately if
// already connected, and will be replayed on every future reconnect.
func (c *Client) Subscribe(topic string, qos byte, handler Handler) error {
	c.mu.Lock()
	c.subs = append(c.subs, subscription{topic: topic, qos: qos, handler: handler})
	c.mu.Unlock()

	if !c.client.IsConnectionOpen() {
		return nil
	}
	return c.subscribeOne(c.client, subscription{topic: topic, qos: qos, handler: handler})
}

func (c *Client) resubscribeAll(mc mqtt.Client) {
	c.mu.Lock()
	subs := append([]subscription(nil), c.subs...)
	c.mu.Unlock()

	for _, s := range subs {
		if err := c.subscribeOne(mc, s); err != nil {
			c.logger.Error("mqtt subscribe failed", "topic", s.topic, "error", err)
		}
	}
}

func (c *Client) subscribeOne(mc mqtt.Client, s subscription) error {
	token := mc.Subscribe(s.topic, s.qos, func(_ mqtt.Client, msg mqtt.Message) {
		s.handler(msg.Topic(), msg.Payload(), time.Now())
	})
	token.Wait()
	return token.Error()
}

// Publish sends payload to topic and waits for the publish to complete.
func (c *Client) Publish(topic string, qos byte, retained bool, payload []byte) error {
	token := c.client.Publish(topic, qos, retained, payload)
	token.Wait()
	return token.Error()
}

// Close disconnects, waiting up to 250ms for in-flight work to settle.
func (c *Client) Close() {
	c.client.Disconnect(250)
}