package gozstd

import (
	"bytes"
	"fmt"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
	"testing"
)

// These tests prove the load-bearing interop property of the shim: the pure-Go
// klauspost backend decodes zstd frames produced by upstream libzstd — plain
// AND dictionary-compressed — which is exactly what dolt's NBS archive reader
// asks of gozstd at runtime (it decompresses .darc archives authored by the
// real `dolt`/libzstd). They shell out to the `zstd` CLI (libzstd) to author
// the fixtures and skip cleanly when it is absent, so CI without the tool still
// passes. The internal round-trip test needs no external tool.

func zstdBin(t *testing.T) string {
	t.Helper()
	p, err := exec.LookPath("zstd")
	if err != nil {
		t.Skip("zstd CLI (libzstd) not found; skipping libzstd-interop test")
	}
	return p
}

// TestDecompressLibzstdPlain: libzstd compresses, the shim decompresses.
func TestDecompressLibzstdPlain(t *testing.T) {
	zstd := zstdBin(t)
	orig := bytes.Repeat([]byte("the quick brown fox jumps over the lazy dog\n"), 500)

	dir := t.TempDir()
	in := filepath.Join(dir, "in")
	out := filepath.Join(dir, "in.zst")
	if err := os.WriteFile(in, orig, 0o644); err != nil {
		t.Fatal(err)
	}
	run(t, zstd, "-q", "-19", in, "-o", out)

	comp, err := os.ReadFile(out)
	if err != nil {
		t.Fatal(err)
	}
	got, err := Decompress(nil, comp)
	if err != nil {
		t.Fatalf("Decompress: %v", err)
	}
	if !bytes.Equal(got, orig) {
		t.Fatalf("plain: round-trip mismatch (%d vs %d bytes)", len(got), len(orig))
	}
}

// TestDecompressDictLibzstd: libzstd trains a dictionary and compresses with it,
// then the shim decompresses with the same dictionary via NewDDict —
// the exact archive-read path dolt exercises. This is the interop risk the
// feasibility review flagged (klauspost matches decoder dicts by embedded
// dict-ID, structured ZDICT format), proven end-to-end here.
func TestDecompressDictLibzstd(t *testing.T) {
	zstd := zstdBin(t)
	dir := t.TempDir()

	// Enough varied-but-similar samples for the trainer to build a dictionary.
	var samples []string
	for i := 0; i < 512; i++ {
		p := filepath.Join(dir, fmt.Sprintf("s%03d", i))
		line := fmt.Sprintf("issue memestudio-%04d: lorem ipsum dolor sit amet consectetur %d\n", i, i*7)
		if err := os.WriteFile(p, []byte(line), 0o644); err != nil {
			t.Fatal(err)
		}
		samples = append(samples, p)
	}
	dictPath := filepath.Join(dir, "dict")
	if out, err := exec.Command(zstd, append([]string{"--train", "--maxdict=8192", "-o", dictPath}, samples...)...).CombinedOutput(); err != nil {
		t.Skipf("zstd --train unavailable/failed (%v): %s", err, out)
	}
	dict, err := os.ReadFile(dictPath)
	if err != nil {
		t.Fatalf("read trained dict: %v", err)
	}

	orig := []byte("issue memestudio-0042: lorem ipsum dolor sit amet consectetur 294 — payload body\n")
	in := filepath.Join(dir, "payload")
	out := filepath.Join(dir, "payload.zst")
	if err := os.WriteFile(in, orig, 0o644); err != nil {
		t.Fatal(err)
	}
	run(t, zstd, "-q", "-19", "-D", dictPath, in, "-o", out)

	comp, err := os.ReadFile(out)
	if err != nil {
		t.Fatal(err)
	}

	dd, err := NewDDict(dict)
	if err != nil {
		t.Fatalf("NewDDict on libzstd-trained dict: %v", err)
	}
	defer dd.Release()
	got, err := DecompressDict(nil, comp, dd)
	if err != nil {
		t.Fatalf("DecompressDict on libzstd frame: %v", err)
	}
	if !bytes.Equal(got, orig) {
		t.Fatalf("dict: round-trip mismatch\n got=%q\nwant=%q", got, orig)
	}
}

