package lib

import (
	"errors"
	"math/big"
)

func Totient(n *big.Int) *big.Int {
	N := new(big.Int).Set(n)

	phi := new(big.Int).Set(N)

	sqrtn := new(big.Int).Sqrt(N)
	for i := big.NewInt(2); i.Cmp(sqrtn) != 1; i.Add(i, big.NewInt(1)) {
		mod := new(big.Int).Mod(N, i)
		if mod.Cmp(big.NewInt(0)) != 0 {
			continue
		}

		// phi -= phi // i
		tmp := new(big.Int).Div(phi, i)
		phi.Sub(phi, tmp)

		for mod.Cmp(big.NewInt(0)) == 0 {
			N.Div(N, i)
			mod.Mod(N, i)
		}
	}

	if N.Cmp(big.NewInt(1)) == 1 {
		// phi -= phi // N
		tmp := new(big.Int).Div(phi, N)
		phi.Sub(phi, tmp)
	}

	return phi
}

func MultiplicativeOrder(g *big.Int, modulus *big.Int) *big.Int {
	e := new(big.Int).Set(g)
	var k *big.Int
	for k = big.NewInt(1); e.Cmp(big.NewInt(1)) != 0; k.Add(k, big.NewInt(1)) {
		e.Mul(e, g)
		e.Mod(e, modulus)
	}

	return k
}

func PrimitiveRoot(modulus *big.Int) (*big.Int, error) {
	if modulus.Cmp(big.NewInt(1)) == 0 {
		return big.NewInt(0), nil
	}

	phi := Totient(modulus)

	for g := big.NewInt(1); g.Cmp(modulus) == -1; g.Add(g, big.NewInt(1)) {
		gcd := new(big.Int).GCD(nil, nil, g, modulus)
		if gcd.Cmp(big.NewInt(1)) != 0 {
			continue
		}

		order := MultiplicativeOrder(g, modulus)
		if order.Cmp(phi) == 0 {
			return g, nil
		}
	}

	return nil, errors.New("no primitive root")
}

func PrimitiveRootFast(modulus *big.Int, tpf map[string]*big.Int) (*big.Int, error) {
	phi := big.NewInt(1)
	for p, exp := range tpf {
		pow, ok := new(big.Int).SetString(p, 10)
		if !ok {
			return nil, errors.New("invalid factor " + p)
		}

		pow.Exp(pow, exp, nil)
		phi.Mul(phi, pow)
	}

	for g := big.NewInt(1); g.Cmp(modulus) == -1; g.Add(g, big.NewInt(1)) {
		gcd := new(big.Int).GCD(nil, nil, g, modulus)
		if gcd.Cmp(big.NewInt(1)) != 0 {
			continue
		}

		if isPrimitiveRoot(g, modulus, phi, tpf) {
			return g, nil
		}
	}

	return nil, errors.New("no primitive root")
}

func isPrimitiveRoot(g *big.Int, modulus *big.Int, phi *big.Int, tpf map[string]*big.Int) bool {
	for p := range tpf {
		// we already know factors are valid from computing phi
		k, _ := new(big.Int).SetString(p, 10)
		k.Div(phi, k)
		k.Exp(g, k, modulus)
		if k.Cmp(big.NewInt(1)) == 0 {
			return false
		}
	}

	return true
}