package kdbx

import (
	"encoding/binary"
	"errors"
	"fmt"
)

// Argon2 (RFC 9106), the key derivation function KeePass 2 uses by default
// for KDBX 4. Implemented here because the standard library has no Argon2
// and this package carries no third-party dependencies.

const (
	argon2BlockSize  = 1024 // bytes
	argon2BlockWords = argon2BlockSize / 8
	argon2SyncPoints = 4
	argon2Version    = 0x13

	// argon2MaxMemory bounds what a file may ask the server to allocate. A
	// malicious or corrupt header must not be able to exhaust memory.
	argon2MaxMemory = 1 << 21 // KiB, i.e. 2 GiB
)

type argon2Block [argon2BlockWords]uint64

// argon2Key derives tagLen bytes. variant is argon2d or argon2id.
func argon2Key(password, salt []byte, iterations, memoryKiB uint32, parallelism uint8,
	tagLen uint32, variant int) ([]byte, error) {
	return argon2Full(password, salt, nil, nil, iterations, memoryKiB, parallelism, tagLen, variant)
}

// argon2Full is the complete function including the secret key and
// associated data, which KDBX never uses but the RFC's test vectors do.
func argon2Full(password, salt, secret, assoc []byte, iterations, memoryKiB uint32,
	parallelism uint8, tagLen uint32, variant int) ([]byte, error) {
	if parallelism == 0 || iterations == 0 || tagLen < 4 {
		return nil, errors.New("kdbx: Argon2 parameters out of range")
	}
	if len(salt) < 8 {
		return nil, errors.New("kdbx: Argon2 salt is too short")
	}
	if memoryKiB > argon2MaxMemory {
		return nil, fmt.Errorf("kdbx: Argon2 asks for %d KiB of memory, refusing above %d",
			memoryKiB, argon2MaxMemory)
	}
	lanes := uint32(parallelism)
	if memoryKiB < 8*lanes {
		memoryKiB = 8 * lanes
	}
	// Round the block count down so every lane holds the same whole segments.
	blocks := memoryKiB / (argon2SyncPoints * lanes) * (argon2SyncPoints * lanes)
	laneLen := blocks / lanes
	segLen := laneLen / argon2SyncPoints

	h0 := argon2InitialHash(password, salt, secret, assoc, iterations, memoryKiB,
		lanes, tagLen, variant)

	mem := make([]argon2Block, blocks)
	for lane := uint32(0); lane < lanes; lane++ {
		for i := uint32(0); i < 2; i++ {
			var tail [8]byte
			binary.LittleEndian.PutUint32(tail[0:], i)
			binary.LittleEndian.PutUint32(tail[4:], lane)
			b := argon2Hash(argon2BlockSize, append(append([]byte{}, h0...), tail[:]...))
			loadBlock(&mem[lane*laneLen+i], b)
		}
	}

	for pass := uint32(0); pass < iterations; pass++ {
		for slice := uint32(0); slice < argon2SyncPoints; slice++ {
			for lane := uint32(0); lane < lanes; lane++ {
				argon2Segment(mem, pass, lane, slice, lanes, laneLen, segLen,
					iterations, variant)
			}
		}
	}

	var final argon2Block
	for lane := uint32(0); lane < lanes; lane++ {
		b := &mem[lane*laneLen+laneLen-1]
		for i := range final {
			final[i] ^= b[i]
		}
	}
	return argon2Hash(int(tagLen), storeBlock(&final)), nil
}

func argon2InitialHash(password, salt, secret, assoc []byte, iterations, memoryKiB,
	lanes, tagLen uint32, variant int) []byte {
	h := newBlake2b(64, nil)
	var u [4]byte
	put := func(v uint32) {
		binary.LittleEndian.PutUint32(u[:], v)
		h.Write(u[:])
	}
	putBytes := func(b []byte) {
		put(uint32(len(b)))
		h.Write(b)
	}
	put(lanes)
	put(tagLen)
	put(memoryKiB)
	put(iterations)
	put(argon2Version)
	put(uint32(variant))
	putBytes(password)
	putBytes(salt)
	putBytes(secret)
	putBytes(assoc)
	return h.Sum()
}

// argon2Hash is the RFC's variable-length hash H', which chains BLAKE2b when
// more than 64 bytes are wanted.
func argon2Hash(outLen int, input []byte) []byte {
	var prefix [4]byte
	binary.LittleEndian.PutUint32(prefix[:], uint32(outLen))
	if outLen <= 64 {
		return blake2bSum(outLen, prefix[:], input)
	}
	out := make([]byte, 0, outLen)
	v := blake2bSum(64, prefix[:], input)
	out = append(out, v[:32]...)
	for len(out)+64 < outLen {
		v = blake2bSum(64, v)
		out = append(out, v[:32]...)
	}
	v = blake2bSum(outLen-len(out), v)
	return append(out, v...)
}

