package testkit

import (
	"bytes"
	"crypto/ecdh"
	"crypto/sha256"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"

	"filippo.io/age"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/hkdf"

	"g.activething.com/go/DateKeys/accesskey"
	"g.activething.com/go/DateKeys/agewrap"
	"g.activething.com/go/DateKeys/capsule"
	"g.activething.com/go/DateKeys/internal/cbortest"
)

// This file rebuilds age files deterministically, with a file key and a nonce
// that are already known: the STREAM of the age payload (C2SP age.md) and an
// X25519 stanza with a chosen ephemeral key. With them, the mutations of the
// format 2 fixtures edit a control, a stanza list or a plaintext and seal it
// again with the file key of the fixture, so that the result is the same on
// every run and keeps every MAC valid, as the holder of that file key could
// make it.

const (
	streamChunkSize = 64 << 10
	streamNonceSize = 16
)

// StreamLen returns the length of the STREAM of an age payload of n bytes:
// the plaintext and a 16-byte tag for each 64 KiB chunk, at least one.
func StreamLen(n int) int {
	return n + chacha20poly1305.Overhead*max(1, (n+streamChunkSize-1)/streamChunkSize)
}

// StreamSeal encrypts plaintext as the payload of an age file whose file key
// and 16-byte nonce are given: ChaCha20-Poly1305 over 64 KiB chunks, with the
// key HKDF-SHA-256(file key, nonce, "payload") and a nonce of an 11-byte
// big-endian counter and a last-chunk flag.
func StreamSeal(fileKey, nonce, plaintext []byte) ([]byte, error) {
	if len(nonce) != streamNonceSize {
		return nil, fmt.Errorf("testkit: STREAM nonce of %d bytes", len(nonce))
	}
	key := make([]byte, chacha20poly1305.KeySize)
	if _, err := io.ReadFull(hkdf.New(sha256.New, fileKey, nonce, []byte("payload")), key); err != nil {
		return nil, err
	}
	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, err
	}
	out := make([]byte, 0, StreamLen(len(plaintext)))
	var counter [chacha20poly1305.NonceSize]byte
	for i := uint64(0); ; i++ {
		chunk := plaintext[:min(len(plaintext), streamChunkSize)]
		plaintext = plaintext[len(chunk):]
		// The counter fills bytes 0 to 10; a payload never needs 2^64 chunks.
		binary.BigEndian.PutUint64(counter[3:11], i)
		if len(plaintext) == 0 {
			counter[11] = 1
			return aead.Seal(out, counter[:], chunk, nil), nil
		}
		out = aead.Seal(out, counter[:], chunk, nil)
	}
}

// AgeParts splits an age file into its header, its nonce and its STREAM.
func AgeParts(file []byte) (header, nonce, stream []byte, err error) {
	n, err := HeaderLen(file)
	if err != nil {
		return nil, nil, nil, err
	}
	if len(file) < n+streamNonceSize {
		return nil, nil, nil, errors.New("testkit: age file without a nonce")
	}
	return file[:n], file[n : n+streamNonceSize], file[n+streamNonceSize:], nil
}

// ResealAge returns file with its payload replaced by plaintext, sealed with
// fileKey and the nonce of file. The header, and with it its MAC, is kept.
func ResealAge(file, fileKey, plaintext []byte) ([]byte, error) {
	header, nonce, _, err := AgeParts(file)
	if err != nil {
		return nil, err
	}
	stream, err := StreamSeal(fileKey, nonce, plaintext)
	if err != nil {
		return nil, err
	}
	return Join(header, nonce, stream), nil
}

// SealAge returns an age file with the given stanzas, a header MAC computed
// with fileKey, the given nonce and plaintext as its payload.
func SealAge(stanzas []*age.Stanza, fileKey, nonce, plaintext []byte) ([]byte, error) {
	header, err := MarshalHeader(stanzas, fileKey)
	if err != nil {
		return nil, err
	}
	stream, err := StreamSeal(fileKey, nonce, plaintext)
	if err != nil {
		return nil, err
	}
	return Join(header, nonce, stream), nil
}

// OpenAgeWithKey decrypts an age file whose file key is known, verifying its
// header MAC and its STREAM.
func OpenAgeWithKey(file, fileKey []byte) ([]byte, error) {
	r, err := age.Decrypt(bytes.NewReader(file), age.NewInjectedFileKeyIdentity(fileKey))
	if err != nil {
		return nil, err
	}
	return io.ReadAll(r)
}

