package main

import (
	"archive/tar"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/klauspost/compress/s2"
	"github.com/nats-io/nats.go"
)

// The JetStream API's subjects this program speaks. Each is granted to the bus's own module and
// nothing else is (mesh-controller internal/broker, BusSnapshotGrants); a test holds the two lists
// equal against a real server.
const (
	apiStreamNames    = "$JS.API.STREAM.NAMES"
	apiStreamInfo     = "$JS.API.STREAM.INFO."
	apiStreamSnapshot = "$JS.API.STREAM.SNAPSHOT."
	apiStreamRestore  = "$JS.API.STREAM.RESTORE."
)

// ChunkSize is what the server is asked to send at a time: its own default. The server sends up to
// an 8 MiB window ahead of the acknowledgements and waits for them past it, so one stream costs the
// server at most that much in flight, whatever its size.
const ChunkSize = 128 * 1024

// SnapshotOptions bounds a snapshot.
type SnapshotOptions struct {
	// StreamTimeout is the longest one stream may take; past it the snapshot fails, naming it.
	StreamTimeout time.Duration
	// Pause is a rest between two streams, so a night's snapshot is never one long burst.
	Pause time.Duration
	// TempDir is where one stream's archive waits until it is written out; the system's when empty.
	TempDir string
	Log     func(format string, args ...any)
}

// Manifest is the first entry of an archive: what was taken, when, how long it took, and the checksum
// of every file beside it. It is what `verify` and `restore` hold the archive to.
type Manifest struct {
	Format   int       `json:"format"`
	Server   string    `json:"server_version"`
	Started  time.Time `json:"started"`
	Finished time.Time `json:"finished"`
	Seconds  float64   `json:"seconds"`
	// Messages and StreamBytes are the streams' own totals at their snapshots; ArchiveBytes what the
	// archives take, compressed.
	Messages     uint64        `json:"messages"`
	StreamBytes  int64         `json:"stream_bytes"`
	ArchiveBytes int64         `json:"archive_bytes"`
	Streams      []StreamEntry `json:"streams"`
	// Skipped are streams that hold nothing across a restart (memory storage), said rather than
	// silently left out.
	Skipped []Skipped `json:"skipped,omitempty"`
}

// StreamEntry is one stream as it was snapshotted.
type StreamEntry struct {
	Name      string   `json:"name"`
	Subjects  []string `json:"subjects,omitempty"`
	Retention string   `json:"retention,omitempty"`
	Messages  uint64   `json:"messages"`
	Bytes     uint64   `json:"bytes"`
	FirstSeq  uint64   `json:"first_seq"`
	LastSeq   uint64   `json:"last_seq"`
	Deleted   int      `json:"deleted,omitempty"`
	Consumers int      `json:"consumers"`
	// Meta is the server's own description of the stream (its configuration and state), in the
	// shape the restore API and the nats CLI take back: backup.json.
	Meta       string `json:"meta"`
	MetaSHA256 string `json:"meta_sha256"`
	// Archive is the server's snapshot of the stream, as the server sent it: stream.tar.s2.
	Archive       string  `json:"archive"`
	ArchiveBytes  int64   `json:"archive_bytes"`
	ArchiveSHA256 string  `json:"archive_sha256"`
	Seconds       float64 `json:"seconds"`
}

// Skipped is a stream not snapshotted, and why.
type Skipped struct {
	Name string `json:"name"`
	Why  string `json:"why"`
}

// ManifestName is the archive's first entry.
const ManifestName = "manifest.json"

// apiError is the error every JetStream API answer may carry.
type apiError struct {
	Code        int    `json:"code"`
	ErrCode     int    `json:"err_code"`
	Description string `json:"description"`
}

func (e *apiError) Error() string { return fmt.Sprintf("%s (%d)", e.Description, e.Code) }

// streamState is what this program reads of a stream's state; the whole of it is kept as the server
// gave it, in the stream's meta.
type streamState struct {
	Messages  uint64 `json:"messages"`
	Bytes     uint64 `json:"bytes"`
	FirstSeq  uint64 `json:"first_seq"`
	LastSeq   uint64 `json:"last_seq"`
	NumDelete int    `json:"num_deleted"`
	Consumers int    `json:"consumer_count"`
}

