package blobgen

import (
	"crypto/sha256"
	"errors"
	"fmt"
	"hash"
	"io"

	"filippo.io/age"
	"github.com/klauspost/compress/zstd"
)

// Reader wraps decompression and decryption with SHA256 verification
type Reader struct {
	reader       io.Reader
	decompressor *zstd.Decoder
	decryptor    io.Reader
	hasher       hash.Hash
	teeReader    io.Reader
	bytesRead    int64
}

// NewReader creates a new Reader that decrypts, decompresses, and verifies
// data. Every supplied identity is offered to age.Decrypt, so a blob
// encrypted to any one of them can be read.
func NewReader(r io.Reader, identities ...age.Identity) (*Reader, error) {
	// Create decryption reader
	decReader, err := age.Decrypt(r, identities...)
	if err != nil {
		return nil, fmt.Errorf("creating decryption reader: %w", err)
	}

	// Create decompression reader
	decompressor, err := zstd.NewReader(decReader)
	if err != nil {
		return nil, fmt.Errorf("creating decompression reader: %w", err)
	}

	// Create SHA256 hasher
	hasher := sha256.New()

	// Create tee reader that reads from decompressor and writes to hasher
	teeReader := io.TeeReader(decompressor, hasher)

	return &Reader{
		reader:       r,
		decompressor: decompressor,
		decryptor:    decReader,
		hasher:       hasher,
		teeReader:    teeReader,
	}, nil
}

// Read implements io.Reader
func (r *Reader) Read(p []byte) (int, error) {
	n, err := r.teeReader.Read(p)
	r.bytesRead += int64(n)

	// When the ciphertext is cut right after the age header plus its
	// 16-byte nonce, the age reader's first read fails with
	// io.ErrUnexpectedEOF, and the zstd decoder maps that to a clean
	// io.EOF at frame start. That makes a truncated stream look like a
	// valid empty one. Distinguish the two: on EOF, read once more from
	// the age reader. A genuine end leaves it at (0, io.EOF); a truncated
	// stream leaves its stored io.ErrUnexpectedEOF, which we surface.
	if errors.Is(err, io.EOF) {
		var probe [1]byte

		m, ageErr := r.decryptor.Read(probe[:])
		if m != 0 || !errors.Is(ageErr, io.EOF) {
			return n, io.ErrUnexpectedEOF
		}
	}

	return n, err
}

// Close closes the decompressor
func (r *Reader) Close() error {
	r.decompressor.Close()

	return nil
}

// Sum256 returns the single SHA-256 of the plaintext read so far. This is the
// first hash only; the stored object name is its double hash, which callers
// obtain by passing this digest to DoubleSHA256.
func (r *Reader) Sum256() []byte {
	return r.hasher.Sum(nil)
}

// BytesRead returns the number of uncompressed bytes read
func (r *Reader) BytesRead() int64 {
	return r.bytesRead
}