package lafsession

import (
	"crypto/ecdh"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"hash"
)

const (
	ProtocolVersion                   uint16 = 1
	CipherSuiteX25519ChaCha20Poly1305 uint16 = 1

	msgClientHello    byte = 1
	msgServerHello    byte = 2
	msgClientFinished byte = 3
	msgServerFinished byte = 4

	ephPubLen = 32
	idPubLen  = 32
	sigLen    = ed25519.SignatureSize

	helloMsgLen = 1 + 2 + 2 + ephPubLen + idPubLen + sigLen

	finishedMsgLen = 1 + sigLen

	transcriptCap = helloMsgLen*2 + finishedMsgLen*2
)

type HandshakeRole int

const (
	RoleClient HandshakeRole = iota
	RoleServer
)

type peerAcceptFunc func(peerPub []byte) bool

type handshakeState struct {
	role       HandshakeRole
	identity   *Identity
	acceptPeer peerAcceptFunc

	ephPriv *ecdh.PrivateKey
	ephPub  [ephPubLen]byte

	peerEphPub [ephPubLen]byte
	peerIDPub  [idPubLen]byte
	hasPeer    bool

	transcript []byte
	hasher     hash.Hash
	hashBuf    [32]byte
	labelBuf   [1 + 32]byte
	msgBuf     []byte

	version uint16
	suite   uint16
}

func newHandshake(role HandshakeRole, id *Identity, accept peerAcceptFunc) (*handshakeState, error) {
	curve := ecdh.X25519()
	eph, err := curve.GenerateKey(rand.Reader)
	if err != nil {
		return nil, err
	}
	h := &handshakeState{
		role:       role,
		identity:   id,
		acceptPeer: accept,
		ephPriv:    eph,
		transcript: make([]byte, 0, transcriptCap),
		hasher:     sha256.New(),
		msgBuf:     make([]byte, 0, helloMsgLen),
		version:    ProtocolVersion,
		suite:      CipherSuiteX25519ChaCha20Poly1305,
	}
	copy(h.ephPub[:], eph.PublicKey().Bytes())
	return h, nil
}

func (h *handshakeState) resetHasherFromTranscript() {
	h.hasher.Reset()
	_, _ = h.hasher.Write(h.transcript)
}

func (h *handshakeState) transcriptSum(dst []byte) []byte {
	return h.hasher.Sum(h.hashBuf[:0])
}

func (h *handshakeState) writeTranscript(p []byte) {
	h.transcript = append(h.transcript, p...)
	_, _ = h.hasher.Write(p)
}

func (h *handshakeState) buildClientHello() ([]byte, error) {
	if h.role != RoleClient {
		return nil, ErrHandshakeState
	}
	msg := h.msgBuf[:0]
	msg = append(msg, msgClientHello)
	msg = binary.BigEndian.AppendUint16(msg, h.version)
	msg = binary.BigEndian.AppendUint16(msg, h.suite)
	msg = append(msg, h.ephPub[:]...)
	msg = append(msg, h.identity.Public...)

	baseLen := len(msg)
	h.writeTranscript(msg)
	sum := h.transcriptSum(nil)
	sig := h.identity.Sign(sum)
	msg = append(msg, sig...)
	h.writeTranscript(sig)
	h.msgBuf = msg
	out := make([]byte, len(msg))
	copy(out, msg)
	_ = baseLen
	return out, nil
}

func (h *handshakeState) processClientHello(msg []byte) error {
	if h.role != RoleServer {
		return ErrHandshakeState
	}
	if len(msg) < helloMsgLen {
		return ErrShortMessage
	}
	if msg[0] != msgClientHello {
		return ErrHandshakeState
	}
	ver := binary.BigEndian.Uint16(msg[1:3])
	suite := binary.BigEndian.Uint16(msg[3:5])
	if ver != ProtocolVersion {
		return ErrBadVersion
	}
	if suite != CipherSuiteX25519ChaCha20Poly1305 {
		return ErrBadCipher
	}
	off := 5
	peerEph := msg[off : off+ephPubLen]
	off += ephPubLen
	peerID := msg[off : off+idPubLen]
	off += idPubLen
	sig := msg[off : off+sigLen]

	if h.acceptPeer != nil && !h.acceptPeer(peerID) {
		return ErrIdentityRejected
	}

	base := msg[:5+ephPubLen+idPubLen]
	h.hasher.Reset()
	_, _ = h.hasher.Write(base)
	sum := h.transcriptSum(nil)
	if !VerifyPeer(peerID, sum, sig) {
		h.hasher.Reset()
		return ErrBadSignature
	}

	h.version = ver
	h.suite = suite
	copy(h.peerEphPub[:], peerEph)
	copy(h.peerIDPub[:], peerID)
	h.hasPeer = true

	h.transcript = append(h.transcript[:0], base...)
	h.transcript = append(h.transcript, sig...)
	_, _ = h.hasher.Write(sig)
	return nil
}

