package main

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha512"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"

	"log/slog"

	"filippo.io/edwards25519"
)

// https://www.cossacklabs.com/blog/introducing_secure_comparator/
// https://www.cossacklabs.com/files/secure-comparator-paper-rev12.pdf

type Server struct {
	a2, a3         ed25519.PrivateKey
	g2a, g3a       ed25519.PublicKey
	pA, qA, pB, qB *edwards25519.Point
}

func (p *Server) Bootstrap() ([]byte, error) {
	// Generate a2, a3
	var err error
	p.g2a, p.a2, err = ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("unable to generate ephemeral key pair: %w", err)
	}
	p.g3a, p.a3, err = ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("unable to generate ephemeral key pair: %w", err)
	}

	// G2a || SIGN(a2, G2a) || G3a || SIGN(a3, G3a)
	var out []byte
	out = append(out, p.g2a...)
	out = append(out, ed25519.Sign(p.a2, p.g2a)...)
	out = append(out, p.g3a...)
	out = append(out, ed25519.Sign(p.a3, p.g3a)...)

	return out, nil
}

func (p *Server) Prove(proof, secret []byte) ([]byte, error) {
	// Ensure expected length
	// G2b || G2b signature || G3b || G3b signature || Pb || Qb
	if len(proof) < 2*(ed25519.PublicKeySize+ed25519.SignatureSize)+2*ed25519.PublicKeySize {
		return nil, errors.New("message too short")
	}

	// Ensure valid signatures
	var (
		g2b    = proof[:ed25519.PublicKeySize]
		g2bSig = proof[ed25519.PublicKeySize : ed25519.PublicKeySize+ed25519.SignatureSize]
		g3b    = proof[ed25519.PublicKeySize+ed25519.SignatureSize : ed25519.PublicKeySize+ed25519.SignatureSize+ed25519.PublicKeySize]
		g3bSig = proof[ed25519.PublicKeySize+ed25519.SignatureSize+ed25519.PublicKeySize : ed25519.PublicKeySize+ed25519.SignatureSize+ed25519.PublicKeySize+ed25519.SignatureSize]
		pb     = proof[2*(ed25519.PublicKeySize+ed25519.SignatureSize) : 2*(ed25519.PublicKeySize+ed25519.SignatureSize)+ed25519.PublicKeySize]
		qb     = proof[2*(ed25519.PublicKeySize+ed25519.SignatureSize)+ed25519.PublicKeySize : 2*(ed25519.PublicKeySize+ed25519.SignatureSize)+2*ed25519.PublicKeySize]
	)

	if !ed25519.Verify(ed25519.PublicKey(g2b), g2b, g2bSig) {
		return nil, errors.New("invalid proof of possession")
	}
	if !ed25519.Verify(ed25519.PublicKey(g3b), g3b, g3bSig) {
		return nil, errors.New("invalid proof of possession")
	}

	// Compute generator groups
	a2, err := privateToScalar(p.a2)
	if err != nil {
		return nil, fmt.Errorf("unable to initialize a scalar from ephemeral key: %w", err)
	}
	g2bp, err := (&edwards25519.Point{}).SetBytes(g2b)
	if err != nil {
		return nil, fmt.Errorf("unable to initialize a point from ephemeral public key: %w", err)
	}
	G2 := (&edwards25519.Point{}).ScalarMult(a2, g2bp)

	a3, err := privateToScalar(p.a3)
	if err != nil {
		return nil, fmt.Errorf("unable to initialize a scalar from ephemeral key: %w", err)
	}
	g3bp, err := (&edwards25519.Point{}).SetBytes(g3b)
	if err != nil {
		return nil, fmt.Errorf("unable to initialize a point from ephemeral public key: %w", err)
	}
	G3 := (&edwards25519.Point{}).ScalarMult(a3, g3bp)

	var challenge [32]byte
	if _, err := io.ReadFull(rand.Reader, challenge[:]); err != nil {
		return nil, fmt.Errorf("unable to generate random challenge: %w", err)
	}
	s, err := edwards25519.NewScalar().SetBytesWithClamping(challenge[:])
	if err != nil {
		return nil, fmt.Errorf("unable to create a scalar from random challenge: %w", err)
	}

	// Compute secret
	h := sha512.Sum512_256(secret)
	x, err := edwards25519.NewScalar().SetBytesWithClamping(h[:])
	if err != nil {
		return nil, fmt.Errorf("unable to prepare secret as scalar: %w", err)
	}

	// Convert to point for arithmetic
	p.qB, err = (&edwards25519.Point{}).SetBytes(qb)
	if err != nil {
		return nil, fmt.Errorf("unable to prepare secret public key as point: %w", err)
	}
	p.pB, err = (&edwards25519.Point{}).SetBytes(pb)
	if err != nil {
		return nil, fmt.Errorf("unable to prepare secret public key as point: %w", err)
	}

	// Pa = s * G3
	p.pA = (&edwards25519.Point{}).ScalarMult(s, G3)
	// Qa = x * G2 + s * G
	p.qA = (&edwards25519.Point{}).VarTimeDoubleScalarBaseMult(x, G2, s)
	// Ra = a3 * (Qa - Qb)
	Ra := (&edwards25519.Point{}).ScalarMult(a3, (&edwards25519.Point{}).Subtract(p.qA, p.qB))

	var out []byte
	out = append(out, p.pA.Bytes()...)
	out = append(out, p.qA.Bytes()...)
	out = append(out, Ra.Bytes()...)

	return out, nil
}