type streamConfig struct {
	Name      string   `json:"name"`
	Subjects  []string `json:"subjects"`
	Retention string   `json:"retention"`
	Storage   string   `json:"storage"`
}

// request asks the JetStream API and decodes its answer, refusing an answer that carries an error.
func request(ctx context.Context, nc *nats.Conn, subject string, body any, into any) error {
	var payload []byte
	if body != nil {
		var err error
		if payload, err = json.Marshal(body); err != nil {
			return err
		}
	}
	rctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	msg, err := nc.RequestWithContext(rctx, subject, payload)
	if err != nil {
		if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, nats.ErrTimeout) || errors.Is(err, nats.ErrNoResponders) {
			return fmt.Errorf("%s was not answered (%v): the bus refuses a subject this user is not granted by not answering it", subject, err)
		}
		return fmt.Errorf("%s: %w", subject, err)
	}
	var wrapped struct {
		Error *apiError `json:"error"`
	}
	if err := json.Unmarshal(msg.Data, &wrapped); err != nil {
		return fmt.Errorf("%s answered something that is not JSON: %q", subject, truncate(msg.Data))
	}
	if wrapped.Error != nil {
		return fmt.Errorf("%s: %w", subject, wrapped.Error)
	}
	if into != nil {
		return json.Unmarshal(msg.Data, into)
	}
	return nil
}

func truncate(b []byte) string {
	if len(b) > 200 {
		return string(b[:200]) + "…"
	}
	return string(b)
}

// StreamNames is every stream in the account, in name order, all pages of them.
func StreamNames(ctx context.Context, nc *nats.Conn) ([]string, error) {
	var names []string
	for {
		var page struct {
			Total   int      `json:"total"`
			Offset  int      `json:"offset"`
			Limit   int      `json:"limit"`
			Streams []string `json:"streams"`
		}
		if err := request(ctx, nc, apiStreamNames, map[string]int{"offset": len(names)}, &page); err != nil {
			return nil, err
		}
		names = append(names, page.Streams...)
		if len(page.Streams) == 0 || len(names) >= page.Total {
			break
		}
	}
	sort.Strings(names)
	return names, nil
}

