// Package batcher buffers measurements and traceroutes, and delivers them as gzip-compressed signed batches.
//
// Delivery rules (docs/PROBE-PROTOCOL.md § POST /batch): flush every batch_flush_seconds or when max_batch is
// reached; health block at most once per minute; network error / 5xx → spool to disk and back off 5 s → 5 min,
// drain oldest first when the server is back (re-signed with a fresh timestamp, bytes unchanged); 4xx → drop
// (except "401 skew": resync clock and retry once).
package batcher

import (
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"sync"
	"sync/atomic"
	"time"

	"internetpressure.io/probe-agent/internal/client"
	"internetpressure.io/probe-agent/internal/protocol"
	"internetpressure.io/probe-agent/internal/spool"
)

// Defaults and limits.
const (
	DefaultFlushEvery = 10 * time.Second
	DefaultMaxBatch   = 500
	HealthEvery       = time.Minute
	BackoffMin        = 5 * time.Second
	BackoffMax        = 5 * time.Minute
	FinalFlushTimeout = 5 * time.Second
	drainPerCycle     = 20
)

// Poster is the subset of client.Client used here (interface for tests).
type Poster interface {
	PostBatch(ctx context.Context, gz []byte, timeout time.Duration) (*protocol.BatchResponse, error)
}

// Hooks are optional callbacks.
type Hooks struct {
	Health         func() *protocol.Health // builds the health block (called ≤ once/minute)
	OnResponse     func(*protocol.BatchResponse)
	OnFlushFailure func()
	OnDelivered    func(t time.Time)
}

// Batcher is the queue + flusher.
type Batcher struct {
	probeID string
	version string
	poster  Poster
	spool   *spool.Spool
	hooks   Hooks
	log     *slog.Logger

	flushEvery atomic.Int64 // ns
	maxBatch   atomic.Int32

	mu    sync.Mutex
	meas  []protocol.Measurement
	trs   []protocol.Traceroute
	kick  chan struct{}
	timer *time.Timer

	lastHealth time.Time
	backoff    time.Duration
	nextTry    time.Time
	stopped    atomic.Bool
}

// New creates a Batcher.
func New(probeID, version string, poster Poster, sp *spool.Spool, hooks Hooks, log *slog.Logger) *Batcher {
	if log == nil {
		log = slog.Default()
	}
	b := &Batcher{probeID: probeID, version: version, poster: poster, spool: sp, hooks: hooks, log: log,
		kick: make(chan struct{}, 1)}
	b.flushEvery.Store(int64(DefaultFlushEvery))
	b.maxBatch.Store(DefaultMaxBatch)
	return b
}

// Configure updates flush interval / batch size (from the remote schedule). Zero keeps the default.
func (b *Batcher) Configure(flushSeconds, maxBatch int) {
	if flushSeconds > 0 {
		b.flushEvery.Store(int64(time.Duration(flushSeconds) * time.Second))
	}
	if maxBatch > 0 {
		b.maxBatch.Store(int32(maxBatch))
	}
}

// Add queues measurements; triggers an immediate flush when max_batch is reached.
func (b *Batcher) Add(ms ...protocol.Measurement) {
	if len(ms) == 0 {
		return
	}
	b.mu.Lock()
	b.meas = append(b.meas, ms...)
	full := len(b.meas) >= int(b.maxBatch.Load())
	b.mu.Unlock()
	if full {
		b.Kick()
	}
}

// AddTraceroute queues a traceroute.
func (b *Batcher) AddTraceroute(t protocol.Traceroute) {
	b.mu.Lock()
	b.trs = append(b.trs, t)
	b.mu.Unlock()
}

// Kick requests a flush as soon as possible.
func (b *Batcher) Kick() {
	select {
	case b.kick <- struct{}{}:
	default:
	}
}

// Buffered returns the number of queued measurements + traceroutes.
func (b *Batcher) Buffered() int {
	b.mu.Lock()
	defer b.mu.Unlock()
	return len(b.meas) + len(b.trs)
}

