package lem

import (
	"os"
	"path/filepath"
	"testing"
)

func TestDecompressZstd(t *testing.T) {
	// Create a temp .jsonl file and compress it
	dir := t.TempDir()
	src := filepath.Join(dir, "test.jsonl")
	os.WriteFile(src, []byte(`{"messages":[{"role":"user","content":"hello"}]}`+"\n"), 0644)

	// Compress with Go zstd
	zstPath := src + ".zst"
	if err := compressFileZstd(src, zstPath); err != nil {
		t.Fatal(err)
	}
	os.Remove(src) // remove original

	// Decompress
	outPath := filepath.Join(dir, "out.jsonl")
	if err := decompressZstd(zstPath, outPath); err != nil {
		t.Fatal(err)
	}

	data, err := os.ReadFile(outPath)
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != `{"messages":[{"role":"user","content":"hello"}]}`+"\n" {
		t.Fatalf("unexpected content: %q", string(data))
	}
}

func TestWalkZstFiles(t *testing.T) {
	dir := t.TempDir()
	// Create a nested .jsonl.zst
	sub := filepath.Join(dir, "sub")
	os.MkdirAll(sub, 0755)
	src := filepath.Join(sub, "data.jsonl")
	os.WriteFile(src, []byte("line1\nline2\n"), 0644)
	compressFileZstd(src, src+".zst")
	os.Remove(src)

	// Also create a .json file (should be ignored)
	os.WriteFile(filepath.Join(sub, "probes.json"), []byte("[]"), 0644)

	var found []string
	walkZstFiles(dir, func(zstPath string) error {
		found = append(found, zstPath)
		return nil
	})
	if len(found) != 1 {
		t.Fatalf("expected 1 .zst file, got %d: %v", len(found), found)
	}
}