// zstd.go — zstd compression/decompression for cold storage JSONL files.
package lem

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"github.com/klauspost/compress/zstd"
)

// decompressZstd reads a .zst file and writes the decompressed content to dst.
func decompressZstd(src, dst string) error {
	f, err := os.Open(src)
	if err != nil {
		return fmt.Errorf("open %s: %w", src, err)
	}
	defer f.Close()

	dec, err := zstd.NewReader(f)
	if err != nil {
		return fmt.Errorf("zstd reader: %w", err)
	}
	defer dec.Close()

	out, err := os.Create(dst)
	if err != nil {
		return fmt.Errorf("create %s: %w", dst, err)
	}
	defer out.Close()

	if _, err := io.Copy(out, dec); err != nil {
		return fmt.Errorf("decompress %s: %w", src, err)
	}
	return nil
}

// compressFileZstd compresses src to dst with zstd level 3.
func compressFileZstd(src, dst string) error {
	f, err := os.Open(src)
	if err != nil {
		return fmt.Errorf("open %s: %w", src, err)
	}
	defer f.Close()

	out, err := os.Create(dst)
	if err != nil {
		return fmt.Errorf("create %s: %w", dst, err)
	}
	defer out.Close()

	enc, err := zstd.NewWriter(out, zstd.WithEncoderLevel(zstd.SpeedDefault))
	if err != nil {
		return fmt.Errorf("zstd writer: %w", err)
	}

	if _, err := io.Copy(enc, f); err != nil {
		enc.Close()
		return fmt.Errorf("compress %s: %w", src, err)
	}
	return enc.Close()
}

// walkZstFiles walks a directory tree and calls fn for each .jsonl.zst file.
func walkZstFiles(root string, fn func(zstPath string) error) error {
	return filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil || info.IsDir() {
			return err
		}
		if strings.HasSuffix(path, ".jsonl.zst") {
			return fn(path)
		}
		return nil
	})
}