// Copyright (c) Tailscale Inc & contributors
// SPDX-License-Identifier: BSD-3-Clause

package compserve

import (
	"io"
	"io/fs"

	"github.com/klauspost/compress/zstd"
	"tailscale.com/util/zstdframe"
)

// Zstd serves zstd-precompressed variants ("*.zst", as generated by
// util/precompress). [Encoding.Decompress] lets ServeFile transcode the
// variant for clients that do not accept zstd.
var Zstd = Encoding{
	Token:      "zstd",
	Ext:        ".zst",
	Decompress: decompressZstd,
	Compress:   compressZstdLive,
}

// compressZstdLive wraps w in a streaming zstd compressor using the fastest
// compression level and a small window suitable for low-latency live
// compression. Close on the returned writer finishes the zstd frame and
// returns the encoder to the shared pool.
func compressZstdLive(w io.Writer) (io.WriteCloser, error) {
	enc, put := zstdframe.GetStreamingEncoder(
		zstdframe.FastestCompression,
		zstdframe.MaxWindowSize(64<<10),
	)
	// Encoder.Reset (streaming) cannot fail; only Decoder.Reset can.
	enc.Reset(w)
	return &pooledZstdWriter{Encoder: enc, put: put}, nil
}

// pooledZstdWriter is a zstd.Encoder writing to the underlying writer;
// Close finishes the frame and returns the encoder to the shared pool.
type pooledZstdWriter struct {
	put    func()
	closed bool
	*zstd.Encoder
}

func (p *pooledZstdWriter) Close() error {
	if p.closed {
		return nil
	}
	p.closed = true
	err := p.Encoder.Close()
	p.put()
	return err
}

// decompressZstd wraps f in a streaming zstd decompressor backed by the
// shared decoder pool in util/zstdframe (concurrency=1, so streaming decode
// neither spawns goroutines nor retains unbounded memory). It takes
// ownership of f.
func decompressZstd(f fs.File) (io.ReadCloser, error) {
	dec, put := zstdframe.GetDecoder()
	if err := dec.Reset(f); err != nil {
		put()
		f.Close()
		return nil, err
	}
	return &pooledZstdReader{f: f, put: put, Decoder: dec}, nil
}

// pooledZstdReader is a zstd.Decoder reading from f; Close returns the
// decoder to the shared pool and closes f.
type pooledZstdReader struct {
	f   fs.File
	put func()
	*zstd.Decoder
}

func (p *pooledZstdReader) Close() error {
	p.put()
	return p.f.Close()
}