// TestShimRoundTripPlain exercises the shim's own compress+decompress with no
// external tool and no dictionary.
func TestShimRoundTripPlain(t *testing.T) {
	orig := bytes.Repeat([]byte("parade lane rolling lined-up stalled past-stand\n"), 200)
	comp := Compress(nil, orig)
	got, err := Decompress(nil, comp)
	if err != nil {
		t.Fatalf("Decompress: %v", err)
	}
	if !bytes.Equal(got, orig) {
		t.Fatal("plain shim round-trip mismatch")
	}
}

// TestShimRoundTripDict exercises the shim's CompressDict + DecompressDict over
// a real structured dictionary. The dictionary is trained by libzstd (the zstd
// CLI) rather than the shim's own BuildDict, because BuildDict is off this
// project's runtime path (dolt only calls it from archive-writer/gc, which this
// binary never runs) and klauspost's trainer is stricter about corpus size than
// libzstd's ZDICT — training quality is irrelevant here; what matters is that
// the shim's encode+decode agree over a standard dictionary.
func TestShimRoundTripDict(t *testing.T) {
	zstd := zstdBin(t)
	dir := t.TempDir()

	var samples []string
	for i := 0; i < 512; i++ {
		p := filepath.Join(dir, fmt.Sprintf("s%03d", i))
		line := fmt.Sprintf("bead memestudio-%04d status open priority %d assignee eugene\n", i, i%4)
		if err := os.WriteFile(p, []byte(line), 0o644); err != nil {
			t.Fatal(err)
		}
		samples = append(samples, p)
	}
	dictPath := filepath.Join(dir, "dict")
	if out, err := exec.Command(zstd, append([]string{"--train", "--maxdict=8192", "-o", dictPath}, samples...)...).CombinedOutput(); err != nil {
		t.Skipf("zstd --train unavailable/failed (%v): %s", err, out)
	}
	dict, err := os.ReadFile(dictPath)
	if err != nil {
		t.Fatal(err)
	}

	orig := bytes.Repeat([]byte("bead memestudio-0042 status open priority 2 assignee eugene\n"), 50)
	cd, err := NewCDict(dict)
	if err != nil {
		t.Fatalf("NewCDict: %v", err)
	}
	defer cd.Release()
	dd, err := NewDDict(dict)
	if err != nil {
		t.Fatalf("NewDDict: %v", err)
	}
	defer dd.Release()

	dcomp := CompressDict(nil, orig, cd)
	dgot, err := DecompressDict(nil, dcomp, dd)
	if err != nil {
		t.Fatalf("DecompressDict: %v", err)
	}
	if !bytes.Equal(dgot, orig) {
		t.Fatal("dict shim round-trip mismatch")
	}
}

// TestCompressReadableByLibzstd proves the reverse interop direction: frames the
// shim produces are decodable by upstream libzstd (the zstd CLI), both plain and
// dictionary-compressed. dolt never consumes shim-written archives at runtime,
// but this confirms the shim emits standard RFC-8878 frames, not a klauspost
// dialect.
func TestCompressReadableByLibzstd(t *testing.T) {
	zstd := zstdBin(t)
	dir := t.TempDir()
	orig := bytes.Repeat([]byte("shim output must be standard zstd\n"), 300)

	// plain: shim compresses, zstd -d decompresses.
	comp := Compress(nil, orig)
	czst := filepath.Join(dir, "c.zst")
	if err := os.WriteFile(czst, comp, 0o644); err != nil {
		t.Fatal(err)
	}
	cout := filepath.Join(dir, "c.out")
	run(t, zstd, "-q", "-d", czst, "-o", cout)
	if got, _ := os.ReadFile(cout); !bytes.Equal(got, orig) {
		t.Fatal("plain: libzstd could not round-trip shim output")
	}

	// dict: shim compresses with a libzstd-trained dict, zstd -d -D decompresses.
	var samples []string
	for i := 0; i < 512; i++ {
		p := filepath.Join(dir, fmt.Sprintf("s%03d", i))
		if err := os.WriteFile(p, []byte(fmt.Sprintf("bead memestudio-%04d status open priority %d\n", i, i%4)), 0o644); err != nil {
			t.Fatal(err)
		}
		samples = append(samples, p)
	}
	dictPath := filepath.Join(dir, "dict")
	if out, err := exec.Command(zstd, append([]string{"--train", "--maxdict=8192", "-o", dictPath}, samples...)...).CombinedOutput(); err != nil {
		t.Skipf("zstd --train failed (%v): %s", err, out)
	}
	dict, err := os.ReadFile(dictPath)
	if err != nil {
		t.Fatal(err)
	}
	cd, err := NewCDict(dict)
	if err != nil {
		t.Fatalf("NewCDict: %v", err)
	}
	defer cd.Release()
	dorig := bytes.Repeat([]byte("bead memestudio-0042 status open priority 2\n"), 40)
	dcomp := CompressDict(nil, dorig, cd)
	dzst := filepath.Join(dir, "d.zst")
	if err := os.WriteFile(dzst, dcomp, 0o644); err != nil {
		t.Fatal(err)
	}
	dout := filepath.Join(dir, "d.out")
	run(t, zstd, "-q", "-d", "-D", dictPath, dzst, "-o", dout)
	if got, _ := os.ReadFile(dout); !bytes.Equal(got, dorig) {
		t.Fatal("dict: libzstd could not round-trip shim dict output")
	}
}