func (p *Server) Verify(proof []byte) (bool, error) {
	// Ensure expected length
	// Rb
	if len(proof) < ed25519.PublicKeySize {
		return false, errors.New("message too short")
	}

	a3, err := privateToScalar(p.a3)
	if err != nil {
		return false, fmt.Errorf("unable to initialize a scalar from ephemeral key: %w", err)
	}

	Rb, _ := (&edwards25519.Point{}).SetBytes(proof)
	// Rab = a3 * Rb
	Rab := (&edwards25519.Point{}).ScalarMult(a3, Rb)

	// CHECK( Rb == (Pa - Pb) )
	if (&edwards25519.Point{}).Subtract(p.pA, p.pB).Equal(Rab) != 1 {
		return false, errors.New("invalid proof")
	}

	return true, nil
}

type Client struct {
	b2, b3   ed25519.PrivateKey
	g2b, g3b ed25519.PublicKey
	p, q     *edwards25519.Point
}

func (v *Client) Proof(in, secret []byte) ([]byte, error) {
	// Ensure expected length
	// g2a || g2a signature || g3a || g3a signature
	if len(in) < 2*(ed25519.PublicKeySize+ed25519.SignatureSize) {
		return nil, errors.New("message too short")
	}

	// Ensure valid signatures
	var (
		g2a    = in[:ed25519.PublicKeySize]
		g2aSig = in[ed25519.PublicKeySize : ed25519.PublicKeySize+ed25519.SignatureSize]
		g3a    = in[ed25519.PublicKeySize+ed25519.SignatureSize : ed25519.PublicKeySize+ed25519.SignatureSize+ed25519.PublicKeySize]
		g3aSig = in[ed25519.PublicKeySize+ed25519.SignatureSize+ed25519.PublicKeySize : ed25519.PublicKeySize+ed25519.SignatureSize+ed25519.PublicKeySize+ed25519.SignatureSize]
	)
	if !ed25519.Verify(ed25519.PublicKey(g2a), g2a, g2aSig) {
		return nil, errors.New("invalid proof of possession")
	}
	if !ed25519.Verify(ed25519.PublicKey(g3a), g3a, g3aSig) {
		return nil, errors.New("invalid proof of possession")
	}

	// Generate b2, b3
	var err error
	v.g2b, v.b2, err = ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("unable to generate ephemeral key pair: %w", err)
	}
	v.g3b, v.b3, err = ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, fmt.Errorf("unable to generate ephemeral key pair: %w", err)
	}

	// Convert to scalars for arithmetric
	b2, err := privateToScalar(v.b2)
	if err != nil {
		return nil, fmt.Errorf("unable to initialize a scalar from ephemeral key: %w", err)
	}
	g2ap, _ := (&edwards25519.Point{}).SetBytes(g2a)
	G2 := (&edwards25519.Point{}).ScalarMult(b2, g2ap)

	b3, err := privateToScalar(v.b3)
	if err != nil {
		return nil, fmt.Errorf("unable to initialize a scalar from ephemeral key: %w", err)
	}
	g3ap, _ := (&edwards25519.Point{}).SetBytes(g3a)
	G3 := (&edwards25519.Point{}).ScalarMult(b3, g3ap)

	var challenge [32]byte
	if _, err := io.ReadFull(rand.Reader, challenge[:]); err != nil {
		return nil, fmt.Errorf("unable to generate random challenge: %w", err)
	}
	r, _ := edwards25519.NewScalar().SetBytesWithClamping(challenge[:])

	// Compute secret (use SH512/256)
	h := sha512.Sum512_256(secret)
	y, err := edwards25519.NewScalar().SetBytesWithClamping(h[:])
	if err != nil {
		return nil, fmt.Errorf("unable to prepare secret as scalar: %w", err)
	}

	// Pb = r * G3
	v.p = (&edwards25519.Point{}).ScalarMult(r, G3)
	// Qb = y * G2 + r * G
	v.q = (&edwards25519.Point{}).VarTimeDoubleScalarBaseMult(y, G2, r)

	var out []byte
	out = append(out, v.g2b...)
	out = append(out, ed25519.Sign(v.b2, v.g2b)...)
	out = append(out, v.g3b...)
	out = append(out, ed25519.Sign(v.b3, v.g3b)...)
	out = append(out, v.p.Bytes()...)
	out = append(out, v.q.Bytes()...)

	return out, nil
}

