package testkit

import (
	"bytes"
	"errors"
	"math/big"
)

// Re-encodings of BLS12-381 points in the compressed form of drand (spec
// §12.2), for the mutations and tests of the canonical point encoding. They
// work on the bytes alone: a flag byte whose low five bits start a
// big-endian coordinate, and coordinates of 48 bytes, x in G1 and c1 then c0
// in G2.

// FieldModulus is p, the modulus of the base field of BLS12-381.
var FieldModulus, _ = new(big.Int).SetString("1a0111ea397fe69a4b1ba7b6434bacd764774b84f38512bf6730d2a0f6b0f6241eabfffeb153ffffb9feffffffffaaab", 16)

// Flags of the first byte of a compressed point (spec §12.2).
const (
	FlagCompressed = 0x80
	FlagInfinity   = 0x40
	FlagSort       = 0x20
	flagMask       = FlagCompressed | FlagInfinity | FlagSort
)

// CoordinateLen is the size of a coordinate of Fp, and of a compressed point
// of G1; a compressed point of G2 takes two.
const CoordinateLen = 48

// AddModulus returns enc with p added to the 48-byte coordinate at byte
// offset at: 0 for x in G1 and for c1 in G2, 48 for c0 in G2. The flags are
// kept. The result reduces modulo p to the coordinate of enc, so a decoder
// that reduces coordinates reads the point of enc, but it is not a canonical
// encoding (spec §12.2). It fails when the sum does not fit in the bits of
// the coordinate: at offset 0, the 381 bits below the flags.
func AddModulus(enc []byte, at int) ([]byte, error) {
	if at < 0 || at%CoordinateLen != 0 || at+CoordinateLen > len(enc) {
		return nil, errors.New("testkit: no coordinate at that offset")
	}
	c := bytes.Clone(enc[at : at+CoordinateLen])
	bits := 8 * CoordinateLen
	if at == 0 {
		c[0] &^= flagMask
		bits -= 3
	}
	sum := new(big.Int).Add(new(big.Int).SetBytes(c), FieldModulus)
	if sum.BitLen() > bits {
		return nil, errors.New("testkit: the coordinate plus p does not fit")
	}
	out := bytes.Clone(enc)
	sum.FillBytes(out[at : at+CoordinateLen])
	if at == 0 {
		out[0] |= enc[0] & flagMask
	}
	return out, nil
}

// ReduceCoordinate returns enc with the 48-byte coordinate at byte offset at
// reduced modulo p, the flags kept: what a decoder that reduces coordinates
// reads, and the inverse of AddModulus.
func ReduceCoordinate(enc []byte, at int) []byte {
	out := bytes.Clone(enc)
	c := out[at : at+CoordinateLen]
	flags := byte(0)
	if at == 0 {
		flags = c[0] & flagMask
		c[0] &^= flagMask
	}
	new(big.Int).Mod(new(big.Int).SetBytes(c), FieldModulus).FillBytes(c)
	c[0] |= flags
	return out
}

// Negated returns the encoding of the negated point: the same x, the sort
// flag flipped (spec §12.2). It is canonical when enc is the canonical
// encoding of a point other than the point at infinity.
func Negated(enc []byte) []byte {
	out := bytes.Clone(enc)
	out[0] ^= FlagSort
	return out
}

// InfinityWithPayload returns enc with the flags of the point at infinity,
// compression and infinity without sort, over its own coordinate bits: an
// encoding of the point at infinity with a payload, which spec §12.2
// forbids.
func InfinityWithPayload(enc []byte) []byte {
	out := bytes.Clone(enc)
	out[0] = out[0]&^flagMask | FlagCompressed | FlagInfinity
	return out
}

// Infinity returns the canonical encoding of the point at infinity in n
// bytes, 48 for G1 and 96 for G2: 0xc0, then zeros (spec §12.2).
func Infinity(n int) []byte {
	out := make([]byte, n)
	out[0] = FlagCompressed | FlagInfinity
	return out
}