package mls

// MLS group types (RFC 9420 §11, §12).
// Data types and serialization from group.go.
// Crypto operations (sign, encrypt, join) go in group_crypto.mx.

// --- Commit ---

type commit struct {
	proposals []proposalOrRef
	path      *updatePath // optional
}

func (c *commit) unmarshal(r *Reader) error {
	*c = commit{}

	err := r.readVector(func(r *Reader) error {
		var por proposalOrRef
		if err := por.unmarshal(r); err != nil {
			return err
		}
		c.proposals = append(c.proposals, por)
		return nil
	})
	if err != nil {
		return err
	}

	present, ok := r.readOptional()
	if !ok {
		return errUnexpectedEOF
	}
	if present {
		c.path = &updatePath{}
		if err := c.path.unmarshal(r); err != nil {
			return err
		}
	}
	return nil
}

func (c *commit) marshal(w *Writer) {
	w.writeVector(len(c.proposals), func(w *Writer, i int) {
		c.proposals[i].marshal(w)
	})
	w.writeOptional(c.path != nil)
	if c.path != nil {
		c.path.marshal(w)
	}
}

// --- GroupInfo ---

type groupInfo struct {
	groupContext    groupContext
	extensions      []extension
	confirmationTag []byte
	signer          leafIndex
	signature       []byte
}

func (info *groupInfo) unmarshal(r *Reader) error {
	*info = groupInfo{}

	if err := info.groupContext.unmarshal(r); err != nil {
		return err
	}

	exts, err := unmarshalExtensionVec(r)
	if err != nil {
		return err
	}
	info.extensions = exts

	var ok bool
	info.confirmationTag, ok = r.readOpaqueVec()
	if !ok {
		return errUnexpectedEOF
	}
	v, ok := r.readUint32()
	if !ok {
		return errUnexpectedEOF
	}
	info.signer = leafIndex(v)
	info.signature, ok = r.readOpaqueVec()
	if !ok {
		return errUnexpectedEOF
	}
	return nil
}

func (info *groupInfo) marshal(w *Writer) {
	info.marshalTBS(w)
	w.writeOpaqueVec(info.signature)
}

func (info *groupInfo) marshalTBS(w *Writer) {
	info.groupContext.marshal(w)
	marshalExtensionVec(w, info.extensions)
	w.writeOpaqueVec(info.confirmationTag)
	w.addUint32(uint32(info.signer))
}

// --- GroupSecrets ---

type groupSecrets struct {
	joinerSecret []byte
	pathSecret   []byte // optional
	psks         []preSharedKeyID
}

func (sec *groupSecrets) unmarshal(r *Reader) error {
	*sec = groupSecrets{}

	var ok bool
	sec.joinerSecret, ok = r.readOpaqueVec()
	if !ok {
		return errUnexpectedEOF
	}

	present, ok := r.readOptional()
	if !ok {
		return errUnexpectedEOF
	}
	if present {
		sec.pathSecret, ok = r.readOpaqueVec()
		if !ok {
			return errUnexpectedEOF
		}
	}

	return r.readVector(func(r *Reader) error {
		var psk preSharedKeyID
		if err := psk.unmarshal(r); err != nil {
			return err
		}
		sec.psks = append(sec.psks, psk)
		return nil
	})
}

func (sec *groupSecrets) marshal(w *Writer) {
	w.writeOpaqueVec(sec.joinerSecret)
	w.writeOptional(sec.pathSecret != nil)
	if sec.pathSecret != nil {
		w.writeOpaqueVec(sec.pathSecret)
	}
	w.writeVector(len(sec.psks), func(w *Writer, i int) {
		sec.psks[i].marshal(w)
	})
}

func (sec *groupSecrets) verifySingleReinitOrBranchPSK() bool {
	n := 0
	for _, pskID := range sec.psks {
		if pskID.pskType != pskTypeResumption {
			continue
		}
		switch pskID.usage {
		case resumptionPSKUsageReinit, resumptionPSKUsageBranch:
			n++
		}
	}
	return n <= 1
}

// --- EncryptedGroupSecrets ---

type encryptedGroupSecrets struct {
	newMember             KeyPackageRef
	encryptedGroupSecrets hpkeCiphertext
}

func (sec *encryptedGroupSecrets) unmarshal(r *Reader) error {
	*sec = encryptedGroupSecrets{}
	var ok bool
	sec.newMember, ok = r.readOpaqueVec()
	if !ok {
		return errUnexpectedEOF
	}
	return sec.encryptedGroupSecrets.unmarshal(r)
}

func (sec *encryptedGroupSecrets) marshal(w *Writer) {
	w.writeOpaqueVec([]byte(sec.newMember))
	sec.encryptedGroupSecrets.marshal(w)
}

// --- Welcome ---

// Welcome includes secret keying information to join a group.
type Welcome struct {
	cipherSuite        CipherSuite
	secrets            []encryptedGroupSecrets
	encryptedGroupInfo []byte
}

// UnmarshalWelcome reads a welcome from an MLS message envelope.
func UnmarshalWelcome(raw []byte) (*Welcome, error) {
	var msg mlsMessage
	if err := unmarshalRaw(raw, &msg); err != nil {
		return nil, err
	}
	if msg.wireFormat != wireFormatMLSWelcome {
		return nil, errInvalidWireFormat
	}
	return msg.welcome, nil
}

// Bytes encodes the welcome in an MLSMessage envelope.
func (w *Welcome) Bytes() []byte {
	raw, err := marshalRaw(&mlsMessage{
		version:    protocolVersionMLS10,
		wireFormat: wireFormatMLSWelcome,
		welcome:    w,
	})
	if err != nil {
		panic("mls: failed to marshal welcome")
	}
	return raw
}

func (w *Welcome) unmarshal(r *Reader) error {
	*w = Welcome{}
	v, ok := r.readUint16()
	if !ok {
		return errUnexpectedEOF
	}
	w.cipherSuite = CipherSuite(v)

	err := r.readVector(func(r *Reader) error {
		var sec encryptedGroupSecrets
		if err := sec.unmarshal(r); err != nil {
			return err
		}
		w.secrets = append(w.secrets, sec)
		return nil
	})
	if err != nil {
		return err
	}

	w.encryptedGroupInfo, ok = r.readOpaqueVec()
	if !ok {
		return errUnexpectedEOF
	}
	return nil
}

func (w *Welcome) marshal(wr *Writer) {
	wr.addUint16(uint16(w.cipherSuite))
	wr.writeVector(len(w.secrets), func(wr *Writer, i int) {
		w.secrets[i].marshal(wr)
	})
	wr.writeOpaqueVec(w.encryptedGroupInfo)
}

// NewMembers returns key package refs this welcome contains secrets for.
func (w *Welcome) NewMembers() []KeyPackageRef {
	refs := []KeyPackageRef{:len(w.secrets)}
	for i, sec := range w.secrets {
		refs[i] = sec.newMember
	}
	return refs
}

func (w *Welcome) findSecret(ref KeyPackageRef) *encryptedGroupSecrets {
	for i, sec := range w.secrets {
		if sec.newMember.equal(ref) {
			return &w.secrets[i]
		}
	}
	return nil
}