package testkit

import (
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"

	"filippo.io/age"

	"g.activething.com/go/DateKeys/agewrap"
	"g.activething.com/go/DateKeys/capsule"
	"g.activething.com/go/DateKeys/codec"
	"g.activething.com/go/DateKeys/datekey"
	"g.activething.com/go/DateKeys/extension"
	"g.activething.com/go/DateKeys/profile"
)

// Build assembles a capsule step by step like capsule.Encrypt, but lets a test
// declare one policy and build another structure, and edit the stanzas of each
// age file while keeping every MAC valid. Every other value (header_binding,
// lengths) stays consistent, as a malicious creator would make it.
type Build struct {
	Profile   *profile.Profile // default Quicknet
	Round     uint64           // default 1000
	Declared  capsule.Policy   // access_policy written in PUBLIC_HEADER
	Structure capsule.Policy   // construction actually used
	// AccessRecipients of INNER_ACCESS_AGE when Structure is time_and_key.
	AccessRecipients []age.Recipient
	Plaintext        []byte
	DateKeyString    string // overrides the canonical dk1_ string in PUBLIC_HEADER
	// RawPolicy, when not zero, is the access_policy value written in
	// PUBLIC_HEADER instead of Declared, which may be outside V1. A header
	// written with RawPolicy or DateKeyString carries no extensions.
	RawPolicy uint64

	HeaderCritical, HeaderNoncritical   []extension.Extension
	ControlCritical, ControlNoncritical []extension.Extension

	// Stanza edits, applied with the corresponding file key.
	EditOuter   func(fileKey []byte, s []*age.Stanza) []*age.Stanza
	EditInner   func(fileKey []byte, s []*age.Stanza) []*age.Stanza
	EditPayload func(fileKey []byte, s []*age.Stanza) []*age.Stanza
}

// Built is a capsule produced by Build and its secrets.
type Built struct {
	DKC             []byte
	Prelude         [capsule.PreludeSize]byte
	PublicHeader    []byte
	Sealed          []byte
	Payload         []byte
	Control         []byte
	PayloadIdentity []byte
	CapsuleID       [capsule.CapsuleIDSize]byte
}

// Make builds the capsule.
func (b Build) Make() (*Built, error) {
	p := b.Profile
	if p == nil {
		p = profile.Quicknet()
	}
	round := b.Round
	if round == 0 {
		round = 1000
	}
	out := &Built{}
	if _, err := rand.Read(out.CapsuleID[:]); err != nil {
		return nil, err
	}
	header, err := b.header(p, round, out.CapsuleID)
	if err != nil {
		return nil, err
	}
	out.PublicHeader = header

	payloadID, err := age.GenerateX25519Identity()
	if err != nil {
		return nil, err
	}
	if out.PayloadIdentity, err = agewrap.RawX25519Identity(payloadID); err != nil {
		return nil, err
	}
	payload, fk, err := Encrypt(b.Plaintext, payloadID.Recipient())
	if err != nil {
		return nil, err
	}
	if b.EditPayload != nil {
		if payload, err = RewriteAge(payload, fk, func(s []*age.Stanza) []*age.Stanza { return b.EditPayload(fk, s) }); err != nil {
			return nil, err
		}
	}
	out.Payload = payload

	timeRecipient, err := agewrap.NewTimeRecipient(p, round)
	if err != nil {
		return nil, err
	}
	seal := func(control []byte) ([]byte, error) {
		plaintext := control
		if b.Structure == capsule.TimeAndKey {
			if len(b.AccessRecipients) == 0 {
				return nil, errors.New("testkit: time_and_key structure needs AccessRecipients")
			}
			inner, fk, err := Encrypt(control, b.AccessRecipients...)
			if err != nil {
				return nil, err
			}
			if b.EditInner != nil {
				if inner, err = RewriteAge(inner, fk, func(s []*age.Stanza) []*age.Stanza { return b.EditInner(fk, s) }); err != nil {
					return nil, err
				}
			}
			plaintext = inner
		}
		outer, fk, err := Encrypt(plaintext, timeRecipient)
		if err != nil {
			return nil, err
		}
		if b.EditOuter != nil {
			return RewriteAge(outer, fk, func(s []*age.Stanza) []*age.Stanza { return b.EditOuter(fk, s) })
		}
		return outer, nil
	}

	ctrl := &capsule.Control{Critical: b.ControlCritical, Noncritical: b.ControlNoncritical}
	draft, err := capsule.EncodeControl(ctrl)
	if err != nil {
		return nil, err
	}
	draftSealed, err := seal(draft)
	if err != nil {
		return nil, err
	}
	out.Prelude = capsule.Prelude{PublicHeaderLen: uint32(len(header)), SealedControlLen: uint32(len(draftSealed))}.Bytes()
	ctrl.HeaderBinding = capsule.HeaderBinding(out.Prelude, header)
	copy(ctrl.PayloadIdentity[:], out.PayloadIdentity)
	if out.Control, err = capsule.EncodeControl(ctrl); err != nil {
		return nil, err
	}
	if out.Sealed, err = seal(out.Control); err != nil {
		return nil, err
	}
	if len(out.Sealed) != len(draftSealed) {
		return nil, fmt.Errorf("testkit: SEALED_CONTROL length changed from %d to %d", len(draftSealed), len(out.Sealed))
	}
	out.DKC = Join(out.Prelude[:], out.PublicHeader, out.Sealed, out.Payload)
	return out, nil
}

