package snapshot

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"

	"github.com/klauspost/compress/zstd"
	"sneak.berlin/go/vaultik/internal/blobgen"
)

// Manifest size bounds. A manifest lists one small entry per blob, and
// blobs are large (the default target is 10 GB), so even a manifest for a
// petabyte-scale backup is a few megabytes. These caps are far above any
// manifest the writer can emit, yet stop a crafted, highly compressible
// manifest from expanding without limit when decoded: the manifest is
// fetched from the store, which is not trusted, and json.Decode buffers
// the whole value in memory.
const (
	manifestMaxCompressed   = 256 * 1024 * 1024  // 256 MiB
	manifestMaxDecompressed = 1024 * 1024 * 1024 // 1 GiB
)

// Manifest represents the structure of a snapshot's blob manifest
//
//nolint:tagliatelle // snake_case is the established on-disk manifest format
type Manifest struct {
	SnapshotID          string     `json:"snapshot_id"`
	Timestamp           string     `json:"timestamp"`
	BlobCount           int        `json:"blob_count"`
	TotalCompressedSize int64      `json:"total_compressed_size"`
	Blobs               []BlobInfo `json:"blobs"`
}

// BlobInfo represents information about a single blob in the manifest
//
//nolint:tagliatelle // snake_case is the established on-disk manifest format
type BlobInfo struct {
	Hash           string `json:"hash"`
	CompressedSize int64  `json:"compressed_size"`
}

// DecodeManifest decodes a manifest from a reader containing compressed
// JSON, reading through byte limits on both the compressed input and the
// decompressed output so an untrusted manifest cannot exhaust memory.
func DecodeManifest(r io.Reader) (*Manifest, error) {
	return decodeManifest(r, manifestMaxCompressed, manifestMaxDecompressed)
}

// decodeManifest is DecodeManifest with explicit limits, so tests can drive
// the bounds with small inputs instead of gigabyte-scale ones.
func decodeManifest(
	r io.Reader, maxCompressed, maxDecompressed int64,
) (*Manifest, error) {
	// Decompress using zstd, bounding how many compressed bytes are read.
	zr, err := zstd.NewReader(blobgen.LimitReader(r, maxCompressed))
	if err != nil {
		return nil, fmt.Errorf("creating zstd reader: %w", err)
	}
	defer zr.Close()

	// Decode JSON manifest, bounding how far the compressed input may
	// expand: json.Decode buffers the whole value, so without this a
	// small, highly compressible manifest could expand to gigabytes.
	var manifest Manifest

	err = json.NewDecoder(blobgen.LimitReader(zr, maxDecompressed)).Decode(&manifest)
	if err != nil {
		return nil, fmt.Errorf("decoding manifest: %w", err)
	}

	return &manifest, nil
}

// EncodeManifest encodes a manifest to compressed JSON
func EncodeManifest(manifest *Manifest, compressionLevel int) ([]byte, error) {
	// Marshal to JSON
	jsonData, err := json.MarshalIndent(manifest, "", "  ")
	if err != nil {
		return nil, fmt.Errorf("marshaling manifest: %w", err)
	}

	// Compress using zstd
	var compressedBuf bytes.Buffer

	writer, err := zstd.NewWriter(&compressedBuf,
		zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(compressionLevel)))
	if err != nil {
		return nil, fmt.Errorf("creating zstd writer: %w", err)
	}

	_, err = writer.Write(jsonData)
	if err != nil {
		_ = writer.Close()

		return nil, fmt.Errorf("writing compressed data: %w", err)
	}

	err = writer.Close()
	if err != nil {
		return nil, fmt.Errorf("closing zstd writer: %w", err)
	}

	return compressedBuf.Bytes(), nil
}