package-updater/internal/agent/conn/conn.go

// Package conn keeps an agent connected to the master: it dials over pinned
// TLS, announces itself, sends its inventory, runs the jobs the master sends
// and streams changes and job results until the connection drops, then
// reconnects with backoff.
package conn

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"sync"
	"time"

	"github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"

	"github.com/shcizo/package-updater/internal/agent/inventory"
	"github.com/shcizo/package-updater/internal/agent/updater"
	"github.com/shcizo/package-updater/internal/pin"
	"github.com/shcizo/package-updater/internal/protocol"
)

// readLimit caps a single incoming message. Matches the master.
const readLimit = 1 << 20

// dialTimeout bounds the TLS + WebSocket handshake so a master that accepts
// TCP but never answers does not hang the agent forever.
const dialTimeout = 30 * time.Second

// writeTimeout bounds one websocket write so a stalled peer ends the session
// instead of blocking the writer forever.
const writeTimeout = 10 * time.Second

// defaultProgressInterval is how often job output is batched to the master.
const defaultProgressInterval = 500 * time.Millisecond

// Source supplies what the agent reports. *inventory.Collector satisfies it.
type Source interface {
	Hello(ctx context.Context) (protocol.Hello, error)
	Snapshot(ctx context.Context) ([]protocol.Service, error)
}

// Jobs is the agent's job queue as the connection sees it. *updater.Runner
// satisfies it.
type Jobs interface {
	// Submit queues a job; a job ID already held is ignored.
	Submit(id string, start protocol.JobStart)
	// Ack forgets a finished job the master has stored.
	Ack(id string)
	// Known lists every job ID held: queued, running, or finished unacked.
	Known() []string
	// TakeProgress returns output produced since the previous call.
	TakeProgress() []updater.Progress
	// Pending returns results the master has not acknowledged yet.
	Pending() []updater.Finished
	// Changed receives when there is new progress or a new result.
	Changed() <-chan struct{}
}

// Config holds connection settings. Token is sent as a bearer header and must
// never be logged.
type Config struct {
	URL              string // wss://host:port/agent
	Token            string // AGENT_TOKEN
	Fingerprint      string // MASTER_KEY
	Version          string
	MinBackoff       time.Duration // default 1s
	MaxBackoff       time.Duration // default 30s
	ProgressInterval time.Duration // default 500ms
}

// Agent is the agent's side of the master connection.
type Agent struct {
	cfg     Config
	src     Source
	jobs    Jobs
	log     *slog.Logger
	changes chan struct{}
}

// New returns an Agent. Call Run to connect.
func New(cfg Config, src Source, jobs Jobs, log *slog.Logger) *Agent {
	if cfg.MinBackoff == 0 {
		cfg.MinBackoff = time.Second
	}
	if cfg.MaxBackoff == 0 {
		cfg.MaxBackoff = 30 * time.Second
	}
	if cfg.ProgressInterval == 0 {
		cfg.ProgressInterval = defaultProgressInterval
	}
	return &Agent{cfg: cfg, src: src, jobs: jobs, log: log, changes: make(chan struct{}, 1)}
}

// Notify tells the agent its inventory may have changed. It never blocks;
// several calls before the agent gets to them collapse into one refresh.
func (a *Agent) Notify() {
	select {
	case a.changes <- struct{}{}:
	default:
	}
}

// Run keeps the agent connected until ctx ends, then returns nil.
func (a *Agent) Run(ctx context.Context) error {
	backoff := a.cfg.MinBackoff
	for {
		connected, err := a.session(ctx)
		if ctx.Err() != nil {
			return nil
		}
		if connected {
			backoff = a.cfg.MinBackoff
		}
		a.log.Warn("master_connection_lost", "error", err, "retry_in", backoff.String())
		select {
		case <-ctx.Done():
			return nil
		case <-time.After(backoff):
		}
		backoff = min(backoff*2, a.cfg.MaxBackoff)
	}
}

