package gateway

import (
	"encoding/base64"
	"encoding/xml"

	"dev.narayana.im/narayana/telegabber/e2ee/omemo"
	"dev.narayana.im/narayana/telegabber/xmpp/extensions"

	"gosrc.io/xmpp/stanza"
)

// omemoFallbackBody is the plaintext shown by clients that don't support
// OMEMO (or don't yet trust this conversation's OMEMO session) instead of
// an empty message body - a courtesy hint, not part of the wire protocol
// (XEP-0380's EME, attached alongside this by sendMessageWrapper, is the
// actual machine-readable hint).
const omemoFallbackBody = "[This message is OMEMO end-to-end encrypted]"

// OMEMOSendFailedBody is sent in place of the real message body when OMEMO
// is active for a chat but encryption couldn't go through (no session
// could be established, e.g. the peer never answered a PEP bundle
// request) - callers (telegram/utils.go, telegram/handlers.go) must never
// fall back to sending the real plaintext in this case, only this notice.
const OMEMOSendFailedBody = "[Failed to send: could not establish an OMEMO session]"

// omemoStanzaExtension converts a just-produced OMEMO e2ee.Envelope's raw
// bytes (an omemo.WireEnvelope, gob-encoded - see e2ee/omemo/envelope.go)
// into the wire stanza.MsgExtension for its version, plus the XEP-0380 EME
// hint. This is the backend-specific glue e2ee.Envelope's own doc comment
// defers to the xmpp layer, since stanza XML shape isn't part of the
// backend-agnostic e2ee.Backend interface - a future second backend would
// need its own equivalent, switched on e2ee.Envelope.Backend (see its only
// caller, sendMessageWrapper's SMOMEMOEnvelope handling).
func omemoStanzaExtension(raw []byte) (extension stanza.MsgExtension, eme extensions.EME, err error) {
	wireEnv, err := omemo.DecodeWireEnvelope(raw)
	if err != nil {
		return nil, extensions.EME{}, err
	}

	if wireEnv.Version == omemo.Omemo0 {
		header := extensions.OMEMO0Header{
			SID: wireEnv.SenderSID,
			IV:  base64.StdEncoding.EncodeToString(wireEnv.IV),
		}
		for _, k := range wireEnv.Keys {
			header.Keys = append(header.Keys, extensions.OMEMO0Key{
				RID:    k.DeviceID,
				PreKey: k.IsPreKey,
				Text:   base64.StdEncoding.EncodeToString(k.Ciphertext),
			})
		}
		enc := extensions.OMEMO0Encrypted{
			Header:  header,
			Payload: base64.StdEncoding.EncodeToString(wireEnv.Payload),
		}
		return enc, extensions.EME{EncryptionNamespace: "eu.siacs.conversations.axolotl", Name: "OMEMO"}, nil
	}

	ns := "urn:xmpp:omemo:2"
	if wireEnv.Version == omemo.Omemo1 {
		ns = "urn:xmpp:omemo:1"
	}

	// Keys are grouped per-recipient-JID on the wire (see OMEMOEncrypted's
	// doc comment) - wireEnv.Keys is a flat list, so re-group here,
	// preserving first-seen JID order for deterministic output.
	groups := make(map[string]*extensions.OMEMOKeys)
	var order []string
	for _, k := range wireEnv.Keys {
		group, ok := groups[k.RecipientJID]
		if !ok {
			group = &extensions.OMEMOKeys{JID: k.RecipientJID}
			groups[k.RecipientJID] = group
			order = append(order, k.RecipientJID)
		}
		group.Keys = append(group.Keys, extensions.OMEMOKey{
			RID:  k.DeviceID,
			Kex:  k.IsPreKey,
			Text: base64.StdEncoding.EncodeToString(k.Ciphertext),
		})
	}
	header := extensions.OMEMOHeader{SID: wireEnv.SenderSID}
	for _, jid := range order {
		header.Keys = append(header.Keys, *groups[jid])
	}

	enc := extensions.OMEMOEncrypted{
		XMLName: xml.Name{Space: ns, Local: "encrypted"},
		Header:  header,
		Payload: base64.StdEncoding.EncodeToString(wireEnv.Payload),
	}
	return enc, extensions.EME{EncryptionNamespace: ns, Name: "OMEMO"}, nil
}