package integer

/*
RSAInt: uses big ints to handle big numbers,
and adapts usage to a more fuctional level
against encryption-keys
*/

import (
    "fmt"
    "time"
    "math/big"
    "math/rand"
)

type inttype int

const (
    ANY inttype = 1 << iota
    ODD inttype = 1 << iota
    EVEN inttype = 1 << iota
)

var (
    ZERO = New(0).setProtected(true)
    ONE = New(1).setProtected(true)
    TWO= New(2).setProtected(true)
    THREE = New(3).setProtected(true)
)


type RSAInt struct {
    *big.Int
    p bool
}

func randInt() *rand.Rand {
    return rand.New(rand.NewSource(time.Now().UnixNano()))
}

func New(i int64) *RSAInt {
    k := new(RSAInt)
    k.Int = big.NewInt(i)
    return k.setProtected(false)
}

func (k *RSAInt) Copy() *RSAInt {
    return New(0).Set(k).setProtected(k.p)
}

func (k *RSAInt) IsProtected() bool {
    return k.p
}

func (k *RSAInt) isProtected() bool {
    if k.IsProtected() {
        fmt.Printf("Using protected value %p: %s\n", k, k.Text(10))
        return true
    }
    return false
}

func (k *RSAInt) setProtected(p bool) *RSAInt {
    k.p = p
    return k
}

func (k *RSAInt) Set(o *RSAInt) *RSAInt {
    k.isProtected()
    k.Int.Set(o.Int)
    return k
}

func (k *RSAInt) Set64(i int64) *RSAInt {
    k.isProtected()
    k.Int.SetInt64(i)
    return k
}

func (k *RSAInt) SetString(s string) (*RSAInt, bool) {
    k.isProtected()
    _, e := k.Int.SetString(s, 10)
    return k, e
}

func (k *RSAInt) Next(s int64) *RSAInt {
    k.isProtected()
    return k.Add(New(s))
}

func (k *RSAInt) Add(b *RSAInt) *RSAInt {
    k.isProtected()
    k.Int.Add(k.Int, b.Int)
    return k
}

func (k *RSAInt) Sub(b *RSAInt) *RSAInt {
    k.isProtected()
    k.Int.Sub(k.Int, b.Int)
    return k
}

func (k *RSAInt) Mul(b *RSAInt) *RSAInt {
    k.isProtected()
    k.Int.Mul(k.Int, b.Int)
    return k
}

func (k *RSAInt) Div(b *RSAInt) *RSAInt {
    k.isProtected()
    k.Int.Div(k.Int, b.Int)
    return k
}

func (k *RSAInt) Mod(b *RSAInt) *RSAInt {
    k.isProtected()
    k.Int.Mod(k.Int, b.Int)
    return k
}

func (k *RSAInt) Pow(e *RSAInt) *RSAInt {
    k.isProtected()
    k.Exp(k.Int, e.Int, nil)
    return k
}

func (k *RSAInt) Pow64(e int64) *RSAInt {
    k.isProtected()
    k.Exp(k.Int, New(e).Int, nil)
    return k
}

func (k *RSAInt) PowMod(e, m *RSAInt) *RSAInt {
    k.isProtected()
    k.Exp(k.Int, e.Int, m.Int)
    return k
}

func (k *RSAInt) GCD(b *RSAInt) *RSAInt {
    k.isProtected()
    k.Int.GCD(nil, nil, k.Int, b.Int)
    return k
}

func (k *RSAInt) Rand() *RSAInt {
    k.isProtected()
    k.Int.Rand(randInt(), k.Int)
    return k
}
func (k *RSAInt) Eq(ot *RSAInt) bool {
    return k.Cmp(ot.Int) == 0
}

func (k *RSAInt) Neq(ot *RSAInt) bool {
    return !k.Eq(ot)
}

func (k *RSAInt) Lt(ot *RSAInt) bool {
    return k.Cmp(ot.Int) < 0
}

func (k *RSAInt) LtEq(ot *RSAInt) bool {
    return k.Cmp(ot.Int) <= 0
}

func (k *RSAInt) Gt(ot *RSAInt) bool {
    return !k.LtEq(ot)
}

func (k *RSAInt) GtEq(ot *RSAInt) bool {
    return !k.Lt(ot)
}

func (k *RSAInt) Or(ot *RSAInt) *RSAInt {
    k.Int.Or(k.Int, ot.Int)
    return k
}

func (k *RSAInt) Len() int {
    return len(k.String())
}

func (k *RSAInt) IsPrime() bool {
    return k.ProbablyPrime(100)
}

func (k *RSAInt) isPrime(p chan *RSAInt) {
    if k.IsPrime() {
        p <- k
    }
}

func (k *RSAInt) RandRange(r inttype) *RSAInt {
    s, e := "", New(0)

    for l, i := k.Len(), 0; i < l; i++ {
        s += "0"
    }

    if _, r := e.SetString(s); !r {
        return nil
    }

    e.Sub(k).Rand().Add(k)

    if c := e.Copy().Mod(TWO); r == EVEN && c.Neq(ZERO) {
        e.Add(ONE)
    } else if r == ODD && c.Eq(ZERO) {
        e.Add(ONE)
    }
    return e
}

func (k *RSAInt) NextRandPrime() {
    p := make(chan *RSAInt)
    var r *RSAInt
    for run, e := true, k.RandRange(ODD); run; e.Next(2) {
        select {
            case r = <-p:
                run = false
                *k = *r
            default:
                go e.Copy().isPrime(p)
        }
    }
}

func (k *RSAInt) PollardPm1(base int64) (*RSAInt, *RSAInt, *RSAInt) {
    n, b, q, p := New(0), New(base), New(0), New(base).GCD(k)

    for f, i := New(1), New(1); p.Eq(ONE) && n.Lt(k); n.Next(1) {
        f.Mul(i)
        p = b.Copy().PowMod(f, k).Sub(ONE).GCD(k)
        i.Next(1)
    }

    q.Set(k).Div(p)

    fmt.Printf(
        "POLLARD: n: %s\tb: %s\tp: %s\tq: %s\n",
        n.Set(p).Mul(q), b.Text(10), p.Text(10), q.Text(10),
    )

    return n, p, q
}

func (k *RSAInt) EulerTot(p, q *RSAInt) *RSAInt {
    return k.Add(New(1)).Sub(p.Copy().Add(q))
}

func (k *RSAInt) RelativePrime() *RSAInt {
    for x := New(65537); x.Lt(k); x.Next(1) {
        if x.Copy().GCD(k).Eq(ONE) {
            return k.Set(x)
        }
    }
    return nil
}

func (k *RSAInt) ModMulInv(e, etot *RSAInt) *RSAInt {
    r := etot.Copy()
    nr := e.Copy().Mod(etot)

    for t, nt := New(0), New(1); nr.Neq(ZERO); {
        q := r.Copy().Div(nr).Or(ZERO)
        nt, t = t.Sub(q.Copy().Mul(nt)), nt
        nr, r = r.Sub(q.Copy().Mul(nr)), nr
        k.Set(t)
    }

    if r.Gt(ONE) {
        k.SetInt64(-1)
    } else if k.Lt(ZERO) {
        k.Add(etot)
    }

    return k
}
