package notify

import "bytes"

// AuthMode describes how notify authenticates the peer during bootstrap.
type AuthMode int

const (
	AuthNone AuthMode = iota
	AuthPSK
	AuthExternalPeer
)

// ProtectionMode describes how notify protects steady-state transport payloads.
type ProtectionMode int

const (
	ProtectionManaged ProtectionMode = iota
	ProtectionExternal
	ProtectionNested
)

// SecurityOptions describes the high-level auth/protection policy.
//
// The current implementation still exposes dedicated helper constructors such
// as UseModernPSKClient/Server and UsePSKOverExternalTransportClient/Server.
type SecurityOptions struct {
	AuthMode              AuthMode
	ProtectionMode        ProtectionMode
	SharedSecret          []byte
	RequireForwardSecrecy bool
}

func authModeName(mode AuthMode) string {
	switch mode {
	case AuthNone:
		return "none"
	case AuthPSK:
		return "psk"
	case AuthExternalPeer:
		return "external-peer"
	default:
		return "unknown"
	}
}

func protectionModeName(mode ProtectionMode) string {
	switch mode {
	case ProtectionManaged:
		return "managed"
	case ProtectionExternal:
		return "external"
	case ProtectionNested:
		return "nested"
	default:
		return "unknown"
	}
}

type transportProtectionProfile struct {
	mode                   ProtectionMode
	msgEn                  func([]byte, []byte) []byte
	msgDe                  func([]byte, []byte) []byte
	fastStreamEncode       transportFastStreamEncoder
	fastBulkEncode         transportFastBulkEncoder
	fastPlainEncode        transportFastPlainEncoder
	runtime                *modernPSKCodecRuntime
	secretKey              []byte
	keyMode                string
	sessionID              []byte
	forwardSecrecy         bool
	forwardSecrecyFallback bool
}

func cloneTransportProtectionKey(src []byte) []byte {
	if len(src) == 0 {
		return nil
	}
	return bytes.Clone(src)
}

func newTransportProtectionProfile(mode ProtectionMode, bundle modernPSKTransportBundle, runtime *modernPSKCodecRuntime, secretKey []byte) transportProtectionProfile {
	return transportProtectionProfile{
		mode:             mode,
		msgEn:            bundle.msgEn,
		msgDe:            bundle.msgDe,
		fastStreamEncode: bundle.fastStreamEncode,
		fastBulkEncode:   bundle.fastBulkEncode,
		fastPlainEncode:  bundle.fastPlainEncode,
		runtime:          runtime,
		secretKey:        cloneTransportProtectionKey(secretKey),
		keyMode:          defaultTransportKeyMode(mode, secretKey),
	}
}

func defaultTransportKeyMode(mode ProtectionMode, secretKey []byte) string {
	if len(secretKey) == 0 {
		return ""
	}
	switch mode {
	case ProtectionManaged, ProtectionNested:
		return peerAttachKeyModeStatic
	case ProtectionExternal:
		return transportKeyModeExternal
	default:
		return ""
	}
}

func cloneTransportSessionID(src []byte) []byte {
	if len(src) == 0 {
		return nil
	}
	return bytes.Clone(src)
}

func (p transportProtectionProfile) clone() transportProtectionProfile {
	p.secretKey = cloneTransportProtectionKey(p.secretKey)
	p.sessionID = cloneTransportSessionID(p.sessionID)
	return p
}

func (p transportProtectionProfile) withForwardSecrecyFallback(fallback bool) transportProtectionProfile {
	p = p.clone()
	p.forwardSecrecy = false
	p.forwardSecrecyFallback = fallback
	p.sessionID = nil
	return p
}

func passthroughTransportCodec(_ []byte, payload []byte) []byte {
	return payload
}

func passthroughFastPlainEncode(_ []byte, plainLen int, fill func([]byte) error) ([]byte, error) {
	if plainLen < 0 {
		return nil, errTransportPayloadEncryptFailed
	}
	buf := make([]byte, plainLen)
	if fill != nil {
		if err := fill(buf); err != nil {
			return nil, err
		}
	}
	return buf, nil
}