// FixedX25519Stanza returns the X25519 stanza that wraps fileKey for
// recipient with the ephemeral scalar SHA-256(seed), as age builds it with a
// random one: the share X25519(e, base point), and the file key sealed with
// the key HKDF-SHA-256(X25519(e, recipient), share || recipient,
// "age-encryption.org/v1/X25519") and a zero nonce.
func FixedX25519Stanza(fileKey []byte, recipient *age.X25519Recipient, seed string) (*age.Stanza, error) {
	their, err := agewrap.RawX25519Recipient(recipient)
	if err != nil {
		return nil, err
	}
	scalar := sha256.Sum256([]byte(seed))
	e, err := ecdh.X25519().NewPrivateKey(scalar[:])
	if err != nil {
		return nil, err
	}
	pub, err := ecdh.X25519().NewPublicKey(their)
	if err != nil {
		return nil, err
	}
	shared, err := e.ECDH(pub)
	if err != nil {
		return nil, err
	}
	share := e.PublicKey().Bytes()
	key := make([]byte, chacha20poly1305.KeySize)
	if _, err := io.ReadFull(hkdf.New(sha256.New, shared, Join(share, their), []byte("age-encryption.org/v1/X25519")), key); err != nil {
		return nil, err
	}
	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, err
	}
	body := aead.Seal(nil, make([]byte, chacha20poly1305.NonceSize), fileKey, nil)
	return &age.Stanza{Type: agewrap.StanzaX25519, Args: []string{base64.RawStdEncoding.EncodeToString(share)}, Body: body}, nil
}

// withSealedLen returns a copy of prelude with SEALED_CONTROL_LEN set to n.
func withSealedLen(prelude []byte, n int) []byte {
	p := bytes.Clone(prelude[:capsule.PreludeSize])
	binary.BigEndian.PutUint32(p[12:16], uint32(n))
	return p
}

// bind returns control, a CONTROL_CBOR map, with its header_binding (key 2)
// computed for prelude and header, and encoded; without key 2 it is only
// encoded.
func bind(control map[uint64]any, prelude, header []byte) ([]byte, error) {
	if _, ok := control[2]; ok {
		var p [capsule.PreludeSize]byte
		copy(p[:], prelude)
		b := capsule.HeaderBinding(p, header)
		control[2] = b[:]
	}
	return cbortest.Marshal(control)
}

// AccessIdentity returns the identity of the .dkk of a time_and_key fixture.
func (f *LoadedFixture) AccessIdentity() (*age.X25519Identity, error) {
	if f.DKK == nil {
		return nil, errors.New("testkit: the fixture has no .dkk")
	}
	k, err := accesskey.Decode(bytes.NewReader(f.DKK))
	if err != nil {
		return nil, err
	}
	defer k.Wipe()
	return agewrap.X25519IdentityFromRaw(k.Material)
}

// withIdentity makes in offer the identity of the .dkk of f instead of the
// .dkk itself, whose capsule_digest a mutation of the .dkc breaks.
func (f *LoadedFixture) withIdentity(in *MutationInput) (*MutationInput, error) {
	id, err := f.AccessIdentity()
	if err != nil {
		return nil, err
	}
	in.DKK, in.Identities = nil, []string{id.String()}
	return in, nil
}

// WithControl returns f, a time_only fixture, with its CONTROL_CBOR decoded
// as a map, changed by edit and sealed again with FK_TIME and the nonce of
// OUTER_TIME_AGE, as anyone can once the round is published (spec §36.1).
// When the length of SEALED_CONTROL changes, the PRELUDE says so, and the
// header_binding of the control, if edit keeps one, covers the new PRELUDE.
func (f *LoadedFixture) WithControl(edit func(control map[uint64]any)) (*MutationInput, error) {
	if f.AccessPolicy != capsule.TimeOnly.String() {
		return nil, errors.New("testkit: WithControl needs a time_only fixture")
	}
	fk, err := f.TimeFileKey()
	if err != nil {
		return nil, err
	}
	control, err := OpenAgeWithKey(f.Parts.Sealed, fk)
	if err != nil {
		return nil, err
	}
	m, err := cbortest.UnmarshalMap(control)
	if err != nil {
		return nil, err
	}
	edit(m)
	draft, err := cbortest.Marshal(m)
	if err != nil {
		return nil, err
	}
	header, _, _, err := AgeParts(f.Parts.Sealed)
	if err != nil {
		return nil, err
	}
	prelude := withSealedLen(f.Parts.Prelude, len(header)+streamNonceSize+StreamLen(len(draft)))
	control, err = bind(m, prelude, f.Parts.Header)
	if err != nil {
		return nil, err
	}
	sealed, err := ResealAge(f.Parts.Sealed, fk, control)
	if err != nil {
		return nil, err
	}
	return f.input(Join(prelude, f.Parts.Header, sealed, f.Parts.Payload)), nil
}

