// Copyright (c) 2022 Tulir Asokan
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

package whatsmeow

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"fmt"
	"time"

	"google.golang.org/protobuf/proto"

	waBinary "go.mau.fi/whatsmeow/binary"
	waProto "go.mau.fi/whatsmeow/binary/proto"
	"go.mau.fi/whatsmeow/types"
	"go.mau.fi/whatsmeow/types/events"
	"go.mau.fi/whatsmeow/util/hkdfutil"
)

func getMediaRetryKey(mediaKey []byte) (cipherKey []byte) {
	return hkdfutil.SHA256(mediaKey, nil, []byte("WhatsApp Media Retry Notification"), 32)
}

func prepareMediaRetryGCM(mediaKey []byte) (cipher.AEAD, error) {
	block, err := aes.NewCipher(getMediaRetryKey(mediaKey))
	if err != nil {
		return nil, fmt.Errorf("failed to initialize AES cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize GCM: %w", err)
	}
	return gcm, nil
}

func encryptMediaRetryReceipt(messageID types.MessageID, mediaKey []byte) (ciphertext, iv []byte, err error) {
	receipt := &waProto.ServerErrorReceipt{
		StanzaId: proto.String(messageID),
	}
	var plaintext []byte
	plaintext, err = proto.Marshal(receipt)
	if err != nil {
		err = fmt.Errorf("failed to marshal payload: %w", err)
		return
	}
	var gcm cipher.AEAD
	gcm, err = prepareMediaRetryGCM(mediaKey)
	if err != nil {
		return
	}
	iv = make([]byte, 12)
	_, err = rand.Read(iv)
	if err != nil {
		panic(err)
	}
	ciphertext = gcm.Seal(plaintext[:0], iv, plaintext, []byte(messageID))
	return
}

// SendMediaRetryReceipt sends a request to the phone to re-upload the media in a message.
//
// The response will come as an *events.MediaRetry. The response will then have to be decrypted
// using DecryptMediaRetryNotification and the same media key passed here.
func (cli *Client) SendMediaRetryReceipt(message *types.MessageInfo, mediaKey []byte) error {
	ciphertext, iv, err := encryptMediaRetryReceipt(message.ID, mediaKey)
	if err != nil {
		return fmt.Errorf("failed to prepare encrypted retry receipt: %w", err)
	}

	rmrAttrs := waBinary.Attrs{
		"jid":     message.Chat,
		"from_me": message.IsFromMe,
	}
	if message.IsGroup {
		rmrAttrs["participant"] = message.Sender
	}

	encryptedRequest := []waBinary.Node{
		{Tag: "enc_p", Content: ciphertext},
		{Tag: "enc_iv", Content: iv},
	}

	err = cli.sendNode(waBinary.Node{
		Tag: "receipt",
		Attrs: waBinary.Attrs{
			"id":   message.ID,
			"to":   cli.Store.ID.ToNonAD(),
			"type": "server-error",
		},
		Content: []waBinary.Node{
			{Tag: "encrypt", Content: encryptedRequest},
			{Tag: "rmr", Attrs: rmrAttrs},
		},
	})
	if err != nil {
		return err
	}
	return nil
}

// DecryptMediaRetryNotification decrypts a media retry notification using the media key.
func DecryptMediaRetryNotification(evt *events.MediaRetry, mediaKey []byte) (*waProto.MediaRetryNotification, error) {
	var notif waProto.MediaRetryNotification
	var plaintext []byte
	if evt.Error != nil && evt.Ciphertext == nil {
		if evt.Error.Code == 2 {
			return nil, ErrMediaNotAvailableOnPhone
		}
		return nil, fmt.Errorf("%w (code: %d)", ErrUnknownMediaRetryError, evt.Error.Code)
	} else if gcm, err := prepareMediaRetryGCM(mediaKey); err != nil {
		return nil, err
	} else if plaintext, err = gcm.Open(nil, evt.IV, evt.Ciphertext, []byte(evt.MessageID)); err != nil {
		return nil, fmt.Errorf("failed to decrypt notification: %w", err)
	} else if err = proto.Unmarshal(plaintext, &notif); err != nil {
		return nil, fmt.Errorf("failed to unmarshal notification (invalid encryption key?): %w", err)
	} else {
		return &notif, nil
	}
}

func parseMediaRetryNotification(node *waBinary.Node) (*events.MediaRetry, error) {
	ag := node.AttrGetter()
	var evt events.MediaRetry
	evt.Timestamp = time.Unix(ag.Int64("t"), 0)
	evt.MessageID = types.MessageID(ag.String("id"))
	if !ag.OK() {
		return nil, ag.Error()
	}
	rmr, ok := node.GetOptionalChildByTag("rmr")
	if !ok {
		return nil, &ElementMissingError{Tag: "rmr", In: "retry notification"}
	}
	rmrAG := rmr.AttrGetter()
	evt.ChatID = rmrAG.JID("jid")
	evt.FromMe = rmrAG.Bool("from_me")
	evt.SenderID = rmrAG.OptionalJIDOrEmpty("participant")
	if !rmrAG.OK() {
		return nil, fmt.Errorf("missing attributes in <rmr> tag: %w", rmrAG.Error())
	}

	errNode, ok := node.GetOptionalChildByTag("error")
	if ok {
		evt.Error = &events.MediaRetryError{
			Code: errNode.AttrGetter().Int("code"),
		}
		return &evt, nil
	}

	evt.Ciphertext, ok = node.GetChildByTag("encrypt", "enc_p").Content.([]byte)
	if !ok {
		return nil, &ElementMissingError{Tag: "enc_p", In: fmt.Sprintf("retry notification %s", evt.MessageID)}
	}
	evt.IV, ok = node.GetChildByTag("encrypt", "enc_iv").Content.([]byte)
	if !ok {
		return nil, &ElementMissingError{Tag: "enc_iv", In: fmt.Sprintf("retry notification %s", evt.MessageID)}
	}
	return &evt, nil
}

func (cli *Client) handleMediaRetryNotification(node *waBinary.Node) {
	evt, err := parseMediaRetryNotification(node)
	if err != nil {
		cli.Log.Warnf("Failed to parse media retry notification: %v", err)
		return
	}
	cli.dispatchEvent(evt)
}