func buildExternalTransportBundle() modernPSKTransportBundle {
	fastPlainEncode := passthroughFastPlainEncode
	fastStreamEncode := func(secretKey []byte, dataID uint64, seq uint64, payload []byte) ([]byte, error) {
		return encodeStreamFastFramePayloadFast(fastPlainEncode, secretKey, streamFastDataFrame{
			DataID:  dataID,
			Seq:     seq,
			Payload: payload,
		})
	}
	fastBulkEncode := func(secretKey []byte, dataID uint64, seq uint64, payload []byte) ([]byte, error) {
		return encodeBulkFastFramePayloadFast(fastPlainEncode, secretKey, bulkFastFrame{
			Type:    bulkFastPayloadTypeData,
			DataID:  dataID,
			Seq:     seq,
			Payload: payload,
		})
	}
	return modernPSKTransportBundle{
		msgEn:            passthroughTransportCodec,
		msgDe:            passthroughTransportCodec,
		fastStreamEncode: fastStreamEncode,
		fastBulkEncode:   fastBulkEncode,
		fastPlainEncode:  fastPlainEncode,
	}
}

func defaultTransportProtectionProfile() transportProtectionProfile {
	return newTransportProtectionProfile(ProtectionManaged, defaultModernPSKTransportBundle(), nil, nil)
}

func (c *ClientCommon) clientTransportProtectionSnapshot() transportProtectionProfile {
	if c == nil {
		return transportProtectionProfile{}
	}
	if state := c.transportProtection.Load(); state != nil {
		return *state
	}
	return transportProtectionProfile{
		mode:             ProtectionManaged,
		msgEn:            c.msgEn,
		msgDe:            c.msgDe,
		fastStreamEncode: c.fastStreamEncode,
		fastBulkEncode:   c.fastBulkEncode,
		fastPlainEncode:  c.fastPlainEncode,
		runtime:          c.modernPSKRuntime,
		secretKey:        c.SecretKey,
		keyMode:          defaultTransportKeyMode(ProtectionManaged, c.SecretKey),
	}
}

func (c *ClientCommon) setClientTransportProtectionProfile(profile transportProtectionProfile) {
	if c == nil {
		return
	}
	profile.secretKey = cloneTransportProtectionKey(profile.secretKey)
	profile.sessionID = cloneTransportSessionID(profile.sessionID)
	c.msgEn = profile.msgEn
	c.msgDe = profile.msgDe
	c.fastStreamEncode = profile.fastStreamEncode
	c.fastBulkEncode = profile.fastBulkEncode
	c.fastPlainEncode = profile.fastPlainEncode
	c.modernPSKRuntime = profile.runtime
	c.SecretKey = profile.secretKey
	c.transportProtection.Store(&profile)
}

func (c *ClientCommon) clearClientSecurityProfiles() {
	if c == nil {
		return
	}
	c.securityConfigured = false
	c.securityAuthMode = AuthNone
	c.securityProtectionMode = ProtectionManaged
	c.securityBootstrap = transportProtectionProfile{}
	c.securitySteady = transportProtectionProfile{}
	c.securitySteadyNegotiated = transportProtectionProfile{}
	c.securityRequireForwardSecrecy = false
	c.peerAttachAuthenticated = false
	c.peerAttachAuthFallback = false
	c.peerAttachAt = 0
}

func (c *ClientCommon) configureClientSecurityProfiles(authMode AuthMode, protectionMode ProtectionMode, bootstrap transportProtectionProfile, steady transportProtectionProfile, requireForwardSecrecy bool) {
	if c == nil {
		return
	}
	c.securityConfigured = true
	c.securityAuthMode = authMode
	c.securityProtectionMode = protectionMode
	c.securityBootstrap = bootstrap.clone()
	c.securitySteady = steady.clone()
	c.securitySteadyNegotiated = steady.clone()
	c.securityRequireForwardSecrecy = requireForwardSecrecy
	c.setClientTransportProtectionProfile(bootstrap)
	c.securityReadyCheck = len(bootstrap.secretKey) != 0
	c.skipKeyExchange = true
}

func (c *ClientCommon) activateClientBootstrapTransportProtection() {
	if c == nil || !c.securityConfigured {
		return
	}
	c.resetClientNegotiatedSteadyTransportProtection()
	c.setClientTransportProtectionProfile(c.securityBootstrap)
}

func (c *ClientCommon) activateClientSteadyTransportProtection() {
	if c == nil || !c.securityConfigured {
		return
	}
	c.setClientTransportProtectionProfile(c.securitySteadyNegotiated)
}

func (c *ClientCommon) resetClientNegotiatedSteadyTransportProtection() {
	if c == nil {
		return
	}
	c.securitySteadyNegotiated = c.securitySteady.clone().withForwardSecrecyFallback(false)
}

func (c *ClientCommon) setClientNegotiatedSteadyTransportProtection(profile transportProtectionProfile) {
	if c == nil {
		return
	}
	c.securitySteadyNegotiated = profile.clone()
}

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

func (c *ClientCommon) clientRequiresForwardSecrecy() bool {
	if c == nil {
		return false
	}
	return c.securityRequireForwardSecrecy
}