// Snapshot writes every stream, one at a time, as a tar: the manifest first, then each stream's
// meta and archive under streams/<name>/. Any stream failing fails the whole: a partial copy of the
// bus written where a whole one is expected is the silent failure this exists to prevent.
func Snapshot(ctx context.Context, nc *nats.Conn, out io.Writer, o SnapshotOptions) (*Manifest, error) {
	if o.StreamTimeout <= 0 {
		o.StreamTimeout = 10 * time.Minute
	}
	if o.Log == nil {
		o.Log = func(string, ...any) {}
	}
	m := &Manifest{Format: 1, Server: nc.ConnectedServerVersion(), Started: time.Now().UTC()}
	names, err := StreamNames(ctx, nc)
	if err != nil {
		return nil, err
	}
	if len(names) == 0 {
		// The mesh's own streams always exist once the controller has run; a bus with none is not
		// one to call backed up.
		return nil, errors.New("the bus holds no streams; a snapshot of nothing would pass for a backup")
	}

	type taken struct {
		entry   StreamEntry
		meta    []byte
		archive *os.File
	}
	var files []taken
	defer func() {
		for _, t := range files {
			t.archive.Close()
			os.Remove(t.archive.Name())
		}
	}()

	for i, name := range names {
		if i > 0 && o.Pause > 0 {
			select {
			case <-time.After(o.Pause):
			case <-ctx.Done():
				return nil, ctx.Err()
			}
		}
		var info struct {
			Config streamConfig `json:"config"`
		}
		if err := request(ctx, nc, apiStreamInfo+name, nil, &info); err != nil {
			return nil, err
		}
		if info.Config.Storage == "memory" {
			m.Skipped = append(m.Skipped, Skipped{Name: name, Why: "memory storage: it holds nothing across a restart, and the server cannot snapshot it"})
			o.Log("%s: skipped, memory storage", name)
			continue
		}
		f, err := os.CreateTemp(o.TempDir, "mesh-nats-snapshot-*")
		if err != nil {
			return nil, err
		}
		files = append(files, taken{archive: f})
		sctx, cancel := context.WithTimeout(ctx, o.StreamTimeout)
		entry, meta, err := snapshotStream(sctx, nc, name, f)
		cancel()
		if err != nil {
			return nil, fmt.Errorf("stream %s: %w", name, err)
		}
		files[len(files)-1].entry, files[len(files)-1].meta = entry, meta
		m.Messages += entry.Messages
		m.StreamBytes += int64(entry.Bytes)
		m.ArchiveBytes += entry.ArchiveBytes
		o.Log("%s: %d message(s), %s, last sequence %d, %d consumer(s) → %s in %.2fs",
			name, entry.Messages, human(int64(entry.Bytes)), entry.LastSeq, entry.Consumers, human(entry.ArchiveBytes), entry.Seconds)
	}
	m.Finished = time.Now().UTC()
	m.Seconds = m.Finished.Sub(m.Started).Seconds()
	for _, t := range files {
		m.Streams = append(m.Streams, t.entry)
	}

	tw := tar.NewWriter(out)
	manifest, err := json.MarshalIndent(m, "", "  ")
	if err != nil {
		return nil, err
	}
	if err := writeEntry(tw, ManifestName, append(manifest, '\n'), m.Finished); err != nil {
		return nil, err
	}
	for _, t := range files {
		if err := writeEntry(tw, t.entry.Meta, t.meta, m.Finished); err != nil {
			return nil, err
		}
		if _, err := t.archive.Seek(0, io.SeekStart); err != nil {
			return nil, err
		}
		hdr := &tar.Header{Name: t.entry.Archive, Mode: 0o600, Size: t.entry.ArchiveBytes, ModTime: m.Finished, Format: tar.FormatPAX}
		if err := tw.WriteHeader(hdr); err != nil {
			return nil, err
		}
		if _, err := io.Copy(tw, t.archive); err != nil {
			return nil, err
		}
	}
	if err := tw.Close(); err != nil {
		return nil, err
	}
	return m, nil
}

func writeEntry(tw *tar.Writer, name string, body []byte, at time.Time) error {
	if err := tw.WriteHeader(&tar.Header{Name: name, Mode: 0o600, Size: int64(len(body)), ModTime: at, Format: tar.FormatPAX}); err != nil {
		return err
	}
	_, err := tw.Write(body)
	return err
}

