(import srfi-1 (math base) (math number-theory) srfi-4 (chicken bitwise))

(define (make-polynom x0 coefficients)
  (lambda (x)
    (apply + `(,x0 ,@(map (lambda (c i) (* c (expt x i))) coefficients (iota (length coefficients) 1))))))

(define (make-polynom-in-field x0 coefficients p)
  (let ((f (make-polynom x0 coefficients)))
    (lambda (x) (modulo (f x) p))))


(define (split-secret secret shares needed-for-recovery)
  (let* ((number-of-coefficients (- shares needed-for-recovery))
         (p (next-prime secret))
         (cs (list-tabulate number-of-coefficients (lambda (i) (random-bits 128))))
         (f (make-polynom-in-field secret cs p)))
    (list `(n ,shares) `(k ,needed-for-recovery) `(p ,p) `(shares
                                                           ,(list-tabulate shares (lambda (i) (cons (add1 i) (f (add1 i)))))))))

(define (recover-secret p shares)
  (modulo (fold (lambda (j acc) (+ acc (* (cdr j)
                               (fold (lambda (m acc)
                                       (if (not (equal? m j))
                                           (* acc (/ (car m) (- (car m) (car j))))
                                           acc))
                                     1 shares))))
          0 shares)
          p))

(define (test-it)
  (let* ((secret 123456)
         (stuff (split-secret secret 5 3))
         (p (car (alist-ref 'p stuff)))
         (prt (car (alist-ref 'shares stuff)))
         (secret_restored (recover-secret p prt)))
    (eq? secret_restored secret)))


;; These vectors are in reverse network byte ordering (LSB)
(define (u8vector->number v)
  (let loop ((i (u8vector-length v))
             (num 0))
    (if (zero? i)
        num
        (loop (sub1 i)
              (+ (arithmetic-shift num 8)
                 (u8vector-ref v (sub1 i)))))))



(define (number->u8vector number)
  (define (determine-vector-length number)
    (let loop ((n number)
               (i 1))
      (if (zero? (quotient n 256))
          i
          (loop (quotient n 256) (add1 i)))))
  (let* ((size (determine-vector-length number))
         (v (make-u8vector size 0)))
    (let loop ((i size)
               (num number))
      (if (zero? i)
          (if (zero? num)
              v
              ;; Internal error
              (error (sprintf "Number too large: ~A can't be split into an u8vector of ~A entries" number size)))
          (begin
            (u8vector-set! v (- size i) (inexact->exact (modulo num 256))) ; XXX
            (loop (sub1 i) (quotient num 256)))))))

(define (string->secret str)
  (u8vector->number (list->u8vector (map char->integer (string->list str)))))

(define (secret->string s)
  (list->string (map integer->char (u8vector->list (number->u8vector s)))))

(define (test-it-again)
  (let* ((secret "my secret passphrase, that I want to share with my friends")
         (s (string->secret secret))
         (shares (split-secret s 5 3)))
    (print "Length secret: " (string-length secret) " s: " s)
    (and (equal? s (recover-secret (car (alist-ref 'p shares)) (car (alist-ref 'shares shares))))
         (secret->string (recover-secret (car (alist-ref 'p shares)) (car (alist-ref 'shares shares)))))))