func (v *Client) Verify(proof []byte) ([]byte, error) {
	// Ensure expected length
	// Pa || Qa || Ra
	if len(proof) < 3*(ed25519.PublicKeySize) {
		return nil, errors.New("message too short")
	}

	// Ensure valid signatures
	var (
		pa = proof[:ed25519.PublicKeySize]
		qa = proof[ed25519.PublicKeySize : 2*ed25519.PublicKeySize]
		ra = proof[2*ed25519.PublicKeySize : 3*ed25519.PublicKeySize]
	)

	b3, err := privateToScalar(v.b3)
	if err != nil {
		return nil, fmt.Errorf("unable to initialize a scalar from ephemeral key: %w", err)
	}

	Pa, _ := (&edwards25519.Point{}).SetBytes(pa)
	Qa, _ := (&edwards25519.Point{}).SetBytes(qa)
	Ra, _ := (&edwards25519.Point{}).SetBytes(ra)

	// Rb = b3 * (Qa - Qb)
	Rb := (&edwards25519.Point{}).ScalarMult(b3, (&edwards25519.Point{}).Subtract(Qa, v.q))
	// Rab = b3 * Ra
	Rab := (&edwards25519.Point{}).ScalarMult(b3, Ra)

	// CHECK( Rab == (Pa - Pb) )
	if (&edwards25519.Point{}).Subtract(Pa, v.p).Equal(Rab) != 1 {
		return nil, errors.New("invalid proof")
	}

	return Rb.Bytes(), nil
}

func privateToScalar(pk ed25519.PrivateKey) (*edwards25519.Scalar, error) {
	h := sha512.Sum512(pk.Seed())
	s, err := edwards25519.NewScalar().SetBytesWithClamping(h[:32])
	return s, err
}

func main() {
	logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{}))
	slog.SetDefault(logger)

	c1 := &Client{}
	logger.Info("C -> S: Session boostrap")

	session := &Server{}
	step1, err := session.Bootstrap()
	if err != nil {
		logger.Error("unable to initialize prover", "error", err)
		return
	}

	logger.Info("S -> C: Session public keys")
	fmt.Println(hex.Dump(step1))

	step2, err := c1.Proof(step1, []byte("very-secret-password"))
	if err != nil {
		logger.Error("unable to initialize prover", "error", err)
		return
	}

	logger.Info("C -> S: Client proof")
	fmt.Println(hex.Dump(step2))

	step3, err := session.Prove(step2, []byte("very-secret-password"))
	if err != nil {
		logger.Error("unable to compute proof", "error", err)
		return
	}

	logger.Info("S -> C: Session proof")
	fmt.Println(hex.Dump(step3))

	step4, err := c1.Verify(step3)
	if err != nil {
		logger.Error("unable to validate proof", "error", err)
		return
	}

	logger.Info("C -> S : Authenticate session proof")
	fmt.Println(hex.Dump(step4))

	valid, err := session.Verify(step4)
	if err != nil {
		logger.Error("unable to authenticate proof", "error", err)
		return
	}

	logger.Info("S: Client Authenticated?")
	fmt.Println(valid)
}