// Run flushes on the timer and on kicks until ctx is cancelled. The caller then invokes Stop once the
// producers (scheduler) have finished, so the final flush contains every completed measurement.
func (b *Batcher) Run(ctx context.Context) {
	for {
		wait := time.Duration(b.flushEvery.Load())
		select {
		case <-ctx.Done():
			return
		case <-b.kick:
		case <-time.After(wait):
		}
		b.Flush(ctx)
	}
}

// Stop performs the graceful final flush: one attempt with a 5 s timeout, everything else spooled.
func (b *Batcher) Stop() {
	if !b.stopped.CompareAndSwap(false, true) {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), FinalFlushTimeout)
	defer cancel()
	for b.Buffered() > 0 {
		body, n := b.take()
		if n == 0 {
			break
		}
		gz, err := encode(body)
		if err != nil {
			b.log.Error("final flush: encode", "err", err)
			return
		}
		if _, err := b.poster.PostBatch(ctx, gz, FinalFlushTimeout); err != nil {
			b.toSpool(gz, "final flush failed: "+err.Error())
			// After one failure, spool the rest without trying the network again.
			for b.Buffered() > 0 {
				body, n := b.take()
				if n == 0 {
					break
				}
				if gz, err := encode(body); err == nil {
					b.toSpool(gz, "final flush: spooled remainder")
				}
			}
			return
		}
		b.log.Info("final flush delivered", "measurements", len(body.Measurements), "traceroutes", len(body.Traceroutes))
	}
}

// Flush runs one delivery cycle: drain the spool (oldest first), then send the live queue.
func (b *Batcher) Flush(ctx context.Context) {
	now := time.Now()
	inBackoff := b.backoff > 0 && now.Before(b.nextTry)

	if !inBackoff && b.spool != nil {
		b.drainSpool(ctx)
		inBackoff = b.backoff > 0 && time.Now().Before(b.nextTry)
	}

	for {
		body, n := b.take()
		if n == 0 {
			return
		}
		gz, err := encode(body)
		if err != nil {
			b.log.Error("batch encode failed, dropping", "err", err, "n", n)
			return
		}
		if inBackoff {
			b.toSpool(gz, "server unavailable (backoff)")
		} else if !b.send(ctx, gz, body) {
			inBackoff = true
		}
		// Keep going until the queue is empty (in max_batch-sized chunks).
		b.mu.Lock()
		more := len(b.meas) > 0 || len(b.trs) > 0
		b.mu.Unlock()
		if !more {
			return
		}
	}
}

// take removes up to max_batch measurements (and traceroutes) from the queue and builds the batch body,
// attaching the health block when due. Returns (body, number of items taken).
func (b *Batcher) take() (*protocol.Batch, int) {
	max := int(b.maxBatch.Load())
	b.mu.Lock()
	nm := len(b.meas)
	if nm > max {
		nm = max
	}
	nt := len(b.trs)
	if nt > max {
		nt = max
	}
	body := &protocol.Batch{ProbeID: b.probeID, AgentVersion: b.version,
		Measurements: make([]protocol.Measurement, nm)}
	copy(body.Measurements, b.meas[:nm])
	b.meas = append(b.meas[:0:0], b.meas[nm:]...) // fresh backing array so memory is released
	if nt > 0 {
		body.Traceroutes = make([]protocol.Traceroute, nt)
		copy(body.Traceroutes, b.trs[:nt])
		b.trs = append(b.trs[:0:0], b.trs[nt:]...)
	}
	b.mu.Unlock()

	healthDue := b.hooks.Health != nil && time.Since(b.lastHealth) >= HealthEvery
	if nm+nt == 0 && !healthDue {
		return nil, 0
	}
	if healthDue {
		body.Health = b.hooks.Health()
		b.lastHealth = time.Now()
	}
	body.SentAt = protocol.FormatTime(time.Now())
	n := nm + nt
	if n == 0 {
		n = 1 // health-only batch
	}
	return body, n
}

