package omemo

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

	"dev.narayana.im/narayana/telegabber/e2ee/omemo/libsignal"
)

// Bundle document parsing/encoding for all three OMEMO versions. As with
// devicelist.go, Omemo1 and Omemo2 share an identical element shape (spk/
// spks/ik/prekeys/pk) - only the namespace differs, per the same
// changelog-silence reasoning - so one xepBundle struct serves both,
// namespace-agnostic on Unmarshal, explicit on Marshal.
//
// Omemo0's element names (signedPreKeyPublic/signedPreKeySignature/
// identityKey/prekeys/preKeyPublic) do NOT carry over to Omemo1/Omemo2 -
// confirmed against the XEP-0384 XML schema directly, not assumed.
//
// None of the three bundle formats carry a registration id (unlike
// vanilla Signal protocol's own prekey-bundle API, which does) - OMEMO
// simply doesn't use it (XEP-0384's OMEMOKeyExchange has no such field
// either). parseBundle leaves RegistrationID as 0; callers pass that
// straight through to libsignal.RemoteBundle, which is the standard,
// harmless practice for this field in OMEMO implementations - device_id
// already does what registration_id would otherwise disambiguate.

type omemo0Bundle struct {
	XMLName               xml.Name         `xml:"eu.siacs.conversations.axolotl bundle"`
	SignedPreKeyPublic     omemo0IDText     `xml:"signedPreKeyPublic"`
	SignedPreKeySignature  string           `xml:"signedPreKeySignature"`
	IdentityKey            string           `xml:"identityKey"`
	PreKeys                omemo0PreKeysWrap `xml:"prekeys"`
}

type omemo0IDText struct {
	ID   uint32 `xml:"signedPreKeyId,attr"`
	Text string `xml:",chardata"`
}

type omemo0PreKeysWrap struct {
	List []omemo0PreKey `xml:"preKeyPublic"`
}

type omemo0PreKey struct {
	ID   uint32 `xml:"preKeyId,attr"`
	Text string `xml:",chardata"`
}

type xepBundle struct {
	XMLName xml.Name       `xml:"bundle"`
	SPK     xepIDText      `xml:"spk"`
	SPKS    string         `xml:"spks"`
	IK      string         `xml:"ik"`
	PreKeys xepPreKeysWrap `xml:"prekeys"`
}

type xepIDText struct {
	ID   uint32 `xml:"id,attr"`
	Text string `xml:",chardata"`
}

type xepPreKeysWrap struct {
	List []xepIDText `xml:"pk"`
}

// parsedBundle is the common, namespace-agnostic result of parsing any of
// the three bundle formats, in the raw (still base64-decoded) bytes
// libsignal expects.
type parsedBundle struct {
	Version               Version // which of the three OMEMO versions this bundle was actually fetched/published under - determines both which libsignal protocol version (V3 vs V4) a session built from it must use, AND (Omemo1 vs Omemo2 specifically) what to record as the chat's negotiated version, since that's a property of the bundle actually fetched, not of this backend's own prior state (see IngestRemoteBundle)
	SignedPreKeyID        uint32
	SignedPreKeyPublic    []byte
	SignedPreKeySignature []byte
	IdentityKeyPublic     []byte
	PreKeys               []parsedPreKey
}

type parsedPreKey struct {
	ID     uint32
	Public []byte
}

func parseBundle(raw []byte) (*parsedBundle, error) {
	var omemo0 omemo0Bundle
	if err := xml.Unmarshal(raw, &omemo0); err == nil {
		return decodeOmemo0Bundle(omemo0)
	}

	var xep xepBundle
	if err := xml.Unmarshal(raw, &xep); err != nil {
		return nil, err
	}
	return decodeXepBundle(xep)
}