func (b Build) header(p *profile.Profile, round uint64, id [capsule.CapsuleIDSize]byte) ([]byte, error) {
	h := &capsule.Header{
		CapsuleID:   id,
		DateKey:     datekey.DateKey{ProfileID: p.ID, Round: round},
		Policy:      b.Declared,
		Critical:    b.HeaderCritical,
		Noncritical: b.HeaderNoncritical,
	}
	if b.DateKeyString == "" && b.RawPolicy == 0 {
		return capsule.EncodeHeader(h)
	}
	dk, policy := b.DateKeyString, uint64(b.Declared)
	if dk == "" {
		dk = h.DateKey.Compact()
	}
	if b.RawPolicy != 0 {
		policy = b.RawPolicy
	}
	return RawHeader(id, dk, policy)
}

// RawHeader encodes a PUBLIC_HEADER with an arbitrary DateKey string, which
// must be valid UTF-8, and policy value, bypassing the validation of
// capsule.EncodeHeader.
func RawHeader(id [capsule.CapsuleIDSize]byte, dk string, policy uint64) ([]byte, error) {
	var e codec.Encoder
	e.Map(5)
	e.Uint(0)
	e.Text(capsule.HeaderTypeTag)
	e.Uint(1)
	e.Uint(capsule.HeaderVersion)
	e.Uint(2)
	e.Bstr(id[:])
	e.Uint(3)
	e.Text(dk)
	e.Uint(4)
	e.Uint(policy)
	return e.Out()
}

// Parts is a .dkc split into its four sections.
type Parts struct {
	Prelude, Header, Sealed, Payload []byte
}

// Split splits a .dkc using the lengths in its prelude.
func Split(dkc []byte) (Parts, error) {
	if len(dkc) < capsule.PreludeSize {
		return Parts{}, errors.New("testkit: short capsule")
	}
	hl := int(binary.BigEndian.Uint32(dkc[8:12]))
	sl := int(binary.BigEndian.Uint32(dkc[12:16]))
	if len(dkc) < capsule.PreludeSize+hl+sl {
		return Parts{}, errors.New("testkit: capsule shorter than its prelude says")
	}
	h := dkc[capsule.PreludeSize : capsule.PreludeSize+hl]
	s := dkc[capsule.PreludeSize+hl : capsule.PreludeSize+hl+sl]
	return Parts{Prelude: dkc[:capsule.PreludeSize], Header: h, Sealed: s, Payload: dkc[capsule.PreludeSize+hl+sl:]}, nil
}

// Join concatenates sections into a new slice.
func Join(parts ...[]byte) []byte {
	var out []byte
	for _, p := range parts {
		out = append(out, p...)
	}
	return out
}

// Reframe joins sections with a prelude whose lengths match them, keeping
// the version, flags and reserved bytes of prelude.
func Reframe(prelude, header, sealed, payload []byte) []byte {
	pre := append([]byte(nil), prelude[:capsule.PreludeSize]...)
	binary.BigEndian.PutUint32(pre[8:12], uint32(len(header)))
	binary.BigEndian.PutUint32(pre[12:16], uint32(len(sealed)))
	return Join(pre, header, sealed, payload)
}