package goxifrat

import (
	"crypto"
	"crypto/subtle"
	"strconv"
	//"crypto/sha512"
	"errors"
	"io"
)

const (
	P = 16
	N = 64
	R = 64
)

var st = [P][P]uint8{
	{7, 9, 13, 10, 15, 2, 0, 6, 3, 12, 8, 4, 1, 5, 14, 11},
	{1, 15, 6, 3, 9, 4, 11, 13, 10, 5, 14, 2, 7, 12, 8, 0},
	{3, 0, 12, 1, 11, 8, 9, 5, 7, 13, 2, 14, 10, 6, 4, 15},
	{4, 6, 15, 8, 13, 1, 5, 9, 14, 11, 10, 7, 2, 0, 3, 12},
	{0, 3, 8, 15, 10, 12, 7, 14, 9, 2, 13, 5, 11, 4, 6, 1},
	{10, 11, 5, 7, 0, 14, 15, 12, 1, 6, 4, 8, 3, 13, 2, 9},
	{5, 14, 10, 13, 8, 11, 4, 3, 6, 1, 15, 0, 12, 7, 9, 2},
	{15, 1, 4, 0, 7, 6, 10, 2, 11, 14, 5, 13, 9, 8, 12, 3},
	{12, 8, 3, 6, 14, 0, 2, 10, 13, 7, 9, 11, 5, 1, 15, 4},
	{13, 2, 7, 5, 4, 9, 8, 1, 12, 3, 0, 15, 6, 10, 11, 14},
	{6, 4, 1, 12, 2, 15, 14, 7, 5, 10, 11, 9, 13, 3, 0, 8},
	{9, 7, 2, 11, 1, 13, 3, 4, 0, 8, 12, 6, 15, 14, 5, 10},
	{11, 10, 14, 9, 3, 5, 1, 8, 15, 4, 6, 12, 0, 2, 13, 7},
	{14, 5, 11, 2, 12, 10, 6, 0, 4, 15, 1, 3, 8, 9, 7, 13},
	{8, 12, 0, 4, 5, 3, 13, 11, 2, 9, 7, 10, 14, 15, 1, 6},
	{2, 13, 9, 14, 6, 7, 12, 15, 8, 0, 3, 1, 4, 11, 10, 5},
}

func m(a, b [N]uint8) (r [N]uint8) {
	for j := 0; j < N; j++ {
		r[j] = st[a[j]][b[j]]
	}
	e := rn[0]
	for i := 0; i < R*N; i++ {
		d := rn[i+1]
		r[d] = st[r[d]][r[e]]
		e = d
	}
	for j := 0; j < N; j++ {
		r[j] = st[r[j]][b[j]]
	}
	return
}

type KeyAgreement struct {
	P    [N]uint8
	C    [N]uint8
	S1   [N]uint8
	S2   [N]uint8
	init bool
}

func NewKeyAgreement(random io.Reader) (k *KeyAgreement, err error) {
	p, err := randomArray(random)
	if err != nil {
		return nil, err
	}
	return &KeyAgreement{P: p}, nil
}

func (k *KeyAgreement) InitPhase(random io.Reader) (b []byte, err error) {
	k.C, err = randomArray(random)
	if err != nil {
		return
	}
	k.S1 = m(k.C, k.P)
	k.S2 = m(k.P, k.C)
	b = append(b, encode(k.C)...)
	b = append(b, encode(k.S1)...)
	k.init = true
	return
}

func (k *KeyAgreement) SharedSecret(buf []byte) (sec, b []byte, err error) {
	if !k.init {
		if l := len(buf); l != 64 {
			return nil, nil, errors.New("xifrat: bad buffer length: " + strconv.Itoa(l))
		}
		k.C = decode(buf[:32])
		buf = buf[32:]
		k.S1 = m(k.C, k.P)
		k.S2 = m(k.P, k.C)
		k.init = true
	}
	if l := len(buf); l != KeySize {
		return nil, nil, errors.New("xifrat: bad buffer length: " + strconv.Itoa(l))
	}
	b = encode(k.S1)
	sec = encode(m(decode(buf), k.S2))
	return
}

func encode(a [N]uint8) (buf []byte) {
	buf = make([]byte, len(a)/2)
	c := 0
	for i := 0; i < len(a); i++ {
		b := byte(a[i])
		i++
		buf[c] = (b << 4) | a[i]
		c++
	}
	return
}

const (
	KeySize = 32
	SigSize = 96
)

func decode(buf []byte) (a [N]uint8) {
	c := 0
	for _, b := range buf {
		a[c] = uint8(b >> 4)
		c++
		a[c] = uint8((b << 4) >> 4)
		c++
	}
	return
}

type PublicKey struct {
	C [N]uint8
}

type PrivateKey struct {
	PublicKey
	Q [N]uint8
}

func (priv *PrivateKey) Public() crypto.PublicKey {
	return &priv.PublicKey
}

func (priv *PrivateKey) Sign(random io.Reader, digest []byte) (sig []byte, err error) {
	if l := len(digest); l != KeySize {
		return nil, errors.New("xifrat: bad digest length: " + strconv.Itoa(l))
	}
	k, err := randomArray(random)
	if err != nil {
		return
	}
	hash := decode(digest)
	s := m(hash, priv.Q)
	ck := m(priv.PublicKey.C, k)
	qk := m(priv.Q, k)
	sig = append(sig, encode(s)...)
	sig = append(sig, encode(ck)...)
	sig = append(sig, encode(qk)...)
	return
}

func Verify(pub *PublicKey, digest, sig []byte) bool {
	if l := len(digest); l != KeySize {
		panic("xifrat: bad digest length: " + strconv.Itoa(l))
	}
	if l := len(sig); l != SigSize {
		panic("xifrat: bad signature length: " + strconv.Itoa(l))
	}
	s := decode(sig[:32])
	ck := decode(sig[32:64])
	qk := decode(sig[64:96])
	hash := decode(digest)
	v1 := m(s, ck)
	v2 := m(m(hash, pub.C), qk)
	return subtle.ConstantTimeCompare(encode(v1), encode(v2)) == 1
}

func randomArray(random io.Reader) (a [N]uint8, err error) {
	buf := make([]byte, 32)
	_, err = random.Read(buf)
	if err != nil {
		return
	}
	a = decode(buf)
	return
}

func GenerateKey(random io.Reader) (p *PrivateKey, err error) {
	pub, err := randomArray(random)
	if err != nil {
		return nil, err
	}
	priv, err := randomArray(random)
	if err != nil {
		return nil, err
	}
	return &PrivateKey{PublicKey{pub}, priv}, nil
}