func decodeOmemo0Bundle(b omemo0Bundle) (*parsedBundle, error) {
	spk, err := b64Decode(b.SignedPreKeyPublic.Text)
	if err != nil {
		return nil, fmt.Errorf("omemo: omemo0 bundle signedPreKeyPublic: %w", err)
	}
	sig, err := b64Decode(b.SignedPreKeySignature)
	if err != nil {
		return nil, fmt.Errorf("omemo: omemo0 bundle signedPreKeySignature: %w", err)
	}
	ik, err := b64Decode(b.IdentityKey)
	if err != nil {
		return nil, fmt.Errorf("omemo: omemo0 bundle identityKey: %w", err)
	}

	result := &parsedBundle{
		Version:               Omemo0,
		SignedPreKeyID:        b.SignedPreKeyPublic.ID,
		SignedPreKeyPublic:    spk,
		SignedPreKeySignature: sig,
		IdentityKeyPublic:     ik,
	}
	for _, pk := range b.PreKeys.List {
		data, err := b64Decode(pk.Text)
		if err != nil {
			return nil, fmt.Errorf("omemo: omemo0 bundle preKeyPublic %d: %w", pk.ID, err)
		}
		result.PreKeys = append(result.PreKeys, parsedPreKey{ID: pk.ID, Public: data})
	}
	return result, nil
}

func decodeXepBundle(b xepBundle) (*parsedBundle, error) {
	spk, err := b64Decode(b.SPK.Text)
	if err != nil {
		return nil, fmt.Errorf("omemo: bundle spk: %w", err)
	}
	sig, err := b64Decode(b.SPKS)
	if err != nil {
		return nil, fmt.Errorf("omemo: bundle spks: %w", err)
	}
	ik, err := b64Decode(b.IK)
	if err != nil {
		return nil, fmt.Errorf("omemo: bundle ik: %w", err)
	}

	version := Omemo2
	if b.XMLName.Space == omemo1NS {
		version = Omemo1
	}
	result := &parsedBundle{
		Version:               version,
		SignedPreKeyID:        b.SPK.ID,
		SignedPreKeyPublic:    spk,
		SignedPreKeySignature: sig,
		IdentityKeyPublic:     ik,
	}
	for _, pk := range b.PreKeys.List {
		data, err := b64Decode(pk.Text)
		if err != nil {
			return nil, fmt.Errorf("omemo: bundle pk %d: %w", pk.ID, err)
		}
		result.PreKeys = append(result.PreKeys, parsedPreKey{ID: pk.ID, Public: data})
	}
	return result, nil
}

// encodeBundle produces a bundle document for one of this gateway's own
// published devices, in the namespace matching version. Public keys are
// expected in the form DecodeSignedPreKey/DecodePreKey/
// DecodeIdentityPublicKey return - i.e. already picking the signature form
// (Omemo0 vs Omemo1/Omemo2) matching version, per the rule documented on
// libsignal.SignedPreKeyInfo.
func encodeBundle(version Version, identityKeyPublic []byte, spkID uint32, spkPublic, spkSignature []byte, preKeys []libsignal.PreKeyInfo) ([]byte, error) {
	if version == Omemo0 {
		b := omemo0Bundle{
			SignedPreKeyPublic:    omemo0IDText{ID: spkID, Text: b64Encode(spkPublic)},
			SignedPreKeySignature: b64Encode(spkSignature),
			IdentityKey:           b64Encode(identityKeyPublic),
		}
		for _, pk := range preKeys {
			b.PreKeys.List = append(b.PreKeys.List, omemo0PreKey{ID: pk.ID, Text: b64Encode(pk.PublicKey)})
		}
		return xml.Marshal(b)
	}

	ns := omemo1NS
	if version == Omemo2 {
		ns = omemo2NS
	}
	b := xepBundle{
		XMLName: xml.Name{Space: ns, Local: "bundle"},
		SPK:     xepIDText{ID: spkID, Text: b64Encode(spkPublic)},
		SPKS:    b64Encode(spkSignature),
		IK:      b64Encode(identityKeyPublic),
	}
	for _, pk := range preKeys {
		b.PreKeys.List = append(b.PreKeys.List, xepIDText{ID: pk.ID, Text: b64Encode(pk.PublicKey)})
	}
	return xml.Marshal(b)
}

func b64Encode(b []byte) string {
	return base64.StdEncoding.EncodeToString(b)
}

func b64Decode(s string) ([]byte, error) {
	return base64.StdEncoding.DecodeString(s)
}