// argon2Segment fills one lane's slice of blocks for one pass.
func argon2Segment(mem []argon2Block, pass, lane, slice, lanes, laneLen, segLen,
	iterations uint32, variant int) {
	dataIndependent := variant == argon2i ||
		(variant == argon2id && pass == 0 && slice < argon2SyncPoints/2)

	var addresses, input, zero argon2Block
	if dataIndependent {
		input[0] = uint64(pass)
		input[1] = uint64(lane)
		input[2] = uint64(slice)
		input[3] = uint64(len(mem))
		input[4] = uint64(iterations)
		input[5] = uint64(variant)
	}

	start := uint32(0)
	if pass == 0 && slice == 0 {
		start = 2 // the first two blocks of each lane are the initial hashes
		if dataIndependent {
			input[6]++
			argon2Addresses(&addresses, &input, &zero)
		}
	}

	offset := lane*laneLen + slice*segLen + start
	for index := start; index < segLen; index, offset = index+1, offset+1 {
		prev := offset - 1
		if offset%laneLen == 0 {
			prev = offset + laneLen - 1 // wrapped to the end of this lane
		}

		var rand uint64
		if dataIndependent {
			if index%argon2BlockWords == 0 {
				input[6]++
				argon2Addresses(&addresses, &input, &zero)
			}
			rand = addresses[index%argon2BlockWords]
		} else {
			rand = mem[prev][0]
		}

		refLane := uint32(rand>>32) % lanes
		if pass == 0 && slice == 0 {
			refLane = lane
		}
		refIndex := argon2RefIndex(uint32(rand), pass, slice, index, refLane == lane,
			laneLen, segLen)

		argon2Mix(&mem[offset], &mem[prev], &mem[refLane*laneLen+refIndex], pass > 0)
	}
}

// argon2RefIndex picks which earlier block this one mixes with, following
// the RFC's reference-area rules.
func argon2RefIndex(j1 uint32, pass, slice, index uint32, sameLane bool,
	laneLen, segLen uint32) uint32 {
	var refAreaSize uint32
	switch {
	case pass == 0 && slice == 0:
		refAreaSize = index - 1
	case pass == 0 && sameLane:
		refAreaSize = slice*segLen + index - 1
	case pass == 0:
		refAreaSize = slice * segLen
		if index == 0 {
			refAreaSize--
		}
	case sameLane:
		refAreaSize = laneLen - segLen + index - 1
	default:
		refAreaSize = laneLen - segLen
		if index == 0 {
			refAreaSize--
		}
	}
	// The RFC's non-uniform mapping, biased towards recent blocks.
	rel := uint64(j1) * uint64(j1) >> 32
	rel = uint64(refAreaSize) - 1 - (uint64(refAreaSize)*rel)>>32

	start := uint32(0)
	if pass != 0 && slice != argon2SyncPoints-1 {
		start = (slice + 1) * segLen
	}
	return (start + uint32(rel)) % laneLen
}

// argon2Addresses produces one block of data-independent addresses.
func argon2Addresses(addresses, input, zero *argon2Block) {
	var tmp argon2Block
	argon2Mix(&tmp, zero, input, false)
	argon2Mix(addresses, zero, &tmp, false)
}

// argon2Mix is the compression function G: out = Z ^ (X^Y), where Z is the
// BLAKE2b round permutation applied to the rows and then the columns.
func argon2Mix(out, x, y *argon2Block, xorExisting bool) {
	var r argon2Block
	for i := range r {
		r[i] = x[i] ^ y[i]
	}
	var z argon2Block
	copy(z[:], r[:])

	for i := 0; i < 8; i++ { // rows
		row := z[i*16 : i*16+16]
		argon2Permute(row)
	}
	for i := 0; i < 8; i++ { // columns
		var col [16]uint64
		for k := 0; k < 8; k++ {
			col[k*2] = z[k*16+i*2]
			col[k*2+1] = z[k*16+i*2+1]
		}
		argon2Permute(col[:])
		for k := 0; k < 8; k++ {
			z[k*16+i*2] = col[k*2]
			z[k*16+i*2+1] = col[k*2+1]
		}
	}
	for i := range z {
		v := z[i] ^ r[i]
		if xorExisting {
			out[i] ^= v
		} else {
			out[i] = v
		}
	}
}

// argon2Permute is one BLAKE2b round over 16 words, with Argon2's extra
// 32-bit multiplication.
func argon2Permute(v []uint64) {
	ga := func(a, b, c, d int) {
		v[a] = argon2Add(v[a], v[b])
		v[d] = rotr64(v[d]^v[a], 32)
		v[c] = argon2Add(v[c], v[d])
		v[b] = rotr64(v[b]^v[c], 24)
		v[a] = argon2Add(v[a], v[b])
		v[d] = rotr64(v[d]^v[a], 16)
		v[c] = argon2Add(v[c], v[d])
		v[b] = rotr64(v[b]^v[c], 63)
	}
	ga(0, 4, 8, 12)
	ga(1, 5, 9, 13)
	ga(2, 6, 10, 14)
	ga(3, 7, 11, 15)
	ga(0, 5, 10, 15)
	ga(1, 6, 11, 12)
	ga(2, 7, 8, 13)
	ga(3, 4, 9, 14)
}

func argon2Add(a, b uint64) uint64 {
	return a + b + 2*uint64(uint32(a))*uint64(uint32(b))
}

func loadBlock(b *argon2Block, data []byte) {
	for i := range b {
		b[i] = binary.LittleEndian.Uint64(data[i*8:])
	}
}

func storeBlock(b *argon2Block) []byte {
	out := make([]byte, argon2BlockSize)
	for i, v := range b {
		binary.LittleEndian.PutUint64(out[i*8:], v)
	}
	return out
}