// TestAppendSemantics verifies the load-bearing dst-append contract: gozstd's
// Compress/Decompress/CompressDict/DecompressDict append to dst and return the
// grown slice, preserving any existing prefix. dolt passes non-empty buffers, so
// a klauspost mismatch here would silently corrupt its data.
func TestAppendSemantics(t *testing.T) {
	orig := []byte("append into an existing buffer without clobbering the prefix")
	prefix := []byte("PREFIX|")

	comp := Compress(append([]byte{}, prefix...), orig)
	if !bytes.HasPrefix(comp, prefix) {
		t.Fatal("Compress dropped the dst prefix")
	}
	got, err := Decompress(append([]byte{}, prefix...), comp[len(prefix):])
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(got, append(append([]byte{}, prefix...), orig...)) {
		t.Fatal("Decompress did not append after the dst prefix")
	}

	// same for the dict variants
	cd, dd := shimDict(t)
	defer cd.Release()
	defer dd.Release()
	dc := CompressDict(append([]byte{}, prefix...), orig, cd)
	if !bytes.HasPrefix(dc, prefix) {
		t.Fatal("CompressDict dropped the dst prefix")
	}
	dg, err := DecompressDict(append([]byte{}, prefix...), dc[len(prefix):], dd)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(dg, append(append([]byte{}, prefix...), orig...)) {
		t.Fatal("DecompressDict did not append after the dst prefix")
	}
}

// TestRoundTripEdgeCases covers empty, tiny, incompressible-random, and large
// inputs through the plain path.
func TestRoundTripEdgeCases(t *testing.T) {
	rng := rand.New(rand.NewSource(1)) // deterministic
	big := make([]byte, 4<<20)
	rng.Read(big)
	cases := map[string][]byte{
		"empty":          {},
		"one-byte":       {0x42},
		"incompressible": randBytes(rng, 64<<10),
		"large-random":   big,
		"large-text":     bytes.Repeat([]byte("mardi gras parade "), 300000),
	}
	for name, orig := range cases {
		got, err := Decompress(nil, Compress(nil, orig))
		if err != nil {
			t.Fatalf("%s: Decompress: %v", name, err)
		}
		if !bytes.Equal(got, orig) {
			t.Fatalf("%s: round-trip mismatch (%d vs %d bytes)", name, len(got), len(orig))
		}
	}
}

