package vaultik

import (
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	"io"

	"filippo.io/age"
	"sneak.berlin/go/vaultik/internal/blobgen"
)

// errBlobHashMismatch is returned when a fetched blob's content hash does
// not match the expected double-SHA-256 hash.
var errBlobHashMismatch = errors.New("blob hash mismatch")

// errBlobNotFullyRead is returned when the verifying reader is closed
// before its plaintext reached EOF. The hash can only be checked once
// the whole stream has been read, so an early or short-read close must
// fail rather than silently skip verification.
var errBlobNotFullyRead = errors.New(
	"blob closed before fully read; hash not verified")

// hashVerifyReader wraps a blobgen.Reader and verifies the double-SHA-256 hash
// of decrypted plaintext when Close is called. It reuses the hash that
// blobgen.Reader already computes internally via its TeeReader, avoiding
// redundant SHA-256 computation.
type hashVerifyReader struct {
	reader   *blobgen.Reader // underlying decrypted blob reader (has internal hasher)
	limited  io.Reader       // reader bounded to the blob's recorded plaintext size
	fetcher  io.ReadCloser   // raw fetched stream (closed on Close)
	blobHash string          // expected double-SHA-256 hex
	done     bool            // EOF reached
}

func (h *hashVerifyReader) Read(p []byte) (int, error) {
	n, err := h.limited.Read(p)
	if errors.Is(err, io.EOF) {
		h.done = true
	}

	return n, err
}

// Close closes the underlying readers and verifies the blob hash. The
// hash check cannot be skipped: closing before the plaintext reached
// EOF (a short read or an early close) is an error, so a caller can
// never obtain unverified blob bytes.
func (h *hashVerifyReader) Close() error {
	readerErr := h.reader.Close()
	fetcherErr := h.fetcher.Close()

	if !h.done {
		return errBlobNotFullyRead
	}

	actualHashHex := hex.EncodeToString(blobgen.DoubleSHA256(h.reader.Sum256()))
	if actualHashHex != h.blobHash {
		return fmt.Errorf("%w: expected %s, got %s",
			errBlobHashMismatch, shortHash(h.blobHash), shortHash(actualHashHex))
	}

	if readerErr != nil {
		return readerErr
	}

	return fetcherErr
}

// FetchAndDecryptBlob downloads a blob, decrypts and decompresses it, and
// returns a streaming reader that computes the double-SHA-256 hash on the fly.
// The hash is verified when the returned reader is closed (after fully reading).
// This avoids buffering the entire blob in memory.
//
// maxPlaintextSize is the blob's uncompressed_size as recorded in the
// snapshot database. Decompression stops with blobgen.ErrOutputTooLarge
// once the plaintext exceeds it, so a tampered blob cannot expand without
// limit — using the recorded size, not the restoring host's
// blob_size_limit, since that config may differ from the backup host's.
func (v *Vaultik) FetchAndDecryptBlob(
	ctx context.Context, blobHash string, maxPlaintextSize int64,
	identities ...age.Identity,
) (io.ReadCloser, error) {
	rc, err := v.FetchBlob(ctx, blobHash)
	if err != nil {
		return nil, err
	}

	reader, err := blobgen.NewReader(rc, identities...)
	if err != nil {
		_ = rc.Close()

		return nil, fmt.Errorf("creating blob reader: %w", err)
	}

	return &hashVerifyReader{
		reader:   reader,
		limited:  blobgen.LimitReader(reader, maxPlaintextSize),
		fetcher:  rc,
		blobHash: blobHash,
	}, nil
}

// FetchBlob downloads a blob and returns a reader for the encrypted data.
// The reader is closed as soon as ctx is cancelled, which ends a read
// blocked inside it.
func (v *Vaultik) FetchBlob(
	ctx context.Context, blobHash string,
) (io.ReadCloser, error) {
	// blobHash reaches here from the snapshot database, which is not
	// trusted. Reject a malformed hash before it is spliced into a storage
	// path (blobHash[:2]/blobHash[2:4]) or a fetch is attempted.
	if !isBlobHash(blobHash) {
		return nil, fmt.Errorf("%w: %s", errInvalidBlobHash, shortHash(blobHash))
	}

	blobPath := fmt.Sprintf("blobs/%s/%s/%s", blobHash[:2], blobHash[2:4], blobHash)

	rc, err := v.Storage.Get(ctx, blobPath)
	if err != nil {
		return nil, fmt.Errorf("downloading blob %s: %w", shortHash(blobHash), err)
	}

	return newCloseOnCancelReader(ctx, rc), nil
}

// closeOnCancelReader closes the reader it wraps when ctx is cancelled.
// Not every backend's reader stops on cancel by itself (the file
// backend's does not), and a read stalled on a slow or hung source would
// otherwise outlast the shutdown timeout, so the process would exit
// before restore or deep verify removed its decrypted temporary files.
type closeOnCancelReader struct {
	io.ReadCloser

	ctx  context.Context //nolint:containedctx // Read reports its error
	stop func() bool
}

func newCloseOnCancelReader(
	ctx context.Context, rc io.ReadCloser,
) *closeOnCancelReader {
	return &closeOnCancelReader{
		ReadCloser: rc,
		ctx:        ctx,
		stop:       context.AfterFunc(ctx, func() { _ = rc.Close() }),
	}
}

// Read returns the context's error, not the closed reader's, once ctx
// has been cancelled.
func (r *closeOnCancelReader) Read(p []byte) (int, error) {
	n, err := r.ReadCloser.Read(p)
	if err != nil && r.ctx.Err() != nil {
		return n, r.ctx.Err()
	}

	return n, err
}

// Close closes the wrapped reader, unless the cancel already closed it.
func (r *closeOnCancelReader) Close() error {
	if !r.stop() {
		return nil
	}

	return r.ReadCloser.Close()
}