package testkit

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"strings"

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

	"github.com/datekeys/datekeys-go/agewrap"
)

// CaptureRecipient forwards to Recipient and records the file key that age
// asks it to wrap. Knowing the file key lets a test rewrite the age header
// and recompute a valid MAC, which models a malicious creator (spec §27).
type CaptureRecipient struct {
	Recipient age.Recipient
	FileKey   []byte
}

// Wrap implements age.Recipient.
func (c *CaptureRecipient) Wrap(fileKey []byte) ([]*age.Stanza, error) {
	c.FileKey = bytes.Clone(fileKey)
	return c.Recipient.Wrap(fileKey)
}

// WrapWithLabels forwards labels when the wrapped recipient has them.
func (c *CaptureRecipient) WrapWithLabels(fileKey []byte) ([]*age.Stanza, []string, error) {
	c.FileKey = bytes.Clone(fileKey)
	if l, ok := c.Recipient.(age.RecipientWithLabels); ok {
		return l.WrapWithLabels(fileKey)
	}
	s, err := c.Recipient.Wrap(fileKey)
	return s, nil, err
}

// Encrypt returns a complete age file for plaintext and the file key age
// generated for it.
func Encrypt(plaintext []byte, recipients ...age.Recipient) (file, fileKey []byte, err error) {
	if len(recipients) == 0 {
		return nil, nil, errors.New("testkit: no recipients")
	}
	capture := &CaptureRecipient{Recipient: recipients[0]}
	all := append([]age.Recipient{capture}, recipients[1:]...)
	var buf bytes.Buffer
	w, err := age.Encrypt(&buf, all...)
	if err != nil {
		return nil, nil, err
	}
	if _, err := w.Write(plaintext); err != nil {
		return nil, nil, err
	}
	if err := w.Close(); err != nil {
		return nil, nil, err
	}
	return buf.Bytes(), capture.FileKey, nil
}

// HeaderLen returns the length of the age header at the start of file.
func HeaderLen(file []byte) (int, error) {
	hdr, err := age.ExtractHeader(bytes.NewReader(file))
	if err != nil {
		return 0, err
	}
	if !bytes.HasPrefix(file, hdr) {
		return 0, errors.New("testkit: header is not in canonical form")
	}
	return len(hdr), nil
}

// RewriteAge replaces the recipient stanzas of an age file with
// edit(stanzas) and recomputes the header MAC with fileKey, keeping the nonce
// and the STREAM payload. The result is an age file that age itself accepts
// whenever some identity yields fileKey: only structural checks can reject it.
func RewriteAge(file, fileKey []byte, edit func([]*age.Stanza) []*age.Stanza) ([]byte, error) {
	stanzas, err := agewrap.Stanzas(bytes.NewReader(file))
	if err != nil {
		return nil, err
	}
	n, err := HeaderLen(file)
	if err != nil {
		return nil, err
	}
	hdr, err := MarshalHeader(edit(stanzas), fileKey)
	if err != nil {
		return nil, err
	}
	return append(hdr, file[n:]...), nil
}

// MarshalHeader serialises an age v1 header as specified by C2SP age.md: the
// intro line, each stanza, and the footer with
// HMAC-SHA-256(HKDF-SHA-256(file key, "", "header"), header up to "---").
func MarshalHeader(stanzas []*age.Stanza, fileKey []byte) ([]byte, error) {
	var b bytes.Buffer
	b.WriteString("age-encryption.org/v1\n")
	for _, s := range stanzas {
		if !validArg(s.Type) {
			return nil, fmt.Errorf("testkit: invalid stanza type %q", s.Type)
		}
		b.WriteString("-> " + s.Type)
		for _, a := range s.Args {
			if !validArg(a) {
				return nil, fmt.Errorf("testkit: invalid stanza argument %q", a)
			}
			b.WriteString(" " + a)
		}
		b.WriteString("\n")
		body := base64.RawStdEncoding.EncodeToString(s.Body)
		for len(body) >= 64 {
			b.WriteString(body[:64] + "\n")
			body = body[64:]
		}
		b.WriteString(body + "\n") // the final line is always short, possibly empty
	}
	b.WriteString("---")
	key := make([]byte, 32)
	if _, err := io.ReadFull(hkdf.New(sha256.New, fileKey, nil, []byte("header")), key); err != nil {
		return nil, err
	}
	mac := hmac.New(sha256.New, key)
	mac.Write(b.Bytes())
	b.WriteString(" " + base64.RawStdEncoding.EncodeToString(mac.Sum(nil)) + "\n")
	return b.Bytes(), nil
}

func validArg(s string) bool {
	return s != "" && !strings.ContainsFunc(s, func(r rune) bool { return r < 33 || r > 126 })
}

// X25519Stanza returns a well-formed X25519 stanza wrapping fileKey for a
// fresh identity, and that identity. It models an extra decryption path that
// a malicious creator could add.
func X25519Stanza(fileKey []byte) (*age.Stanza, *age.X25519Identity, error) {
	id, err := age.GenerateX25519Identity()
	if err != nil {
		return nil, nil, err
	}
	s, err := id.Recipient().Wrap(fileKey)
	if err != nil {
		return nil, nil, err
	}
	return s[0], id, nil
}