// encode marshals + gzips a batch.
func encode(body *protocol.Batch) ([]byte, error) {
	raw, err := json.Marshal(body)
	if err != nil {
		return nil, err
	}
	var buf bytes.Buffer
	zw, _ := gzip.NewWriterLevel(&buf, gzip.BestSpeed)
	if _, err := zw.Write(raw); err != nil {
		return nil, err
	}
	if err := zw.Close(); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// send posts one live batch. Returns false when the server is unavailable (batch spooled, backoff set).
func (b *Batcher) send(ctx context.Context, gz []byte, body *protocol.Batch) bool {
	resp, err := b.poster.PostBatch(ctx, gz, client.BatchTimeout)
	if err != nil {
		var he *client.HTTPError
		if errors.As(err, &he) && he.IsSkew() {
			// The client resynced its clock from the response; retry once with a fresh signature.
			b.log.Warn("batch rejected for clock skew, resynced and retrying once", "err", err)
			resp, err = b.poster.PostBatch(ctx, gz, client.BatchTimeout)
		}
	}
	if err != nil {
		if client.Retryable(err) {
			b.fail(err)
			b.toSpool(gz, err.Error())
			return false
		}
		b.log.Error("batch rejected, dropping", "err", err, "measurements", len(body.Measurements))
		return true
	}
	b.success(resp, len(body.Measurements), len(body.Traceroutes))
	return true
}

// drainSpool re-sends spooled batches oldest first until empty, a failure, or drainPerCycle files.
func (b *Batcher) drainSpool(ctx context.Context) {
	for i := 0; i < drainPerCycle; i++ {
		name, gz, err := b.spool.Oldest()
		if err != nil {
			return
		}
		resp, err := b.poster.PostBatch(ctx, gz, client.BatchTimeout)
		if err != nil {
			var he *client.HTTPError
			if errors.As(err, &he) && he.IsSkew() {
				resp, err = b.poster.PostBatch(ctx, gz, client.BatchTimeout)
			}
		}
		if err != nil {
			if client.Retryable(err) {
				b.fail(err)
				return
			}
			b.log.Error("spooled batch rejected, dropping", "file", name, "err", err)
			b.spool.Remove(name)
			continue
		}
		b.spool.Remove(name)
		b.success(resp, -1, -1)
		b.log.Info("spooled batch delivered", "file", name, "accepted", resp.Accepted, "rejected", resp.Rejected)
	}
}

func (b *Batcher) toSpool(gz []byte, why string) {
	if b.spool == nil {
		b.log.Error("no spool, batch lost", "why", why)
		return
	}
	name, err := b.spool.Write(gz)
	if err != nil {
		b.log.Error("spool write failed, batch lost", "err", err, "why", why)
		return
	}
	b.log.Warn("batch spooled", "file", name, "bytes", len(gz), "why", why)
}

func (b *Batcher) fail(err error) {
	if b.backoff == 0 {
		b.backoff = BackoffMin
	} else {
		b.backoff *= 2
		if b.backoff > BackoffMax {
			b.backoff = BackoffMax
		}
	}
	b.nextTry = time.Now().Add(b.backoff)
	if b.hooks.OnFlushFailure != nil {
		b.hooks.OnFlushFailure()
	}
	b.log.Warn("batch delivery failed", "err", err, "retry_in", b.backoff)
}

func (b *Batcher) success(resp *protocol.BatchResponse, nm, nt int) {
	if b.backoff > 0 {
		b.log.Info("server reachable again")
	}
	b.backoff = 0
	if b.hooks.OnDelivered != nil {
		b.hooks.OnDelivered(time.Now())
	}
	if resp != nil && b.hooks.OnResponse != nil {
		b.hooks.OnResponse(resp)
	}
	if nm >= 0 {
		b.log.Debug("batch delivered", "measurements", nm, "traceroutes", nt, "accepted", resp.Accepted, "rejected", resp.Rejected)
	}
}

// InBackoff reports whether deliveries are currently suspended.
func (b *Batcher) InBackoff() bool { return b.backoff > 0 && time.Now().Before(b.nextTry) }