// TestConcurrent exercises the concurrency-safety the shim claims (dolt calls
// these from many goroutines): shared coders for plain Compress/Decompress and a
// single shared DDict for DecompressDict.
func TestConcurrent(t *testing.T) {
	cd, dd := shimDict(t)
	defer cd.Release()
	defer dd.Release()

	payload := bytes.Repeat([]byte("concurrent parade "), 500)
	dictFrame := CompressDict(nil, payload, cd)

	const n = 64
	errs := make(chan error, n)
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			// plain round-trip on distinct data
			src := []byte(fmt.Sprintf("goroutine-%d-%s", i, payload))
			if got, err := Decompress(nil, Compress(nil, src)); err != nil || !bytes.Equal(got, src) {
				errs <- fmt.Errorf("plain g%d: err=%v equal=%v", i, err, bytes.Equal(got, src))
				return
			}
			// concurrent DecompressDict on the shared DDict
			if got, err := DecompressDict(nil, dictFrame, dd); err != nil || !bytes.Equal(got, payload) {
				errs <- fmt.Errorf("dict g%d: err=%v equal=%v", i, err, bytes.Equal(got, payload))
				return
			}
		}(i)
	}
	wg.Wait()
	close(errs)
	for err := range errs {
		t.Error(err)
	}
}

// TestErrorPaths confirms bad input yields an error, never a panic — dolt's
// archive reader handles the error return.
func TestErrorPaths(t *testing.T) {
	if _, err := Decompress(nil, []byte("not a zstd frame at all")); err == nil {
		t.Fatal("Decompress of garbage: want error, got nil")
	}
	// Empty input is a benign divergence from libzstd (which errors): klauspost
	// yields empty output, no error, no panic. dolt never decompresses empty
	// input, so this is documented rather than treated as a fault.
	if got, err := Decompress(nil, nil); err != nil || len(got) != 0 {
		t.Fatalf("Decompress of empty input: want (empty, nil), got (%d bytes, %v)", len(got), err)
	}
	// A frame compressed with dict A must not decode under dict B (dict-ID
	// mismatch) — and must fail cleanly rather than return wrong bytes.
	a := buildShimDict(t, "alpha")
	b := buildShimDict(t, "bravo")
	cdA, _ := NewCDict(a)
	defer cdA.Release()
	ddB, _ := NewDDict(b)
	defer ddB.Release()
	frame := CompressDict(nil, bytes.Repeat([]byte("alpha payload "), 50), cdA)
	if _, err := DecompressDict(nil, frame, ddB); err == nil {
		t.Fatal("DecompressDict under the wrong dict: want error, got nil")
	}
}

// --- test helpers ---

// shimDict builds a matching CDict/DDict pair from a libzstd-trained dictionary
// (via the zstd CLI); it skips the calling test if the CLI is unavailable.
func shimDict(t *testing.T) (*CDict, *DDict) {
	t.Helper()
	d := buildShimDict(t, "shared")
	cd, err := NewCDict(d)
	if err != nil {
		t.Fatalf("NewCDict: %v", err)
	}
	dd, err := NewDDict(d)
	if err != nil {
		t.Fatalf("NewDDict: %v", err)
	}
	return cd, dd
}

// buildShimDict trains a distinct structured dictionary via the zstd CLI, seeded
// by tag so different tags yield different dict-IDs.
func buildShimDict(t *testing.T, tag string) []byte {
	t.Helper()
	zstd := zstdBin(t)
	dir := t.TempDir()
	var samples []string
	for i := 0; i < 512; i++ {
		p := filepath.Join(dir, fmt.Sprintf("s%03d", i))
		if err := os.WriteFile(p, []byte(fmt.Sprintf("%s bead-%04d status open priority %d assignee eugene\n", tag, i, i%4)), 0o644); err != nil {
			t.Fatal(err)
		}
		samples = append(samples, p)
	}
	dictPath := filepath.Join(dir, "dict")
	if out, err := exec.Command(zstd, append([]string{"--train", "--maxdict=8192", "-o", dictPath}, samples...)...).CombinedOutput(); err != nil {
		t.Skipf("zstd --train failed (%v): %s", err, out)
	}
	d, err := os.ReadFile(dictPath)
	if err != nil {
		t.Fatal(err)
	}
	return d
}

func randBytes(rng *rand.Rand, n int) []byte {
	b := make([]byte, n)
	rng.Read(b)
	return b
}

func run(t *testing.T, name string, args ...string) {
	t.Helper()
	if out, err := exec.Command(name, args...).CombinedOutput(); err != nil {
		t.Fatalf("%s %v: %v\n%s", name, args, err, out)
	}
}