// session runs one connection until it fails. connected reports whether the
// hello was delivered, so Run resets its backoff only after real progress.
//
// Two goroutines serve the connection: a reader, which handles what the
// master sends, and a writer, the only goroutine that writes to the socket.
func (a *Agent) session(parent context.Context) (connected bool, err error) {
	httpClient := &http.Client{Transport: &http.Transport{TLSClientConfig: pin.ClientConfig(a.cfg.Fingerprint)}}
	// The dial context only bounds the handshake; coder/websocket does not
	// tie the established connection to it, so cancelling afterwards is safe.
	dialCtx, cancelDial := context.WithTimeout(parent, dialTimeout)
	c, _, err := websocket.Dial(dialCtx, a.cfg.URL, &websocket.DialOptions{
		HTTPClient: httpClient,
		HTTPHeader: http.Header{"Authorization": []string{"Bearer " + a.cfg.Token}},
	})
	cancelDial()
	if err != nil {
		return false, fmt.Errorf("dial master: %w", err)
	}
	c.SetReadLimit(readLimit)

	ctx, cancel := context.WithCancelCause(parent)
	out := make(chan protocol.Envelope, 64)
	var wg sync.WaitGroup
	defer func() {
		cancel(nil)
		if parent.Err() != nil {
			_ = c.Close(websocket.StatusNormalClosure, "")
		}
		c.CloseNow()
		wg.Wait()
	}()

	wg.Add(2)
	go func() { defer wg.Done(); a.writeLoop(ctx, cancel, c, out) }()
	go func() { defer wg.Done(); a.readLoop(ctx, cancel, c) }()

	// send queues an envelope for the writer. It gives up when the session ends.
	send := func(typ string, payload any) error {
		env, err := protocol.Encode(typ, payload)
		if err != nil {
			return err
		}
		return enqueue(ctx, out, env)
	}
	sendJob := func(typ, id string, payload any) error {
		env, err := protocol.EncodeJob(typ, id, payload)
		if err != nil {
			return err
		}
		return enqueue(ctx, out, env)
	}
	fail := func(connected bool, err error) (bool, error) {
		if cause := context.Cause(ctx); cause != nil && !errors.Is(cause, context.Canceled) {
			return connected, cause
		}
		return connected, err
	}

	hello, err := a.src.Hello(ctx)
	if err != nil {
		return fail(false, err)
	}
	hello.AgentVersion = a.cfg.Version
	hello.ProtocolVersion = protocol.Version
	hello.KnownJobs = a.jobs.Known()
	if err := send(protocol.TypeHello, hello); err != nil {
		return fail(false, err)
	}

	last, err := a.src.Snapshot(ctx)
	if err != nil {
		return fail(true, err)
	}
	if err := send(protocol.TypeSnapshot, protocol.Snapshot{Services: last}); err != nil {
		return fail(true, err)
	}
	a.log.Info("master_connected", "services", len(last), "known_jobs", len(hello.KnownJobs))

	sent := make(map[string]bool) // job results already sent on this connection
	flush := func() error {
		for _, p := range a.jobs.TakeProgress() {
			if err := sendJob(protocol.TypeJobProgress, p.ID, protocol.JobProgress{Output: p.Output}); err != nil {
				return err
			}
		}
		for _, f := range a.jobs.Pending() {
			if sent[f.ID] {
				continue
			}
			if err := sendJob(protocol.TypeJobResult, f.ID, f.Result); err != nil {
				return err
			}
			sent[f.ID] = true
		}
		return nil
	}

	// flush is throttled to one call per ProgressInterval; a signal that
	// arrives early arms the timer for the remainder instead of being lost.
	lastFlush := time.Time{}
	var timer *time.Timer
	var timerC <-chan time.Time
	defer func() {
		if timer != nil {
			timer.Stop()
		}
	}()
	maybeFlush := func() error {
		wait := a.cfg.ProgressInterval - time.Since(lastFlush)
		if wait > 0 {
			if timerC == nil {
				timer = time.NewTimer(wait)
				timerC = timer.C
			}
			return nil
		}
		lastFlush = time.Now()
		return flush()
	}
	if err := maybeFlush(); err != nil {
		return fail(true, err)
	}

	jobsChanged := a.jobs.Changed()
	for {
		select {
		case <-ctx.Done():
			return fail(true, ctx.Err())
		case <-jobsChanged:
			if err := maybeFlush(); err != nil {
				return fail(true, err)
			}
		case <-timerC:
			timerC = nil
			lastFlush = time.Now()
			if err := flush(); err != nil {
				return fail(true, err)
			}
		case <-a.changes:
			next, err := a.src.Snapshot(ctx)
			if err != nil {
				return fail(true, err)
			}
			upserts, removes := inventory.Diff(last, next)
			for _, s := range upserts {
				if err := send(protocol.TypeUpsert, protocol.Upsert{Service: s}); err != nil {
					return fail(true, err)
				}
			}
			for _, k := range removes {
				if err := send(protocol.TypeRemove, protocol.Remove{Key: k}); err != nil {
					return fail(true, err)
				}
			}
			last = next
		}
	}
}

// enqueue hands env to the writer, or gives up when the session ends.
func enqueue(ctx context.Context, out chan<- protocol.Envelope, env protocol.Envelope) error {
	select {
	case out <- env:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// writeLoop is the only goroutine that writes to c. A failed or timed-out
// write ends the session.
func (a *Agent) writeLoop(ctx context.Context, cancel context.CancelCauseFunc, c *websocket.Conn, out <-chan protocol.Envelope) {
	for {
		select {
		case <-ctx.Done():
			return
		case env := <-out:
			wctx, wcancel := context.WithTimeout(ctx, writeTimeout)
			err := wsjson.Write(wctx, c, env)
			wcancel()
			if err != nil {
				cancel(fmt.Errorf("send %s: %w", env.Type, err))
				return
			}
		}
	}
}

// readLoop handles what the master sends. A message it cannot parse ends the
// session; a message type it does not know is logged and ignored, so a newer
// master can add message types without breaking this agent.
func (a *Agent) readLoop(ctx context.Context, cancel context.CancelCauseFunc, c *websocket.Conn) {
	for {
		var env protocol.Envelope
		if err := wsjson.Read(ctx, c, &env); err != nil {
			cancel(fmt.Errorf("read: %w", err))
			return
		}
		switch env.Type {
		case protocol.TypeJobStart:
			start, err := protocol.Decode[protocol.JobStart](env)
			if err != nil || env.ID == "" {
				cancel(fmt.Errorf("malformed job.start (id %q): %v", env.ID, err))
				return
			}
			a.jobs.Submit(env.ID, start)
		case protocol.TypeJobAck:
			// JobAck has no fields, so the payload is not inspected.
			a.jobs.Ack(env.ID)
		default:
			a.log.Warn("unknown_message_type", "type", env.Type)
		}
	}
}