func (h *handshakeState) buildServerHello() ([]byte, error) {
	if h.role != RoleServer {
		return nil, ErrHandshakeState
	}
	msg := h.msgBuf[:0]
	msg = append(msg, msgServerHello)
	msg = binary.BigEndian.AppendUint16(msg, h.version)
	msg = binary.BigEndian.AppendUint16(msg, h.suite)
	msg = append(msg, h.ephPub[:]...)
	msg = append(msg, h.identity.Public...)

	h.writeTranscript(msg)
	sum := h.transcriptSum(nil)
	sig := h.identity.Sign(sum)
	msg = append(msg, sig...)
	h.writeTranscript(sig)
	h.msgBuf = msg

	out := make([]byte, len(msg))
	copy(out, msg)
	return out, nil
}

func (h *handshakeState) processServerHello(msg []byte) error {
	if h.role != RoleClient {
		return ErrHandshakeState
	}
	if len(msg) < helloMsgLen {
		return ErrShortMessage
	}
	if msg[0] != msgServerHello {
		return ErrHandshakeState
	}
	ver := binary.BigEndian.Uint16(msg[1:3])
	suite := binary.BigEndian.Uint16(msg[3:5])
	if ver != ProtocolVersion {
		return ErrBadVersion
	}
	if suite != CipherSuiteX25519ChaCha20Poly1305 {
		return ErrBadCipher
	}
	off := 5
	peerEph := msg[off : off+ephPubLen]
	off += ephPubLen
	peerID := msg[off : off+idPubLen]
	off += idPubLen
	sig := msg[off : off+sigLen]

	if h.acceptPeer != nil && !h.acceptPeer(peerID) {
		return ErrIdentityRejected
	}

	base := msg[:5+ephPubLen+idPubLen]
	tLen := len(h.transcript)
	h.writeTranscript(base)
	sum := h.transcriptSum(nil)
	if !VerifyPeer(peerID, sum, sig) {
		h.transcript = h.transcript[:tLen]
		h.resetHasherFromTranscript()
		return ErrBadSignature
	}

	copy(h.peerEphPub[:], peerEph)
	copy(h.peerIDPub[:], peerID)
	h.hasPeer = true
	h.writeTranscript(sig)
	return nil
}

func (h *handshakeState) buildFinished(msgType byte) ([]byte, error) {
	sum := h.transcriptSum(nil)
	h.labelBuf[0] = msgType
	copy(h.labelBuf[1:], sum)
	sig := h.identity.Sign(h.labelBuf[:])
	msg := h.msgBuf[:0]
	msg = append(msg, msgType)
	msg = append(msg, sig...)
	h.writeTranscript(msg)
	h.msgBuf = msg

	out := make([]byte, len(msg))
	copy(out, msg)
	return out, nil
}

func (h *handshakeState) processFinished(msg []byte, expectedType byte, peerID []byte) error {
	if len(msg) < finishedMsgLen {
		return ErrShortMessage
	}
	if msg[0] != expectedType {
		return ErrHandshakeState
	}
	sig := msg[1 : 1+sigLen]
	sum := h.transcriptSum(nil)
	h.labelBuf[0] = expectedType
	copy(h.labelBuf[1:], sum)
	if !VerifyPeer(peerID, h.labelBuf[:], sig) {
		return ErrBadSignature
	}
	h.writeTranscript(msg[:finishedMsgLen])
	return nil
}

func (h *handshakeState) deriveKeys() (*trafficKeys, error) {
	curve := ecdh.X25519()
	peerPub, err := curve.NewPublicKey(h.peerEphPub[:])
	if err != nil {
		return nil, err
	}
	shared, err := h.ephPriv.ECDH(peerPub)
	if err != nil {
		return nil, err
	}
	defer zero(shared)
	sum := h.transcriptSum(nil)
	return deriveTrafficKeys(shared, sum, h.version, h.suite)
}

func (h *handshakeState) peerIDSlice() []byte {
	if !h.hasPeer {
		return nil
	}
	return h.peerIDPub[:]
}