(define-module (language python string)
  #:use-module (oop goops)
  #:use-module (oop pf-objects)
  #:use-module (language python hash)
  #:use-module (ice-9 match)
  #:use-module (language python list)
  #:use-module (language python exceptions)
  #:use-module (language python for)
  #: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
                      pystring-listing <py-string> pystring))

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


(define-class <py-string> () str)

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

(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 startswith 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-syntax-rule (a b x y) (b (symbol->string x) (symbol->string y)))

(define-syntax-rule (mkop op)
 (begin
   (define-method (op (s1 <string>) (s2 <py-string>))
     (op s1 (slot-ref s2 'str)))
   (define-method (op (s2 <py-string>) (s1 <string>))
     (op s1 (slot-ref s2 'str)))))

(mkop <)
(mkop <=)
(mkop >)
(mkop >=)
(mkop +)
(mkop *)

(define-method (<  (s1 <string>) (s2 <string>)) (string-ci<  s1 s2))
(define-method (<= (s1 <string>) (s2 <string>)) (string-ci<= s1 s2))
(define-method (>  (s1 <string>) (s2 <string>)) (string-ci>  s1 s2))
(define-method (>= (s1 <string>) (s2 <string>)) (string-ci>= s1 s2))

(define-method (<  (s1 <symbol>) (s2 <symbol>)) (a string-ci<  s1 s2))
(define-method (<= (s1 <symbol>) (s2 <symbol>)) (a string-ci<= s1 s2))
(define-method (>  (s1 <symbol>) (s2 <symbol>)) (a string-ci>  s1 s2))
(define-method (>= (s1 <symbol>) (s2 <symbol>)) (a string-ci>= s1 s2))


(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))))

(define-python-class string (<py-string>)
  (define __init__
    (case-lambda
      ((self s)
       (cond
        ((is-a? s <py-string>)
         (slot-set! self 'str (slot-ref s 'src)))
        ((is-a? s <string>)
         (slot-set! self 'str s)))))))

(define pystring string)

(define-method (py-class (o <string>))    string)
(define-method (py-class (o <py-string>)) string)

(define-method (pyhash (o <py-string>)) (hash (slot-ref o 'str) pyhash-N))

(define-method (py-equal? (o <py-string>) x)
  (equal? (slot-ref o 'str) x))
(define-method (py-equal? x (o <py-string>))
  (equal? (slot-ref o 'str) x))

(define-class <string-iter> (<py-string>) str i d)

(define-method (write (o <string-iter>) . l)
  (define port (if (null? l) #t (car l)))
  (for ((x : o)) ((l '()))
       (cons (string-ref x 0) l)
       #:final
       (format port "iter(~s)" (list->string (reverse l)))))

(define-method (wrap-in (o <string-iter> ))
  (let ((out (make <string-iter>)))
    (slot-set! out 'str (slot-ref o 'str))
    (slot-set! out 'i   (slot-ref o 'i))
    (slot-set! out 'd   (slot-ref o 'd))
    out))

(define-method (wrap-in (s <string>))
  (let ((out (make <string-iter>)))
    (slot-set! out 'str s)
    (slot-set! out 'i   0)
    (slot-set! out 'd   1)
    out))

(define-method (py-reversed (s <string>))
  (let ((out (make <string-iter>)))
    (slot-set! out 'str s)
    (slot-set! out 'i   (- (string-length s) 1))
    (slot-set! out 'd   -1)
    out))

(define-method (next (o <string-iter>))
  (let ((i   (slot-ref o 'i  ))
        (d   (slot-ref o 'd))
        (str (slot-ref o 'str)))
    (if (> d 0)
        (if (< i (string-length str))
            (let ((ret (string-ref str i)))
              (slot-set! o 'i (+ i d))
              (list->string (list ret)))
            (throw StopIteration))
        (if (>= i 0)
            (let ((ret (string-ref str i)))
              (slot-set! o 'i (+ i d))
              (list->string (list ret)))
            (throw StopIteration)))))

(define (pystring-listing)
  (let ((l (to-pylist
            (map symbol->string
                 '(__add__ __class__ __contains__ __delattr__ __doc__
                           __eq__ __format__ __ge__ __getattribute__
                           __getitem__ __getnewargs__ __getslice__ __gt__
                           __hash__ __init__ __le__ __len__ __lt__ __mod__
                           __mul__ __ne__ __new__ __reduce__ __reduce_ex__
                           __repr__ __rmod__ __rmul__ __setattr__ __sizeof__
                           __str__ __subclasshook__
                           _formatter_field_name_split _formatter_parser
                           capitalize center count decode encode endswith
                           expandtabs find format index isalnum isalpha
                           isdigit islower isspace istitle isupper join
                           ljust lower lstrip partition replace rfind rindex
                           rjust rpartition rsplit rstrip split splitlines
                           startswith strip swapcase
                           title translate upper zfill)))))
    (pylist-sort! l)
    l))