package notify

import (
	"bytes"
	"crypto/hmac"
	cryptorand "crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

const (
	peerAttachFeatureExplicitAuth uint64 = 1 << iota
	peerAttachFeatureChannelBinding
	peerAttachFeatureForwardSecrecy
)

const (
	peerAttachNonceSize = 16
	peerAttachReplayTTL = 30 * time.Second
)

var (
	errPeerAttachAuthInvalid               = errors.New("peer attach auth invalid")
	errPeerAttachReplayRejected            = errors.New("peer attach replay rejected")
	errPeerAttachReplayWindowFull          = errors.New("peer attach replay window full")
	errPeerAttachExplicitAuthRequired      = errors.New("peer attach explicit auth required")
	errPeerAttachChannelBindingRequired    = errors.New("peer attach channel binding required")
	errPeerAttachChannelBindingUnavailable = errors.New("peer attach channel binding unavailable")
	errPeerAttachForwardSecrecyRequired    = errors.New("peer attach forward secrecy required")
)

type peerAttachAuthResult struct {
	explicit             bool
	fallback             bool
	clientNonce          []byte
	serverNonce          []byte
	channelBinding       []byte
	clientECDHEPublicKey []byte
}

type peerAttachReplayCache struct {
	mu              sync.Mutex
	entries         map[string]time.Time
	rejects         atomic.Int64
	overflowRejects atomic.Int64
}

func newPeerAttachNonce() ([]byte, error) {
	buf := make([]byte, peerAttachNonceSize)
	if _, err := cryptorand.Read(buf); err != nil {
		return nil, err
	}
	return buf, nil
}

func appendPeerAttachAuthBytes(dst []byte, data []byte) []byte {
	dst = binary.BigEndian.AppendUint32(dst, uint32(len(data)))
	return append(dst, data...)
}

func appendPeerAttachAuthString(dst []byte, value string) []byte {
	return appendPeerAttachAuthBytes(dst, []byte(value))
}

func appendPeerAttachAuthBool(dst []byte, value bool) []byte {
	if value {
		return append(dst, 1)
	}
	return append(dst, 0)
}

func peerAttachRequestAuthPayload(req peerAttachRequest, channelBinding []byte) []byte {
	buf := make([]byte, 0, 96+len(req.PeerID)+len(channelBinding))
	buf = appendPeerAttachAuthString(buf, "notify/peer-attach/request-auth/v1")
	buf = binary.BigEndian.AppendUint64(buf, req.Features)
	buf = appendPeerAttachAuthString(buf, req.PeerID)
	buf = appendPeerAttachAuthBytes(buf, req.ClientNonce)
	if supportsPeerAttachChannelBinding(req.Features) {
		buf = appendPeerAttachAuthBytes(buf, channelBinding)
	}
	return buf
}

func peerAttachResponseAuthPayload(req peerAttachRequest, resp peerAttachResponse, channelBinding []byte) []byte {
	buf := make([]byte, 0, 160+len(req.PeerID)+len(resp.PeerID)+len(resp.Error)+len(channelBinding))
	buf = appendPeerAttachAuthString(buf, "notify/peer-attach/response-auth/v1")
	buf = binary.BigEndian.AppendUint64(buf, req.Features)
	buf = appendPeerAttachAuthString(buf, req.PeerID)
	buf = appendPeerAttachAuthBytes(buf, req.ClientNonce)
	buf = binary.BigEndian.AppendUint64(buf, resp.Features)
	buf = appendPeerAttachAuthString(buf, resp.PeerID)
	buf = appendPeerAttachAuthBool(buf, resp.Accepted)
	buf = appendPeerAttachAuthBool(buf, resp.Reused)
	buf = appendPeerAttachAuthString(buf, resp.Error)
	buf = appendPeerAttachAuthBytes(buf, resp.ServerNonce)
	if supportsPeerAttachChannelBinding(resp.Features) {
		buf = appendPeerAttachAuthBytes(buf, channelBinding)
	}
	return buf
}

func signPeerAttachPayload(secretKey []byte, payload []byte) []byte {
	if len(secretKey) == 0 {
		return nil
	}
	mac := hmac.New(sha256.New, secretKey)
	_, _ = mac.Write(payload)
	return mac.Sum(nil)
}

func computePeerAttachRequestAuthTag(secretKey []byte, req peerAttachRequest, channelBinding []byte) []byte {
	return signPeerAttachPayload(secretKey, peerAttachRequestAuthPayload(req, channelBinding))
}

func computePeerAttachResponseAuthTag(secretKey []byte, req peerAttachRequest, resp peerAttachResponse, channelBinding []byte) []byte {
	return signPeerAttachPayload(secretKey, peerAttachResponseAuthPayload(req, resp, channelBinding))
}

func supportsExplicitPeerAttachAuth(features uint64) bool {
	return features&peerAttachFeatureExplicitAuth != 0
}

func supportsPeerAttachChannelBinding(features uint64) bool {
	return features&peerAttachFeatureChannelBinding != 0
}

func supportsPeerAttachForwardSecrecy(features uint64) bool {
	return features&peerAttachFeatureForwardSecrecy != 0
}

func classifyPeerAttachRejectCounter(s *ServerCommon, err error) {
	if s == nil || err == nil {
		return
	}
	switch {
	case errors.Is(err, errPeerAttachReplayRejected):
		s.peerAttachReplay.rejects.Add(1)
	case errors.Is(err, errPeerAttachReplayWindowFull):
		s.peerAttachReplay.overflowRejects.Add(1)
	case errors.Is(err, errPeerAttachExplicitAuthRequired), errors.Is(err, errPeerAttachChannelBindingRequired), errors.Is(err, errPeerAttachForwardSecrecyRequired):
		s.peerAttachDowngradeRejectCount.Add(1)
	case errors.Is(err, errPeerAttachChannelBindingUnavailable):
		s.peerAttachBindingRejectCount.Add(1)
	default:
		s.peerAttachAuthRejectCount.Add(1)
	}
}

func (c *ClientCommon) shouldUseExplicitPeerAttachAuth() bool {
	if c == nil || !c.securityConfigured {
		return false
	}
	return c.securityAuthMode == AuthPSK && len(c.securityBootstrap.secretKey) != 0
}

func resolvePeerAttachChannelBinding(provider PeerAttachChannelBindingProvider, role PeerAttachChannelBindingRole, peerID string, conn net.Conn) ([]byte, error) {
	if provider == nil {
		return nil, nil
	}
	if conn == nil {
		return nil, errPeerAttachChannelBindingUnavailable
	}
	binding, err := provider(PeerAttachChannelBindingContext{
		Role:   role,
		PeerID: peerID,
		Conn:   conn,
	})
	if err != nil || len(binding) == 0 {
		return nil, errPeerAttachChannelBindingUnavailable
	}
	return bytes.Clone(binding), nil
}

func (c *ClientCommon) buildPeerAttachRequest(peerID string) (peerAttachRequest, peerAttachRequestState, error) {
	cfg := c.peerAttachSecuritySnapshot()
	req := peerAttachRequest{
		PeerID: stringsTrimSpaceNoAlloc(peerID),
	}
	if !c.shouldUseExplicitPeerAttachAuth() {
		if c.clientRequiresForwardSecrecy() {
			return peerAttachRequest{}, peerAttachRequestState{}, errPeerAttachForwardSecrecyRequired
		}
		if cfg.requireExplicitAuth {
			return peerAttachRequest{}, peerAttachRequestState{}, errPeerAttachExplicitAuthRequired
		}
		return req, peerAttachRequestState{}, nil
	}
	nonce, err := newPeerAttachNonce()
	if err != nil {
		return peerAttachRequest{}, peerAttachRequestState{}, err
	}
	req.Features = peerAttachFeatureExplicitAuth
	requestState := peerAttachRequestState{}
	var channelBinding []byte
	if cfg.channelBinding != nil {
		channelBinding, err = resolvePeerAttachChannelBinding(cfg.channelBinding, PeerAttachChannelBindingRoleClient, req.PeerID, c.clientTransportConnSnapshot())
		if err != nil {
			return peerAttachRequest{}, peerAttachRequestState{}, err
		}
	}
	if len(channelBinding) != 0 {
		req.Features |= peerAttachFeatureChannelBinding
	}
	if cfg.requireChannelBinding && !supportsPeerAttachChannelBinding(req.Features) {
		return peerAttachRequest{}, peerAttachRequestState{}, errPeerAttachChannelBindingRequired
	}
	if c.clientSupportsForwardSecrecy() {
		requestState.forwardSecrecy, err = newPeerAttachForwardSecrecyClientState()
		if err != nil {
			return peerAttachRequest{}, peerAttachRequestState{}, err
		}
		req.Features |= peerAttachFeatureForwardSecrecy
		req.ClientECDHEPublicKey = bytes.Clone(requestState.forwardSecrecy.publicKey)
	}
	req.ClientNonce = nonce
	req.AuthTag = computePeerAttachRequestAuthTag(c.securityBootstrap.secretKey, req, channelBinding)
	return req, requestState, nil
}

func (c *ClientCommon) verifyPeerAttachResponse(req peerAttachRequest, resp peerAttachResponse, requestState peerAttachRequestState) (peerAttachResponseVerifyResult, error) {
	cfg := c.peerAttachSecuritySnapshot()
	baseSteady := transportProtectionProfile{}
	if c != nil {
		baseSteady = c.securitySteady.clone().withForwardSecrecyFallback(false)
	}
	result := peerAttachResponseVerifyResult{steadyProfile: baseSteady}
	if c == nil || !c.shouldUseExplicitPeerAttachAuth() {
		if c.clientRequiresForwardSecrecy() {
			return peerAttachResponseVerifyResult{}, errPeerAttachForwardSecrecyRequired
		}
		if cfg.requireExplicitAuth {
			return peerAttachResponseVerifyResult{}, errPeerAttachExplicitAuthRequired
		}
		return result, nil
	}
	if !supportsExplicitPeerAttachAuth(resp.Features) {
		if c.clientRequiresForwardSecrecy() {
			return peerAttachResponseVerifyResult{}, errPeerAttachForwardSecrecyRequired
		}
		if cfg.requireExplicitAuth {
			return peerAttachResponseVerifyResult{}, errPeerAttachExplicitAuthRequired
		}
		if c.clientSupportsForwardSecrecy() {
			result.steadyProfile = result.steadyProfile.withForwardSecrecyFallback(true)
		}
		result.authFallback = true
		return result, nil
	}
	var channelBinding []byte
	if supportsPeerAttachChannelBinding(req.Features) {
		if cfg.channelBinding == nil {
			return peerAttachResponseVerifyResult{}, errPeerAttachChannelBindingUnavailable
		}
		if !supportsPeerAttachChannelBinding(resp.Features) {
			return peerAttachResponseVerifyResult{}, errPeerAttachChannelBindingRequired
		}
		var err error
		channelBinding, err = resolvePeerAttachChannelBinding(cfg.channelBinding, PeerAttachChannelBindingRoleClient, req.PeerID, c.clientTransportConnSnapshot())
		if err != nil {
			return peerAttachResponseVerifyResult{}, err
		}
	} else if cfg.requireChannelBinding {
		return peerAttachResponseVerifyResult{}, errPeerAttachChannelBindingRequired
	}
	if len(resp.ServerNonce) != peerAttachNonceSize {
		return peerAttachResponseVerifyResult{}, errPeerAttachAuthInvalid
	}
	expected := computePeerAttachResponseAuthTag(c.securityBootstrap.secretKey, req, peerAttachResponse{
		PeerID:      resp.PeerID,
		Accepted:    resp.Accepted,
		Reused:      resp.Reused,
		Error:       resp.Error,
		Features:    resp.Features,
		ServerNonce: resp.ServerNonce,
	}, channelBinding)
	if !hmac.Equal(resp.AuthTag, expected) {
		return peerAttachResponseVerifyResult{}, errPeerAttachAuthInvalid
	}
	if requestState.forwardSecrecy == nil || !supportsPeerAttachForwardSecrecy(req.Features) {
		return result, nil
	}
	if !supportsPeerAttachForwardSecrecy(resp.Features) {
		if c.clientRequiresForwardSecrecy() {
			return peerAttachResponseVerifyResult{}, errPeerAttachForwardSecrecyRequired
		}
		result.steadyProfile = result.steadyProfile.withForwardSecrecyFallback(true)
		return result, nil
	}
	if resp.KeyMode != "" && resp.KeyMode != peerAttachKeyModeECDHE {
		return peerAttachResponseVerifyResult{}, errPeerAttachForwardSecrecyInvalid
	}
	profile, err := derivePeerAttachForwardSecrecyTransportProfile(c.securitySteady, c.securityBootstrap.secretKey, requestState.forwardSecrecy.privateKey, resp.ServerECDHEPublicKey, req, resp)
	if err != nil {
		return peerAttachResponseVerifyResult{}, err
	}
	result.steadyProfile = profile
	return result, nil
}

func (s *ServerCommon) validatePeerAttachRequestAuth(logical *LogicalConn, transport net.Conn, req peerAttachRequest) (peerAttachAuthResult, error) {
	cfg := s.peerAttachSecuritySnapshot()
	if !supportsExplicitPeerAttachAuth(req.Features) {
		if s.serverRequiresForwardSecrecy() {
			return peerAttachAuthResult{}, errPeerAttachForwardSecrecyRequired
		}
		if cfg.requireExplicitAuth {
			return peerAttachAuthResult{}, errPeerAttachExplicitAuthRequired
		}
		return peerAttachAuthResult{fallback: true}, nil
	}
	if logical == nil {
		return peerAttachAuthResult{}, errPeerAttachAuthInvalid
	}
	var channelBinding []byte
	if supportsPeerAttachChannelBinding(req.Features) {
		if cfg.channelBinding == nil {
			return peerAttachAuthResult{}, errPeerAttachChannelBindingUnavailable
		}
		var err error
		channelBinding, err = resolvePeerAttachChannelBinding(cfg.channelBinding, PeerAttachChannelBindingRoleServer, req.PeerID, transport)
		if err != nil {
			return peerAttachAuthResult{}, err
		}
	} else if cfg.requireChannelBinding {
		return peerAttachAuthResult{}, errPeerAttachChannelBindingRequired
	}
	secretKey := logical.secretKeySnapshot()
	if len(secretKey) == 0 {
		return peerAttachAuthResult{}, errPeerAttachAuthInvalid
	}
	if len(req.ClientNonce) != peerAttachNonceSize {
		return peerAttachAuthResult{}, errPeerAttachAuthInvalid
	}
	expected := computePeerAttachRequestAuthTag(secretKey, peerAttachRequest{
		PeerID:      req.PeerID,
		Features:    req.Features,
		ClientNonce: req.ClientNonce,
	}, channelBinding)
	if !hmac.Equal(req.AuthTag, expected) {
		return peerAttachAuthResult{}, errPeerAttachAuthInvalid
	}
	if supportsPeerAttachForwardSecrecy(req.Features) {
		if len(req.ClientECDHEPublicKey) != peerAttachECDHEPublicKeySize {
			return peerAttachAuthResult{}, errPeerAttachForwardSecrecyInvalid
		}
	}
	if err := s.acceptPeerAttachReplay(req.PeerID, req.ClientNonce, time.Now(), cfg.replayWindow, cfg.replayCapacity); err != nil {
		return peerAttachAuthResult{}, err
	}
	serverNonce, err := newPeerAttachNonce()
	if err != nil {
		return peerAttachAuthResult{}, err
	}
	return peerAttachAuthResult{
		explicit:             true,
		clientNonce:          bytes.Clone(req.ClientNonce),
		serverNonce:          serverNonce,
		channelBinding:       channelBinding,
		clientECDHEPublicKey: bytes.Clone(req.ClientECDHEPublicKey),
	}, nil
}

func (s *ServerCommon) signPeerAttachResponse(logical *LogicalConn, req peerAttachRequest, resp *peerAttachResponse, auth peerAttachAuthResult) {
	if s == nil || logical == nil || resp == nil || !auth.explicit {
		return
	}
	secretKey := logical.secretKeySnapshot()
	if len(secretKey) == 0 {
		return
	}
	resp.Features |= peerAttachFeatureExplicitAuth
	if len(auth.channelBinding) != 0 {
		resp.Features |= peerAttachFeatureChannelBinding
	}
	resp.ServerNonce = bytes.Clone(auth.serverNonce)
	resp.AuthTag = computePeerAttachResponseAuthTag(secretKey, req, *resp, auth.channelBinding)
}

func (s *ServerCommon) preparePeerAttachSteadyTransportProfile(logical *LogicalConn, req peerAttachRequest, resp *peerAttachResponse, auth peerAttachAuthResult) (transportProtectionProfile, error) {
	if s == nil {
		return transportProtectionProfile{}, nil
	}
	profile := s.securitySteady.clone().withForwardSecrecyFallback(false)
	if resp != nil && auth.explicit {
		resp.Features |= peerAttachFeatureExplicitAuth
		if len(auth.channelBinding) != 0 {
			resp.Features |= peerAttachFeatureChannelBinding
		}
		resp.ServerNonce = bytes.Clone(auth.serverNonce)
	}
	if resp != nil && resp.KeyMode == "" {
		resp.KeyMode = profile.keyMode
	}
	if !s.serverSupportsForwardSecrecy() {
		return profile, nil
	}
	if !auth.explicit || !supportsPeerAttachForwardSecrecy(req.Features) {
		if s.serverRequiresForwardSecrecy() {
			return transportProtectionProfile{}, errPeerAttachForwardSecrecyRequired
		}
		return profile.withForwardSecrecyFallback(true), nil
	}
	fsState, err := newPeerAttachForwardSecrecyClientState()
	if err != nil {
		return transportProtectionProfile{}, err
	}
	if resp != nil {
		resp.Features |= peerAttachFeatureForwardSecrecy
		resp.KeyMode = peerAttachKeyModeECDHE
		resp.ServerECDHEPublicKey = bytes.Clone(fsState.publicKey)
	}
	return derivePeerAttachForwardSecrecyTransportProfile(s.securitySteady, logical.secretKeySnapshot(), fsState.privateKey, auth.clientECDHEPublicKey, req, *resp)
}

func (s *ServerCommon) acceptPeerAttachReplay(peerID string, nonce []byte, now time.Time, window time.Duration, capacity int) error {
	if s == nil || len(nonce) == 0 {
		return nil
	}
	cache := &s.peerAttachReplay
	key := peerID + "\x00" + string(nonce)
	expireBefore := now.Add(-window)
	cache.mu.Lock()
	defer cache.mu.Unlock()
	if cache.entries == nil {
		cache.entries = make(map[string]time.Time)
	}
	for replayKey, seenAt := range cache.entries {
		if seenAt.Before(expireBefore) {
			delete(cache.entries, replayKey)
		}
	}
	if seenAt, ok := cache.entries[key]; ok && !seenAt.Before(expireBefore) {
		return errPeerAttachReplayRejected
	}
	if capacity > 0 && len(cache.entries) >= capacity {
		return errPeerAttachReplayWindowFull
	}
	cache.entries[key] = now
	return nil
}

func (s *ServerCommon) peerAttachReplayRejectCountSnapshot() int64 {
	if s == nil {
		return 0
	}
	return s.peerAttachReplay.rejects.Load()
}

func (s *ServerCommon) peerAttachReplayOverflowRejectCountSnapshot() int64 {
	if s == nil {
		return 0
	}
	return s.peerAttachReplay.overflowRejects.Load()
}

func (c *ClientCommon) markClientPeerAttachAuthenticated(fallback bool, at time.Time) {
	if c == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	c.peerAttachAuthenticated = true
	c.peerAttachAuthFallback = fallback
	c.peerAttachAt = at.UnixNano()
}

func (c *ClientCommon) resetClientPeerAttachAuth() {
	if c == nil {
		return
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	c.peerAttachAuthenticated = false
	c.peerAttachAuthFallback = false
	c.peerAttachAt = 0
}

func (c *ClientCommon) clientPeerAttachAuthSnapshot() (bool, bool, time.Time) {
	if c == nil {
		return false, false, time.Time{}
	}
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.peerAttachAt == 0 {
		return c.peerAttachAuthenticated, c.peerAttachAuthFallback, time.Time{}
	}
	return c.peerAttachAuthenticated, c.peerAttachAuthFallback, time.Unix(0, c.peerAttachAt)
}

func stringsTrimSpaceNoAlloc(value string) string {
	start := 0
	for start < len(value) {
		switch value[start] {
		case ' ', '\t', '\n', '\r':
			start++
		default:
			goto endStart
		}
	}
	return ""
endStart:
	end := len(value)
	for end > start {
		switch value[end-1] {
		case ' ', '\t', '\n', '\r':
			end--
		default:
			return value[start:end]
		}
	}
	return value[start:end]
}