package e2ee

import (
	"context"
	"encoding/xml"
	"fmt"
	"time"

	"gosrc.io/xmpp"
	"gosrc.io/xmpp/stanza"
)

// FetchTimeout bounds how long a single PEP items request may take before
// giving up - telegabber is a background bridge relaying Telegram traffic,
// so a slow or unresponsive remote pubsub service must not stall the
// outbound message path indefinitely.
const FetchTimeout = 15 * time.Second

// FetchAndIngestPeer fetches peer's device-list and every not-yet-known
// device's bundle over PEP - using the already-established
// xmpp.Component.SendIQ IQ-result routing, not a new transport mechanism -
// and ingests them into backend under owner's identity. It is the
// generic, backend-agnostic counterpart to the xmpp package's PEP-serving
// IQ handlers (handleGetOMEMODeviceListIq/handleGetOMEMOBundleIq): driven
// entirely by whatever PEP node names backend itself reports via
// DeviceListNode/BundleNode, it has no OMEMO-specific knowledge of its own,
// so it works unchanged for any future Backend implementation.
//
// ctx should be a long-lived, process-scoped context (e.g.
// gateway.ShutdownCtx), not context.Background() - each PEP request derives
// its own FetchTimeout-bounded child from it (see fetchOneItem), and
// deriving from a shutdown-aware parent means process teardown cancels any
// in-flight request immediately instead of leaving it to run out its own
// timeout regardless of teardown.
func FetchAndIngestPeer(ctx context.Context, component *xmpp.Component, backend Backend, owner PeerID, peer PeerID) error {
	_, ownerBareJID, ok := SplitOwnedPeer(owner)
	if !ok {
		return fmt.Errorf("e2ee: FetchAndIngestPeer: %q is not an OwnedPeer", owner)
	}

	// Fetch under owner's currently negotiated variant (DefaultVariant for
	// a brand new chat) - that's what Encrypt will actually produce, so
	// there's no benefit to establishing a session under a different one.
	variant, err := backend.NegotiatedVariant(owner)
	if err != nil {
		return fmt.Errorf("e2ee: NegotiatedVariant: %w", err)
	}

	deviceListItem, err := fetchOneItem(ctx, component, ownerBareJID, string(peer), backend.DeviceListNode(variant))
	if err != nil {
		return fmt.Errorf("e2ee: fetch device list for %s: %w", peer, err)
	}
	raw, err := xml.Marshal(deviceListItem)
	if err != nil {
		return fmt.Errorf("e2ee: marshal device list for %s: %w", peer, err)
	}
	deviceIDs, err := backend.IngestRemoteDeviceList(owner, peer, DeviceListDoc{Raw: raw})
	if err != nil {
		return fmt.Errorf("e2ee: IngestRemoteDeviceList(%s): %w", peer, err)
	}

	// Only devices we don't already have a session for need their bundle
	// fetched - re-processing an existing device's bundle would rebuild its
	// session from scratch, discarding live ratchet state for no reason.
	known := map[DeviceID]bool{}
	if existing, err := backend.Devices(owner, peer); err == nil {
		for _, d := range existing {
			known[d.ID] = true
		}
	}

	var lastErr error
	fetchedAny := false
	for _, id := range deviceIDs {
		if known[id] {
			continue
		}
		bundleItem, err := fetchOneItem(ctx, component, ownerBareJID, string(peer), backend.BundleNode(variant, id))
		if err != nil {
			lastErr = fmt.Errorf("e2ee: fetch bundle for %s:%s: %w", peer, id, err)
			continue
		}
		bundleRaw, err := xml.Marshal(bundleItem)
		if err != nil {
			lastErr = fmt.Errorf("e2ee: marshal bundle for %s:%s: %w", peer, id, err)
			continue
		}
		if err := backend.IngestRemoteBundle(owner, peer, id, BundleDoc{Raw: bundleRaw}); err != nil {
			lastErr = fmt.Errorf("e2ee: IngestRemoteBundle(%s:%s): %w", peer, id, err)
			continue
		}
		fetchedAny = true
	}

	if !fetchedAny && lastErr != nil {
		return lastErr
	}
	return nil
}

// fetchOneItem requests the single most recent item of node from jid's PEP
// service, sent as if from ownerBareJID (the bridged chat's pseudo-JID doing
// the fetching), and returns its payload node.
func fetchOneItem(parent context.Context, component *xmpp.Component, ownerBareJID, jid, node string) (*stanza.Node, error) {
	iq, err := stanza.NewItemsRequest(jid, node, 1)
	if err != nil {
		return nil, err
	}
	iq.Attrs.From = ownerBareJID

	ctx, cancel := context.WithTimeout(parent, FetchTimeout)
	defer cancel()

	ch, err := component.SendIQ(ctx, iq)
	if err != nil {
		return nil, err
	}

	select {
	case result := <-ch:
		if result.Type == stanza.IQTypeError {
			return nil, fmt.Errorf("e2ee: %s returned an error for node %s", jid, node)
		}
		pubsub, ok := result.Payload.(*stanza.PubSubGeneric)
		if !ok || pubsub.Items == nil || len(pubsub.Items.List) == 0 {
			return nil, fmt.Errorf("e2ee: %s returned no items for node %s", jid, node)
		}
		item := pubsub.Items.List[0]
		if item.Any == nil {
			return nil, fmt.Errorf("e2ee: %s returned an empty item for node %s", jid, node)
		}
		return item.Any, nil
	case <-ctx.Done():
		return nil, fmt.Errorf("e2ee: timed out fetching node %s from %s", node, jid)
	}
}