/*
Copyright © 2025 filifa

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package cmd

import (
	"errors"
	"fmt"
	"log"
	"math/big"

	"github.com/spf13/cobra"
)

var modulus string
var tpf []string

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 computeNaive(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 computeFromFactors(modulus *big.Int, tpf []string) (*big.Int, error) {
	phi := big.NewInt(1)
	factors := make(map[string]bool)
	for _, s := range tpf {
		p, ok := new(big.Int).SetString(s, 10)
		if !ok {
			return nil, errors.New("invalid input " + s)
		}

		phi.Mul(phi, p)
		factors[p.Text(10)] = true
	}

	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
		}

		isPrimitive := true
		for p := range factors {
			e := new(big.Int)
			f, _ := new(big.Int).SetString(p, 10)
			k := new(big.Int).Div(phi, f)
			e.Exp(g, k, modulus)

			if e.Cmp(big.NewInt(1)) == 0 {
				isPrimitive = false
				break
			}
		}

		if isPrimitive {
			return g, nil
		}
	}

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

func primitiveRoot(cmd *cobra.Command, args []string) {
	m, ok := new(big.Int).SetString(modulus, 10)
	if !ok {
		log.Fatal("invalid input " + modulus)
	}

	root := new(big.Int)
	var err error
	if len(tpf) == 0 {
		root, err = computeNaive(m)
		if err != nil {
			log.Fatal(err)
		}
	} else {
		root, err = computeFromFactors(m, tpf)
		if err != nil {
			log.Fatal(err)
		}
	}

	fmt.Println(root)
}

// primitiveRootCmd represents the primitiveRoot command
var primitiveRootCmd = &cobra.Command{
	Use:   "primitive-root -m M",
	Short: "Compute a primitive root modulo n",
	Long: `Compute a primitive root modulo n.

This command computes a value g such that, for all integers a coprime to n, g^k = a (mod n) for some k. In other words, this command computes a generator for the multiplicative group of integers modulo n.

For improved performance, provide the prime factorization of the totient of n with the -t flag.

Note that primitive roots only exist for the moduli 1, 2, 4, p^k, and 2p^k, where p is an odd prime. The totients of these numbers are 1, 1, 2, (p-1)*p^(k-1), and (p-1)*p^(k-1), respectively.`,
	Run: primitiveRoot,
}

func init() {
	rootCmd.AddCommand(primitiveRootCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// primitiveRootCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// primitiveRootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")

	primitiveRootCmd.Flags().StringVarP(&modulus, "modulus", "m", "", "modulus")
	primitiveRootCmd.MarkFlagRequired("modulus")

	primitiveRootCmd.Flags().StringSliceVarP(&tpf, "totient-factorization", "t", make([]string, 0), "prime factorization of the totient of the modulus")
	// TODO: add a check flag for verifying -t input is the totient, test for performance
}