package extensions

import (
	"encoding/xml"

	"gosrc.io/xmpp/stanza"
)

// OMEMO0Encrypted is the legacy/"siacs" OMEMO <encrypted> element
// (eu.siacs.conversations.axolotl, XEP-0384 <= 0.3.0): a flat list of
// per-device <key> elements directly under <header>, plus an explicit
// <iv> - the payload's independently-random GCM nonce. Omemo1/Omemo2 (see
// OMEMOEncrypted) have no such element, since their IV is HKDF-derived
// from the key instead.
//
//	<encrypted xmlns='eu.siacs.conversations.axolotl'>
//	  <header sid='27183'>
//	    <key rid='31415' prekey='true'>BASE64</key>
//	    <iv>BASE64</iv>
//	  </header>
//	  <payload>BASE64</payload>
//	</encrypted>
type OMEMO0Encrypted struct {
	XMLName xml.Name     `xml:"eu.siacs.conversations.axolotl encrypted"`
	Header  OMEMO0Header `xml:"header"`
	Payload string       `xml:"payload,omitempty"`
}

// Namespace implements stanza.MsgExtension.
func (c OMEMO0Encrypted) Namespace() string {
	return c.XMLName.Space
}

type OMEMO0Header struct {
	SID  uint32      `xml:"sid,attr"`
	Keys []OMEMO0Key `xml:"key"`
	IV   string      `xml:"iv"`
}

type OMEMO0Key struct {
	RID    uint32 `xml:"rid,attr"`
	PreKey bool   `xml:"prekey,attr,omitempty"`
	Text   string `xml:",chardata"`
}

// OMEMOEncrypted is the current OMEMO <encrypted> element shape, shared by
// Omemo1 (urn:xmpp:omemo:1) and Omemo2 (urn:xmpp:omemo:2) - identical
// element names, registered under both namespaces (see init()); which one
// an incoming stanza actually used is recovered from XMLName.Space after
// decoding (its tag deliberately omits a namespace, so Go's encoding/xml
// matches either and records whichever was actually present - confirmed
// against gosrc.io/xmpp's own decode path, which is plain
// d.DecodeElement() under the hood, not anything namespace-strict of its
// own). Keys are grouped per-recipient-JID under <keys jid='...'>, and the
// payload is an encrypted SCE envelope rather than the raw body, so there's
// no standalone <iv> (it's HKDF-derived from the key - see e2ee/omemo).
//
//	<encrypted xmlns='urn:xmpp:omemo:2'>
//	  <header sid='27183'>
//	    <keys jid='juliet@capulet.lit'>
//	      <key rid='31415' kex='true'>BASE64</key>
//	    </keys>
//	  </header>
//	  <payload>BASE64</payload>
//	</encrypted>
type OMEMOEncrypted struct {
	XMLName xml.Name    `xml:"encrypted"`
	Header  OMEMOHeader `xml:"header"`
	Payload string      `xml:"payload,omitempty"`
}

// Namespace implements stanza.MsgExtension.
func (c OMEMOEncrypted) Namespace() string {
	return c.XMLName.Space
}

type OMEMOHeader struct {
	SID  uint32      `xml:"sid,attr"`
	Keys []OMEMOKeys `xml:"keys"`
}

type OMEMOKeys struct {
	JID  string     `xml:"jid,attr"`
	Keys []OMEMOKey `xml:"key"`
}

type OMEMOKey struct {
	RID  uint32 `xml:"rid,attr"`
	Kex  bool   `xml:"kex,attr,omitempty"`
	Text string `xml:",chardata"`
}

// EME is the XEP-0380 Explicit Message Encryption hint, telling clients
// that can't decrypt <encrypted> what scheme was used instead of just
// showing nothing. The Go field is EncryptionNamespace (not Namespace) to
// avoid colliding with the stanza.MsgExtension.Namespace() method this
// type also needs - the wire attribute is still called "namespace".
type EME struct {
	XMLName             xml.Name `xml:"urn:xmpp:eme:0 encryption"`
	EncryptionNamespace  string   `xml:"namespace,attr"`
	Name                 string   `xml:"name,attr,omitempty"`
}

// Namespace implements stanza.MsgExtension.
func (c EME) Namespace() string {
	return c.XMLName.Space
}

func init() {
	stanza.TypeRegistry.MapExtension(stanza.PKTMessage, xml.Name{
		Space: "eu.siacs.conversations.axolotl",
		Local: "encrypted",
	}, OMEMO0Encrypted{})

	stanza.TypeRegistry.MapExtension(stanza.PKTMessage, xml.Name{
		Space: "urn:xmpp:omemo:1",
		Local: "encrypted",
	}, OMEMOEncrypted{})

	stanza.TypeRegistry.MapExtension(stanza.PKTMessage, xml.Name{
		Space: "urn:xmpp:omemo:2",
		Local: "encrypted",
	}, OMEMOEncrypted{})

	stanza.TypeRegistry.MapExtension(stanza.PKTMessage, xml.Name{
		Space: "urn:xmpp:eme:0",
		Local: "encryption",
	}, EME{})
}