func (s *ServerCommon) serverSupportsForwardSecrecy() bool {
	if s == nil || !s.securityConfigured {
		return false
	}
	if s.securityAuthMode != AuthPSK {
		return false
	}
	if s.securityProtectionMode == ProtectionExternal {
		return false
	}
	return len(s.securityBootstrap.secretKey) != 0
}

func (s *ServerCommon) serverRequiresForwardSecrecy() bool {
	if s == nil {
		return false
	}
	return s.securityRequireForwardSecrecy
}

func (s *ServerCommon) setServerDefaultTransportProtectionProfile(profile transportProtectionProfile) {
	if s == nil {
		return
	}
	profile.secretKey = cloneTransportProtectionKey(profile.secretKey)
	profile.sessionID = cloneTransportSessionID(profile.sessionID)
	s.defaultMsgEn = profile.msgEn
	s.defaultMsgDe = profile.msgDe
	s.defaultFastStreamEncode = profile.fastStreamEncode
	s.defaultFastBulkEncode = profile.fastBulkEncode
	s.defaultFastPlainEncode = profile.fastPlainEncode
	s.defaultModernPSKRuntime = profile.runtime
	s.SecretKey = profile.secretKey
}

func (s *ServerCommon) clearServerSecurityProfiles() {
	if s == nil {
		return
	}
	s.securityConfigured = false
	s.securityAuthMode = AuthNone
	s.securityProtectionMode = ProtectionManaged
	s.securityBootstrap = transportProtectionProfile{}
	s.securitySteady = transportProtectionProfile{}
	s.securityRequireForwardSecrecy = false
}

func (s *ServerCommon) configureServerSecurityProfiles(authMode AuthMode, protectionMode ProtectionMode, bootstrap transportProtectionProfile, steady transportProtectionProfile, requireForwardSecrecy bool) {
	if s == nil {
		return
	}
	s.securityConfigured = true
	s.securityAuthMode = authMode
	s.securityProtectionMode = protectionMode
	s.securityBootstrap = bootstrap.clone()
	s.securitySteady = steady.clone()
	s.securityRequireForwardSecrecy = requireForwardSecrecy
	s.setServerDefaultTransportProtectionProfile(bootstrap)
	s.securityReadyCheck = len(bootstrap.secretKey) != 0
}

func (s *ServerCommon) applyLogicalSteadyTransportProtection(logical *LogicalConn) {
	if s == nil || logical == nil || !s.securityConfigured {
		return
	}
	logical.applyTransportProtectionProfile(s.securitySteady)
}

func transportProtectionProfileFromAttachmentState(state *clientConnAttachmentState) transportProtectionProfile {
	if state == nil {
		return transportProtectionProfile{}
	}
	return transportProtectionProfile{
		mode:                   state.protectionMode,
		msgEn:                  state.msgEn,
		msgDe:                  state.msgDe,
		fastStreamEncode:       state.fastStreamEncode,
		fastBulkEncode:         state.fastBulkEncode,
		fastPlainEncode:        state.fastPlainEncode,
		runtime:                state.modernPSKRuntime,
		secretKey:              cloneTransportProtectionKey(state.secretKey),
		keyMode:                state.keyMode,
		sessionID:              cloneTransportSessionID(state.sessionID),
		forwardSecrecy:         state.forwardSecrecy,
		forwardSecrecyFallback: state.forwardSecrecyFallback,
	}
}

func (c *LogicalConn) transportProtectionProfileSnapshot() transportProtectionProfile {
	if c == nil {
		return transportProtectionProfile{}
	}
	return transportProtectionProfileFromAttachmentState(c.attachmentStateSnapshot())
}

func (c *LogicalConn) applyTransportProtectionProfile(profile transportProtectionProfile) {
	if c == nil {
		return
	}
	c.updateAttachmentState(func(state *clientConnAttachmentState) {
		state.protectionMode = profile.mode
		state.msgEn = profile.msgEn
		state.msgDe = profile.msgDe
		state.fastStreamEncode = profile.fastStreamEncode
		state.fastBulkEncode = profile.fastBulkEncode
		state.fastPlainEncode = profile.fastPlainEncode
		state.modernPSKRuntime = profile.runtime
		state.secretKey = cloneTransportProtectionKey(profile.secretKey)
		state.keyMode = profile.keyMode
		state.sessionID = cloneTransportSessionID(profile.sessionID)
		state.forwardSecrecy = profile.forwardSecrecy
		state.forwardSecrecyFallback = profile.forwardSecrecyFallback
	})
}