// snapshotStream takes one stream's snapshot into f: the server's chunks, each acknowledged as it
// arrives so the server's window keeps moving, until the server's last, empty message. It answers
// the stream's entry (the archive's size and checksum included) and its meta.
//
// The protocol is the server's (nats-server jetstream_api.go, streamSnapshot): the request names a
// subject to deliver to; the server answers with the stream's configuration and its state at the
// snapshot, then sends the archive in chunks whose reply subject is the acknowledgement it waits for
// past its window — two seconds without one and it gives up with "408 No Flow Response". The end is
// an empty message whose status header says 204, or the error that ended it.
func snapshotStream(ctx context.Context, nc *nats.Conn, name string, f *os.File) (StreamEntry, []byte, error) {
	started := time.Now()
	deliver := nc.NewRespInbox()

	hash := sha256.New()
	var written int64
	var mu sync.Mutex
	var failed error
	done := make(chan struct{})
	var once sync.Once
	finish := func(err error) {
		mu.Lock()
		if failed == nil {
			failed = err
		}
		mu.Unlock()
		once.Do(func() { close(done) })
	}
	sub, err := nc.Subscribe(deliver, func(msg *nats.Msg) {
		if len(msg.Data) == 0 {
			if status := msg.Header.Get("Status"); status != "" && status != "204" {
				finish(fmt.Errorf("the server ended the snapshot: %s %s", status, msg.Header.Get("Description")))
				return
			}
			finish(nil)
			return
		}
		// Acknowledged before it is written: the chunk is already in memory, and a slow disk must not
		// read to the server as a reader that went away.
		if msg.Reply != "" {
			if err := nc.Publish(msg.Reply, nil); err != nil {
				finish(fmt.Errorf("acknowledging a chunk: %w", err))
				return
			}
		}
		if _, err := f.Write(msg.Data); err != nil {
			finish(fmt.Errorf("writing the archive: %w", err))
			return
		}
		hash.Write(msg.Data)
		written += int64(len(msg.Data))
	})
	if err != nil {
		return StreamEntry{}, nil, err
	}
	defer sub.Unsubscribe()
	// Every chunk is held until the callback has it; the window bounds how many there can be.
	if err := sub.SetPendingLimits(-1, -1); err != nil {
		return StreamEntry{}, nil, err
	}
	if err := nc.FlushWithContext(ctx); err != nil {
		return StreamEntry{}, nil, err
	}

	var resp struct {
		Config json.RawMessage `json:"config"`
		State  json.RawMessage `json:"state"`
	}
	req := map[string]any{"deliver_subject": deliver, "chunk_size": ChunkSize}
	if err := request(ctx, nc, apiStreamSnapshot+name, req, &resp); err != nil {
		return StreamEntry{}, nil, err
	}
	select {
	case <-done:
	case <-ctx.Done():
		return StreamEntry{}, nil, fmt.Errorf("not finished within its bound: %w", ctx.Err())
	}
	// The subscription is drained before the archive is read back, so no callback is still writing.
	_ = sub.Unsubscribe()
	mu.Lock()
	err = failed
	mu.Unlock()
	if err != nil {
		return StreamEntry{}, nil, err
	}
	if written == 0 {
		return StreamEntry{}, nil, errors.New("the server sent an empty archive")
	}
	if err := f.Sync(); err != nil {
		return StreamEntry{}, nil, err
	}

	var cfg streamConfig
	var state streamState
	if err := json.Unmarshal(resp.Config, &cfg); err != nil {
		return StreamEntry{}, nil, fmt.Errorf("its configuration: %w", err)
	}
	if err := json.Unmarshal(resp.State, &state); err != nil {
		return StreamEntry{}, nil, fmt.Errorf("its state: %w", err)
	}
	// The archive reads to its end — s2 frames, then a tar — before it is called a snapshot.
	if _, err := f.Seek(0, io.SeekStart); err != nil {
		return StreamEntry{}, nil, err
	}
	if err := readsThrough(f); err != nil {
		return StreamEntry{}, nil, fmt.Errorf("the archive the server sent does not read: %w", err)
	}

	// The meta, in the shape the restore API takes and the nats CLI writes (backup.json): the
	// server's own configuration and state, kept as raw JSON so a field this program does not know
	// is never dropped on the way back.
	meta, err := json.MarshalIndent(map[string]json.RawMessage{"config": resp.Config, "state": resp.State}, "", "  ")
	if err != nil {
		return StreamEntry{}, nil, err
	}
	metaSum := sha256.Sum256(meta)
	return StreamEntry{
		Name: name, Subjects: cfg.Subjects, Retention: cfg.Retention,
		Messages: state.Messages, Bytes: state.Bytes, FirstSeq: state.FirstSeq, LastSeq: state.LastSeq,
		Deleted: state.NumDelete, Consumers: state.Consumers,
		Meta: "streams/" + name + "/backup.json", MetaSHA256: hex.EncodeToString(metaSum[:]),
		Archive: "streams/" + name + "/stream.tar.s2", ArchiveBytes: written,
		ArchiveSHA256: hex.EncodeToString(hash.Sum(nil)), Seconds: time.Since(started).Seconds(),
	}, meta, nil
}

// readsThrough is whether a stream's archive decompresses and lists to its end.
func readsThrough(r io.Reader) error {
	tr := tar.NewReader(s2.NewReader(r))
	for {
		_, err := tr.Next()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		if _, err := io.Copy(io.Discard, tr); err != nil {
			return err
		}
	}
}

// validName is a stream name that is one safe path component.
func validName(name string) bool {
	return name != "" && name != "." && name != ".." && !strings.ContainsAny(name, "/\\\x00 *>")
}