(define-module (language python string)
  #:use-module (oop goops)
  #:use-module (oop pf-objects)
  #:use-module (ice-9 match)
  #:use-module (language python list)
  #:use-module (language python exceptions)
  #:use-module (parser stis-parser)
  #:export (py-format py-capitalize py-center py-endswith
                      py-expandtabs py-find py-rfind
                      py-isalnum py-isalpha py-isdigit py-islower
                      py-isspace py-isupper py-istitle py-join py-ljust
                      py-rljust py-lower py-upper py-lstrip py-rstrip
                      py-partition py-replace py-strip py-title
                      py-rpartitio py-rindex py-split py-rsplit py-splitlines
                      py-startswith py-swapcase py-translate py-zfill))

(define-syntax-rule (aif it p x y) (let ((it p)) (if it x y)))

(define-syntax-rule (define-py (f n o . u) code ...)
  (begin
    (define-method (f (o <string>) . u) code ...)
    (define-method (f (o <p>) . l) (apply (ref o 'n) l))))

(define-py (py-capitalize capitalize s)
  (let* ((n (len s))
         (w (make-string n)))
    (let lp ((i 0) (first? #t))
      (if (< i n)
          (let ((ch (string-ref s i)))
            (if (and first? (char-alphabetic? ch))
                (begin
                  (string-set! w i (char-upcase ch))
                  (lp (+ i 1) #f))
                (begin
                  (string-set! w i ch)
                  (lp (+ i 1) first?))))
          w))))

(define-py (py-center center o w . l)
  (let* ((ws (if (pair? l)
                 (car (string->list (car l)))
                 #\space))
         (n  (string-length o))
         (w  (if (< w n) n w))
         (d  (- w n))
         (e  (floor-quotient (- w n) 2))
         (s  (make-string w #\space)))
    (let lp ((i 0) (j e))
      (if (< i n)
          (begin
            (string-set! s j (string-ref o i))
            (lp (+ i 1) (+ j 1)))))
    s))
    
    

;;;py-decode
;;;py-encode


(define-py (py-endswith endswith o (suff <string>) . l)
  (let* ((n   (string-length o))
         (ns  (string-length suff))
         (f   (lambda (x) (< x 0) (+ n x) x)))
    (call-with-values (lambda ()
                        (match l
                          (()    (values 0      n  ))
                          ((x)   (values (f x)  n  ))
                          ((x y) (values (f x) (f y)))))
      (lambda (start end)
        (string-suffix? suff o 0 ns start end)))))

(define-py (py-startswith endswith o (suff <string>) . l)
  (let* ((n   (string-length o))
         (ns  (string-length suff))
         (f   (lambda (x) (< x 0) (+ n x) x)))
    (call-with-values (lambda ()
                        (match l
                          (()    (values 0      n  ))
                          ((x)   (values (f x)  n  ))
                          ((x y) (values (f x) (f y)))))
      (lambda (start end)
        (string-prefix? suff o 0 ns start end)))))

(define-py (py-expandtabs expandtabs s . l)
  (let* ((tabsize (match l (() 8) ((x) x)))
         (u       (string->list (make-string tabsize #\space)))
         (n       (string-length s)))
    (let lp ((l (string->list s)) (r '()))
      (if (pair? l)
          (let ((x (car l)))
            (if (eq? x #\tab)
                (lp (cdr l) (append u r))
                (lp (cdr l) (cons x r))))
          (list->string (reverse r))))))
      
(define-py (py-find find s sub . l)
  (let* ((n   (string-length s))
         (f   (lambda (x) (< x 0) (+ n x) x)))
    (call-with-values (lambda ()
                        (match l
                          (()    (values 0      n  ))
                          ((x)   (values (f x)  n  ))
                          ((x y) (values (f x) (f y)))))
      (lambda (start end)
        (aif it (string-contains s sub start end)
             it
             -1)))))

(define-py (py-rfind rfind s sub . l)  
  (let* ((n   (string-length s))
         (s   (string-reverse s))
         (sub (string-reverse sub))
         (f   (lambda (x) (< x 0) (+ n x) x)))
    (call-with-values (lambda ()
                        (match l
                          (()    (values 0      n  ))
                          ((x)   (values (f x)  n  ))
                          ((x y) (values (f x) (f y)))))
      (lambda (start end)
        (aif it (string-contains s sub start end)
             (- n it (len sub))
             -1)))))

(define i       (f-list #:i (mk-token (f+ (f-reg! "[0-9]")))))
(define s       (f-list #:s (mk-token (f+ (f-not! (f-tag "}"))))))
(define e       (f-list #:e (f-and (f-tag "}") f-true)))
(define tagbody (f-or! e i s))

(define tag     (f-seq "{" tagbody "}"))
(define nontag  (f-list #:str (mk-token (f+  (f-or! (f-tag "{{") 
                                                    (f-not! tag))))))
(define e       (ff* (f-or! tag nontag)))

(define (compile x args kwargs)
  (let lp ((l x) (r '()) (u '()) (i 0))
    (match l
      (((#:str x) . l)
       (lp l (cons x r) u i))
      (((#:i x)   . l)
       (lp l (cons "~a" r) (cons (list-ref args (string->number x)) u) i))
      (((#:s x)   . l)
       (lp l (cons "~a" r) (cons (hash-ref kwargs x None) u) i))
      (((#:e)     . l)
       (lp l (cons "~a" r) (cons (list-ref args i) u) (+ i 1)))
      (()
       (apply format #f (string-join (reverse r) "") (reverse u))))))

(define-py (py-format format s . l)
  (call-with-values
      (lambda ()
        (let lp ((l l) (args '()) (kwargs (make-hash-table)))
          (match l
            (((? keyword? key) x . l)
             (hash-set! kwargs (symbol->string (keyword->symbol key)) x)
             (lp l args kwargs))
            ((x . l)
             (lp l (cons x args) kwargs))
            (()
             (values (reverse args) kwargs)))))
    (lambda (args kwargs)
      (compile (parse s e) args kwargs))))

(define-syntax-rule (mk-is py-isalnum isalnum x ...)
  (define-py (py-isalnum isalnum s)
    (and (> (len s) 0)
         (string-fold
          (lambda (ch s)
            (if (or (x ch) ...)
                s
                #f))
          #t s))))

(mk-is py-isalnum isalnum char-alphabetic? char-numeric?)
(mk-is py-isalpha isalpha char-alphabetic?)
(mk-is py-isdigit isdigit char-numeric?)
(mk-is py-islower islower char-lower-case?)
(mk-is py-isspace isspace char-whitespace?)
(mk-is py-isupper isupper char-upper-case?)

(define-py (py-istitle istitle s)
  (let ((n (len s)))
    (if ((> n 0))
        (let lp ((i 0) (space? #t))
          (if (< i n)
              (let ((ch (string-ref s i)))
                (if space?
                    (cond
                     ((char-whitespace? ch)
                      (lp (+ i 1) #t))
                     ((char-upper-case? ch)
                      (lp (+ i 1) #f))
                     (else
                      #f))
                    (cond
                     ((char-whitespace? ch)
                      (lp (+ i 1) #t))
                     ((char-upper-case? ch)
                      #f)
                     ((char-lower-case? ch)
                      (lp (+ i 1) #f))
                     (else
                      #f))))
              #t))
        #f)))


(define-py (py-join join s iterator)
  (string-join (to-list iterator) s))

(define-py (py-ljust ljust s width . l)
  (let* ((n  (len s))
         (ch (match l
               ((x . l)
                (if (string? x)
                    (string-ref x 0)
                    x))
                (()
                 #\space))))
    (if (< width n)
        (pylist-slice s 0 width)
        (let ((ret (make-string width ch)))
          (let lp ((i 0))
            (if (< i n)
                (string-set! ret i (string-ref s i))
                ret))))))

(define-py (py-rjust rjust s width . l)
  (let* ((n  (len s))
         (ch (match l
               ((x . l)
                (if (string? x)
                    (string-ref x 0)
                    x))
               (()
                #\space))))
    (if (< width n)
        (pylist-slice s (- width) (len s))
        (let ((ret (make-string width ch)))
          (let lp ((i 0) (j (- width n)))
            (if (< i n)
                (string-set! ret j (string-ref s i))
                ret))))))

(define-py (py-lower lower s)
  (string-downcase s))

(define-py (py-upper upper s)
  (string-upcase s))

(define-py (py-lstrip lstrip s . l)
  (match l
    (()
     (string-trim s))
    ((x . _)
     (let ((l (map (lambda (x) (if (string? x) (string-ref x 0) x)) x)))
       (string-trim s (lambda (ch) (member ch l)))))))

(define-py (py-rstrip rstrip s . l)
  (match l
    (()
     (string-trim-right s))
    ((x . _)
     (let ((l (map (lambda (x) (if (string? x) (string-ref x 0) x)) x)))
       (string-trim-right s (lambda (ch) (member ch l)))))))

(define-py (py-partition partition s (sep <string>))
  (let ((n (len s))
        (m (len sep)))
    (define (test i)
      (let lp ((i i) (j 0))
        (if (< i n)
            (if (< j m)
                (if (eq? (string-ref s i) (string-ref sep j))
                    (lp (+ i 1) (+ j 1))
                    #f)
                #t)
            #f)))
    (let lp ((i 0))
      (if (< i n)
          (if (test i)
              (list (pylist-slice s 0 i) sep (pylist-slice s (+ i m) n))
              (lp (+ i 1)))
          (list s "" "")))))

(define-py (py-rpartition rpartition ss (ssep <string>))
  (let* ((s    (string-reverse ss))
         (sep  (string-reverse ssep))
         (n    (len s))
         (m    (len sep)))
    (define (test i)
      (let lp ((i i) (j 0))
        (if (< i n)
            (if (< j m)
                (if (eq? (string-ref s i) (string-ref sep j))
                    (lp (+ i 1) (+ j 1))
                    #f)
                #t)
            #f)))
    (let lp ((i 0))
      (if (< i n)
          (if (test i)
              (list (string-reverse
                     (pylist-slice s (+ i m) n))
                    ssep
                    (string-reverse
                     (pylist-slice s 0 i)))
              (lp (+ i 1)))
          (list "" "" s)))))
    
(define-py (py-replace replace s old new . l)
  (let ((n (match l (() #f) ((n . _) n))))
    (string-join
     (reverse
      (let lp ((s s) (r '()))
        (let ((l (py-partition s old)))
          (if (equal? (cadr l) "")
              (cons s r)
              (lp (list-ref l 2) (cons (car l) r))))))
     new)))

(define-py (py-strip strip s . l)
  (apply py-rstrip (apply py-lstrip s l) l))

(define-py (py-title title s)
  (string-titlecase s))

(define-py (py-rindex rindex s . l) 
  (let ((n (len s)))
    (- n (apply pylist-index (string-reverse s) l) 1)))



(define-py (py-split split s . l)
  (define ws (f+ (f-reg "[ \t\n]")))
  (define r
    (f-or! (f-seq f-eof (f-out '()))
           (f-cons (f-seq (mk-token (f* (f-reg! "."))) f-eof) (f-out '()))))
  (define (u ws)  (mk-token (f+ (f-not! ws))))
  (define (tok ws i)
    (if (= i 0)
        (f-list (mk-token (f* (f-reg! "."))))
        (let ((e (mk-token (f* (f-not! ws)))))
          (f-seq (f? ws)
                 (f-cons e
                         (let lp ((i i))
                           (if (> (- i 1) 0)
                               (f-or! (f-seq (f? ws) f-eof (f-out '()))
                                      (f-cons (f-seq ws e) (Ds (lp (- i 1)))))
                               r)))))))
  
  (define N 1000000000000)
  (let ((e (call-with-values
               (lambda ()
                 (match l
                   (()      (values ws          N))
                   ((sep)   (values (f-tag sep) N))
                   ((sep n) (values (f-tag sep) n))))
             tok)))
    (parse s e)))

(define-py (py-rsplit rsplit s . l)
  (reverse
   (map string-reverse
        (apply py-split
               (string-reverse s)
               (match l
                 (() '())
                 ((sep . l) (cons (string-reverse sep) l)))))))


(define-py (py-splitlines splitlines s . l)
  (let ((n     (len s))
        (keep? (match l
                 ((#:keepends v)
                  v)
                 ((v)
                  v)
                 (_ #f))))
    (let lp ((i 0) (r '()) (old 0))
      (if (< i n)
          (let ((ch (string-ref s i)))
            (if (eq? ch #\newline)
                (if keep?
                    (lp (+ i 1)
                        (cons
                         (pylist-slice s old (+ i 1) 1)
                         r)
                        (+ i 1))
                    (lp (+ i 1)
                        (cons
                         (pylist-slice s old i 1)
                         r)
                        (+ i 1)))
                (lp (+ i 1) r old)))
          (reverse r)))))

(define-py (py-swapcase swapcase s)
  (list->string
   (string-fold
    (lambda (ch s)
      (cons
       (cond
        ((char-upper-case? ch)
         (char-downcase ch))
        ((char-lower-case? ch)
         (char-upcase ch))
        (else ch))
       s))
    '()
    s)))
        
(define-py (py-translate translate s table . l)
  (let* ((n (len s))
         (w (make-string n))
         (t (if (eq? table None) #f table))
         (d (match l (() #f) ((x) x))))
    (define (tr ch)
      (define (e)
        (if t
            (let ((i (char->integer ch)))
              (if (< i (string-length t))
                  (string-ref t i)
                  ch))
            ch))

      (if d
          (if (string-contains d (list->string (list ch)))
              #f
              (e))
          (e)))
    
    (let lp ((i 0) (k 0))
      (if (< i n)
          (let ((ch (tr (string-ref s i))))
            (if ch
                (begin
                  (string-set! w k ch)
                  (lp (+ i 1) (+ k 1)))
                (lp (+ i 1) k)))
          (if (= k n)
              w
              (pylist-slice w 0 k 1))))))
                
(define-py (py-zfill zfill s width)
  (let* ((n (len s))
         (w (pk (pylist-slice s 0 n 1))))
    (let lp ((i 0))
      (if (< i n)
          (let ((ch (string-ref s i)))
            (if (char-numeric? ch)
                (let lp ((j (max 0 (- i width))))
                  (pk i j)
                  (if (< j i)
                      (begin
                        (string-set! w j #\0)
                        (lp (+ j 1)))
                      w))
                (lp (+ i 1))))
          s))))