// WithPayloadPlaintext returns f with the plaintext of its PAYLOAD_AGE
// replaced by plaintext, sealed with FK_PAYLOAD and the nonce of PAYLOAD_AGE,
// as whoever knows I_PAYLOAD can make it.
func (f *LoadedFixture) WithPayloadPlaintext(plaintext []byte) (*MutationInput, error) {
	payload, err := f.resealPayload(plaintext)
	if err != nil {
		return nil, err
	}
	return f.input(Join(f.Parts.Prelude, f.Parts.Header, f.Parts.Sealed, payload)), nil
}

// WithInnerStanzas returns f, a time_and_key fixture with a .dkk, with the
// stanzas of INNER_ACCESS_AGE replaced by edit(FK_ACCESS, stanzas), as its
// creator could make them: the header MAC is recomputed with FK_ACCESS, the
// PRELUDE follows the new length of SEALED_CONTROL, the control is bound to
// the new PRELUDE, and INNER_ACCESS_AGE and OUTER_TIME_AGE are sealed again
// with FK_ACCESS, FK_TIME and their nonces. The input offers the identity of
// the .dkk, whose capsule_digest no longer matches.
func (f *LoadedFixture) WithInnerStanzas(edit func(fileKey []byte, stanzas []*age.Stanza) ([]*age.Stanza, error)) (*MutationInput, error) {
	fkTime, err := f.TimeFileKey()
	if err != nil {
		return nil, err
	}
	inner, err := OpenAgeWithKey(f.Parts.Sealed, fkTime)
	if err != nil {
		return nil, err
	}
	id, err := f.AccessIdentity()
	if err != nil {
		return nil, err
	}
	stanzas, err := agewrap.Stanzas(bytes.NewReader(inner))
	if err != nil {
		return nil, err
	}
	var fkAccess []byte
	for _, s := range stanzas {
		if fk, err := id.Unwrap([]*age.Stanza{s}); err == nil {
			fkAccess = fk
		}
	}
	if fkAccess == nil {
		return nil, errors.New("testkit: the .dkk opens no stanza of INNER_ACCESS_AGE")
	}
	control, err := OpenAgeWithKey(inner, fkAccess)
	if err != nil {
		return nil, err
	}
	if stanzas, err = edit(fkAccess, stanzas); err != nil {
		return nil, err
	}
	innerHeader, err := MarshalHeader(stanzas, fkAccess)
	if err != nil {
		return nil, err
	}
	_, innerNonce, _, err := AgeParts(inner)
	if err != nil {
		return nil, err
	}
	outerHeader, _, _, err := AgeParts(f.Parts.Sealed)
	if err != nil {
		return nil, err
	}
	innerLen := len(innerHeader) + streamNonceSize + StreamLen(len(control))
	prelude := withSealedLen(f.Parts.Prelude, len(outerHeader)+streamNonceSize+StreamLen(innerLen))
	m, err := cbortest.UnmarshalMap(control)
	if err != nil {
		return nil, err
	}
	if control, err = bind(m, prelude, f.Parts.Header); err != nil {
		return nil, err
	}
	stream, err := StreamSeal(fkAccess, innerNonce, control)
	if err != nil {
		return nil, err
	}
	sealed, err := ResealAge(f.Parts.Sealed, fkTime, Join(innerHeader, innerNonce, stream))
	if err != nil {
		return nil, err
	}
	return f.withIdentity(f.input(Join(prelude, f.Parts.Header, sealed, f.Parts.Payload)))
}