#lang typed/racket
(require "ast.rkt")

(define-type IndentLevel Integer)

(define-type ISeq
  (∪ Null
      String
      (List 'Hang IndentLevel ISeq)
      'Newline
      (List '++ ISeq ISeq)
      (List 'Indent IndentLevel ISeq)))

(define-syntax iseq-append-with-sep
  (syntax-rules ()
    [(iseq-append-with-sep _) '()]
    [(iseq-append-with-sep _ a) a]
    [(iseq-append-with-sep sep a b c ...)
     (iseq-append a sep (iseq-append-with-sep sep b c ...) )]))

(: seq (-> ISeq ISeq ISeq))
(define (seq a b)
  (match (list a b)
    [`(() ,b) b]
    [`(,a ()) a]
    [_ (list '++ a b)]))

(define-syntax iseq-append
  (syntax-rules ()
    [(iseq-append a b) (seq a b)]
    [(iseq-append a b c ...) (seq a (iseq-append b c ...))]))

(: ppr-list (All (A) (-> (-> A ISeq) ISeq (Listof A) ISeq)))
(define (ppr-list f sep xs)
  (match xs
    ['() '()]
    [(list x) (f x)]
    [(list x y xs ...)
     (iseq-append (f x) sep (ppr-list f sep (cons y xs)))]))

(: ppr-alt (-> CoreAlt ISeq))
(define (ppr-alt a)
  (match a
    [(list con xs body)
     (iseq-append-with-sep
      " "
      (iseq-append
       "<"
       (number->string con)
       ">")
      (ppr-list symbol->string " " xs)
      "->"
      (ppr-expr body))]))

(: ppr-bind (-> CoreBind ISeq))
(define (ppr-bind a)
  (match a
    [(list x a)
     (iseq-append-with-sep " " (symbol->string x) "=" (ppr-expr a))]))

(: ppr-binds (-> CoreBinds ISeq))
(define (ppr-binds a)
  (ppr-list ppr-bind (iseq-append ";" 'Newline) a))

(: ppr-alts (-> CoreAlts ISeq))
(define (ppr-alts x) (ppr-list ppr-alt (iseq-append ";" 'Newline) x))

(: ppr-case (-> CoreExpr CoreAlts ISeq))
(define (ppr-case e alts)
  (iseq-append
   (iseq-append-with-sep
    " "
    "case"
    (ppr-expr e)
    "of")
   (list 'Hang -2
         (iseq-append 'Newline (ppr-alts alts)))))

(: ppr-lets (-> Let CoreBinds CoreExpr ISeq))
(define (ppr-lets lt binds e)
  (list 'Hang 0
        (iseq-append
         (symbol->string lt)
         " "
         (list 'Hang 0 (ppr-binds binds))
         (iseq-append 'Newline  "in " (ppr-expr e)))))

(: ppr-expr (-> CoreExpr ISeq))
(define (ppr-expr a)
  (match a
    [`(,f ,x) (iseq-append (ppr-expr f) " " (ppr-aexpr x))]
    [(? integer? n) (number->string n)]
    [(? symbol? x) (symbol->string x)]
    [`(case ,e ,alts) (ppr-case e alts)]
    [`(let ,binds ,e) (ppr-lets 'let binds e)]
    [`(letrec ,binds ,e) (ppr-lets 'letrec binds e)]))

(: ppr-aexpr (-> CoreExpr ISeq))
(define (ppr-aexpr a)
  (if (or (integer? a) (symbol? a))
      (ppr-expr a)
      (iseq-append "(" (ppr-expr a) ")")))

(: iseq->string (-> IndentLevel ISeq String))
(define (iseq->string l a)
  (let ([s (open-output-string)])
    (port-count-lines! s)
    (letrec ([go : (-> IndentLevel ISeq Void)
                 (λ (l i)
                   (match i
                     ['() (void)]
                     [`(++ ,i0 ,i1) (begin
                                      (go l i0)
                                      (go l i1))]
                     [`(Indent ,l0 ,i) (go (max (+ l l0) 0) i)]
                     [`(Hang ,l0 ,i)
                      (let-values ([(linum col-pos pos) (port-next-location s)])
                        (if (not (false? col-pos))
                            (go (max (+ col-pos l0) 0) i)
                            (error "impossible: line counting not enabled")))]
                     ['Newline (display
                                (string-append
                                 "\n"
                                 (make-string l #\space))
                                s)]
                     [(? string?) (display i s)]))])
      (go l a)
      (get-output-string s))))