package glacier

import (
	"crypto/sha256"
	"io"

	"github.com/aws/aws-sdk-go/internal/sdkio"
)

const bufsize = 1024 * 1024

// Hash contains information about the tree-hash and linear hash of a
// Glacier payload. This structure is generated by ComputeHashes().
type Hash struct {
	TreeHash   []byte
	LinearHash []byte
}

// ComputeHashes computes the tree-hash and linear hash of a seekable reader r.
//
// See http://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html for more information.
func ComputeHashes(r io.ReadSeeker) Hash {
	start, _ := r.Seek(0, sdkio.SeekCurrent) // Read the whole stream
	defer r.Seek(start, sdkio.SeekStart)     // Rewind stream at end

	buf := make([]byte, bufsize)
	hashes := [][]byte{}
	hsh := sha256.New()

	for {
		// Build leaf nodes in 1MB chunks
		n, err := io.ReadAtLeast(r, buf, bufsize)
		if n == 0 {
			break
		}

		tmpHash := sha256.Sum256(buf[:n])
		hashes = append(hashes, tmpHash[:])
		hsh.Write(buf[:n]) // Track linear hash while we're at it

		if err != nil {
			break // This is the last chunk
		}
	}

	return Hash{
		LinearHash: hsh.Sum(nil),
		TreeHash:   ComputeTreeHash(hashes),
	}
}

// ComputeTreeHash builds a tree hash root node given a slice of
// hashes. Glacier tree hash to be derived from SHA256 hashes of 1MB
// chucks of the data.
//
// See http://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html for more information.
func ComputeTreeHash(hashes [][]byte) []byte {
	if hashes == nil || len(hashes) == 0 {
		return nil
	}

	for len(hashes) > 1 {
		tmpHashes := [][]byte{}

		for i := 0; i < len(hashes); i += 2 {
			if i+1 <= len(hashes)-1 {
				tmpHash := append(append([]byte{}, hashes[i]...), hashes[i+1]...)
				tmpSum := sha256.Sum256(tmpHash)
				tmpHashes = append(tmpHashes, tmpSum[:])
			} else {
				tmpHashes = append(tmpHashes, hashes[i])
			}
		}

		hashes = tmpHashes
	}

	return hashes[0]
}