#lang racket/base
(require (prefix-in etc: mzlib/etc)
         texpict/mrpict
         (only-in pict pin-line pin-arrow-line)
         (except-in texpict/utils pin-line pin-arrow-line)
         racket/class
         racket/runtime-path
         racket/draw
         racket/contract
         (only-in racket/list last))

(define the-font-size 12)
(define prim-font-family 'swiss)
;; Was 'modern, but we want font smoothing even for small text:
(define font-family (make-object font% the-font-size 'modern 'normal 'normal #f 'smoothed))

(define prim-types '("int" "String" "float" "double" "boolean"))

;; how far a dot is to the right of a class
(define dot-edge-spacing 10)

(define field-arrowhead-size 10)

(define hierarchy-color "navy")
(define type-link-color "firebrick")

#|
(define font-family "Palatino")
(define-runtime-path afm "afm")
(current-ps-afm-file-paths (cons afm (current-ps-afm-file-paths)))

(let ([id (send the-font-name-directory find-or-create-font-id font-family 'default)])
  (send the-font-name-directory set-post-script-name 
        id 'normal 'normal "Palatino-Roman")
  (send the-font-name-directory set-post-script-name 
        id 'bold 'normal "Palatino-Bold")
  (send the-font-name-directory set-post-script-name 
        id 'normal 'italic "Palatino-Italic"))
|#

(define (user-type-font x) (text x font-family the-font-size))
(define (prim-type-font x) (text x  prim-font-family the-font-size))
(define (var-font x) (text x `(bold . ,font-family) the-font-size))
(define (field-name-font x) (text x font-family the-font-size))
(define (comment-font x) (text x font-family the-font-size))
(define (normal-font x) (text x font-family the-font-size))
(define (java-this) (text "this" `(bold . ,font-family) the-font-size))

;; field-spec : string string -> pict
(define (field-spec type fd #:default [default #f] [comment #f])
  (define code-line
    (hbl-append (if type
                    (hbl-append (type-spec type) (normal-font " "))
                    (blank))
                (field-name-font fd)
                (if default
                    (hbl-append (normal-font " = ") (normal-font default))
                    (blank))
                #;(normal-font ";")))
  (if comment
      (hbl-append code-line (normal-font " ") (comment-font (format "[in ~a]" comment)))
      code-line))

(define (method-spec range name #:body [body #f] . args)
  (unless (even? (length args))
    (error 'method-spec "expected a list of types and argument names, but found ~a arguments"
           (length args)))
  (define first-line
    (hbl-append
     (type-spec range)
     (normal-font " ")
     (var-font name)
     (cond
       [(null? args) (normal-font "()")]
       [else
        (hbl-append
         (normal-font "(")
         (let loop ([args args])
           (define type (car args))
           (define param (cadr args))
           (define single-arg
             (if param
                 (hbl-append (type-spec type) (normal-font " ") (var-font param))
                 (type-spec type)))
         
           (cond
             [(null? (cddr args)) (hbl-append single-arg (normal-font ")"))]
             [else (hbl-append single-arg (normal-font ", ") (loop (cddr args)))])))])
     (if body
         (hbl-append (normal-font " {"))
         (blank))))
  (if body
      (vl-append first-line (hbl-append (blank 8 0) body (normal-font "}")))
      first-line))
             
(define (type-spec str)
  (cond
    [(member str prim-types) (prim-type-font str)]
    [else (user-type-font str)]))

;; class-name : string -> pict
(define (class-name txt #:spacing-word [spacing-word txt]) 
  (define p (colorize (lt-superimpose (ghost (var-font spacing-word))
                                      (apply vl-append (map var-font (regexp-split #rx"\n" txt))))
                      "white"))
  (refocus (cc-superimpose (colorize (filled-rectangle (+ class-box-margin class-box-margin (pict-width p))
                                                       (+ class-box-margin class-box-margin (pict-height p)))
                                     "black")
                           p)
           p))

(define class-box-margin 4)

;; class-box : pict (or/c #f (listof pict)) (or/c #f (listof pict)) -> pict
(define (class-box name fields methods)
  (define (mk-blank)
    (blank 0 (+ class-box-margin class-box-margin)))
  (cond
    [(and methods fields)
     (let* ([top-spacer (mk-blank)]
            [bottom-spacer (mk-blank)]
            [main (vl-append name
                             top-spacer
                             (if (null? fields)
                                 (blank 0 4)
                                 (apply vl-append fields))
                             bottom-spacer
                             (if (null? methods)
                                 (blank 0 4)
                                 (apply vl-append methods)))])
       (add-hline (add-hline (frame (inset main class-box-margin)) top-spacer) bottom-spacer))]
    [fields
     (define top-spacer (mk-blank))
     (define main
       (vl-append name
                  top-spacer
                  (if (null? fields)
                      (blank)
                      (apply vl-append fields))))
     (add-hline (frame (inset main class-box-margin)) top-spacer)]
    [methods (class-box name methods fields)]
    [else (frame (inset name class-box-margin))]))

(define (add-hline main sub)
  (define-values (x y) (cc-find main sub))
  (pin-line main sub (λ (p1 p2) (values 0 y)) sub (λ (p1 p2) (values (pict-width main) y))))

;; hierarchy : pict (cons pict (listof pict)) (cons pict (listof pict)) -> pict
(define (hierarchy main supers subs)
  (define supers-bottoms
    (apply max
           (for/list ([x (in-list supers)])
             (let-values ([(x y) (cb-find main x)])
               y))))
  (define subs-tops
    (apply min
           (for/list ([x (in-list subs)])
             (let-values ([(x y) (ct-find main x)])
               y))))
  (define sorted-subs (sort subs (λ (x y) (< (left-edge-x main x) (left-edge-x main y)))))
  (unless (< supers-bottoms subs-tops)
    (error 'hierarchy
           "expected supers to be on top of subs, supers bottom is at ~a, and subs tops is at ~a"
           supers-bottoms
           subs-tops))
  (define main-line-y (max (- subs-tops 20) (/ (+ supers-bottoms subs-tops) 2)))
  (define main-line-start-x (center-x main (car sorted-subs)))
  (define main-line-end-x (center-x main (last sorted-subs)))
  (define w/main-line
    (pin-line main
              main
              (λ (_1 _2) (values main-line-start-x main-line-y))
              main
              (λ (_1 _2) (values main-line-end-x main-line-y))
              #:color hierarchy-color))
  (define super-lines
    (for/list ([super (in-list supers)])
      (define-values (x y) (cb-find main super))
      (pin-over (pin-line (ghost main) super cb-find main (λ (_1 _2) (values x main-line-y)))
                (- x (/ (pict-width triangle) 2))
                (- (/ (+ y main-line-y) 2) (/ (pict-height triangle) 2))
                triangle)))
  (define sub-lines
    (for/list ([sub (in-list subs)])
      (define-values (x y) (ct-find main sub))
      (pin-line (ghost main)
                sub
                ct-find
                main
                (λ (_1 _2) (values x main-line-y))
                #:color hierarchy-color)))
  (apply cc-superimpose w/main-line (append sub-lines super-lines)))

(define triangle-width 12)
(define triangle-height 12)
(define triangle
  (let ([points (list (make-object point% (/ triangle-width 2) 0)
                      (make-object point% 0 triangle-height)
                      (make-object point% triangle-width triangle-height))])
    (colorize 
     (dc (λ (dc dx dy)
           (define brush (send dc get-brush))
           (send dc set-brush (send brush get-color) 'solid)
           (send dc draw-polygon points dx dy)
           (send dc set-brush brush))
         triangle-width
         triangle-height)
     hierarchy-color)))

(define (center-x main pict)
  (define-values (x y) (cc-find main pict))
  x)

(define (left-edge-x main pict)
  (define-values (x y) (lc-find main pict))
  x)


(define (add-dot-right main class field) (add-dot-left-right/offset main class field 0 rc-find))
(define (add-dot-right/space main class field [count 1])
  (add-dot-right/offset main class field (* count dot-edge-spacing)))
(define (add-dot-right/offset main class field offset)
  (add-dot-left-right/offset main class field offset rc-find))

(define (add-dot-left main class field) (add-dot-left-right/offset main class field 0 lc-find))
(define (add-dot-left/space main class field [count 1])
  (add-dot-left/offset main class field (* count (- dot-edge-spacing))))
(define (add-dot-left/offset main class field offset)
  (add-dot-left-right/offset main class field offset lc-find))

(define (add-dot-left-right/offset main class field offset finder)
  (define-values (_1 y) (cc-find main field))
  (define-values (x-edge _2) (finder main class))
  (add-dot main (+ x-edge offset) y))

(define add-dot-junction
  (case-lambda
    [(main x-pict y-pict) (add-dot-junction main x-pict cc-find y-pict cc-find)]
    [(main x-pict x-find y-pict y-find)
     (define-values (x _1) (x-find main x-pict))
     (define-values (_2 y) (y-find main y-pict))
     (add-dot main x y)]))

(define (add-dot-offset pict dot dx dy)
  (define-values (x y) (cc-find pict dot))
  (add-dot pict (+ x dx) (+ y dy)))

(define dot-δx (make-parameter 0))
(define dot-δy (make-parameter 0))

(define (add-dot pict dx dy)
  (define dot (blank))
  (values (pin-over pict (+ dx (dot-δx)) (+ dy (dot-δy)) dot) dot))

(define (connect-dots show-arrowhead? main dot1 . dots)
  (let loop ([prev-dot dot1]
             [dots dots]
             [pict main])
    (cond
      [(null? dots) pict]
      [else 
       (loop (car dots) 
             (cdr dots)
             (connect-two-dots pict prev-dot (car dots) (null? (cdr dots)) show-arrowhead?))])))

(define (connect-two-dots pict dot1 dot2 arrowhead? show-arrowhead?)
  (if arrowhead?
      (pin-arrow-line field-arrowhead-size pict
                      dot1 cc-find
                      dot2 cc-find
                      #:hide-arrowhead? (not show-arrowhead?)
                      #:color type-link-color)
      (pin-line pict
                dot1 cc-find
                dot2 cc-find
                #:color type-link-color)))

(define (hierarchy/layout tops bottoms 
                          #:every-other-space [every-other-space 0]
                          #:top-space [top-space 40]
                          #:bottom-space [bottom-space 40]
                          #:vertical-space [vertical-space 60])
  (hierarchy 
   (vc-append (apply ht-append top-space tops)
              (blank 0 vertical-space)
              (apply ht-append bottom-space 
                     (let loop ([bottoms bottoms]
                                [every-other? #f])
                       (cond
                         [(null? bottoms) '()]
                         [else
                          (cons (if every-other?
                                    (vc-append (blank 0 every-other-space)
                                               (car bottoms))
                                    (car bottoms))
                                (loop (cdr bottoms)
                                      (not every-other?)))]))))
   tops
   bottoms))

(define (add-dot-delta f dx dy)
  (parameterize ([dot-δx dx]
                 [dot-δy dy])
    (f)))


(define (right-right-reference main0 start-class start-field finish-class finish-name
                               [count 1] 
                               #:connect-dots [connect-dots connect-dots]
                               #:dot-delta [dot-delta 0])
  (define going-down?
    (let-values ([(_1 start-y) (find-cc main0 start-field)]
                 [(_2 finish-y) (find-cc main0 finish-name)])
      (< start-y finish-y)))
  (define-values (main1 dot1)
    (add-dot-delta (λ () (add-dot-right main0 start-class start-field))
                   0
                   (if going-down?
                       dot-delta
                       (- dot-delta))))
  (define-values (main2 dot2)
    (add-dot-delta (λ () (add-dot-right/space main1 start-class start-field count))
                   dot-delta
                   (if going-down?
                       dot-delta
                       (- dot-delta))))
  (define-values (main3 dot3)
    (add-dot-delta (λ () (add-dot-right main2 finish-class finish-name))
                   0
                   (if going-down?
                       (- dot-delta)
                       dot-delta)))
  (define-values (main4 dot4) (add-dot-delta (λ () (add-dot-junction main3 dot2 dot3)) 0 0))
  
  ;; these last two dots are just there for the delta-less arrowhead
  (define-values (main5 dot5) (add-dot-right main4 finish-class finish-name))
  (define-values (main6 dot6)
    (add-dot-delta
     (λ () (add-dot-right main5 finish-class finish-name))
     1 ;; just enough to get the arrowhead going the right direction; not enough to see the line
     0))
  
  (connect-dots #t (connect-dots #f main6 dot1 dot2 dot4 dot3) dot6 dot5))

(define (left-left-reference main0
                             start-class
                             start-field
                             finish-class
                             finish-name
                             [count 1]
                             #:connect-dots [connect-dots connect-dots]
                             #:dot-delta [dot-delta 0])
  (let ([going-down? (let-values ([(_1 start-y) (find-cc main0 start-field)]
                                  [(_2 finish-y) (find-cc main0 finish-name)])
                       (< start-y finish-y))])
    (define-values (main1 dot1)
      (add-dot-delta (λ () (add-dot-left main0 start-class start-field))
                     0
                     (if going-down?
                         dot-delta
                         (- dot-delta))))
    (define-values (main2 dot2)
      (add-dot-delta (λ () (add-dot-left/space main1 start-class start-field count))
                     (- dot-delta)
                     (if going-down?
                         dot-delta
                         (- dot-delta))))
    (define-values (main3 dot3)
      (add-dot-delta (λ () (add-dot-left main2 finish-class finish-name))
                     0
                     (if going-down?
                         (- dot-delta)
                         dot-delta)))
    (define-values (main4 dot4) (add-dot-delta (λ () (add-dot-junction main3 dot2 dot3)) 0 0))
    (define-values (main5 dot5) (add-dot-left main4 finish-class finish-name))
    (define-values (main6 dot6)
      (add-dot-delta
       (λ () (add-dot-left main5 finish-class finish-name))
       -1 ;; just enough to get the arrowhead going the right direction; not enough to see the line
       0))

    (connect-dots #t (connect-dots #f main6 dot1 dot2 dot4 dot3) dot6 dot5)))

(define (left-top-reference main0
                            start-class
                            start-field
                            finish-class
                            [count 1]
                            #:connect-dots [connect-dots connect-dots])
  (define-values (main1 dot1) (add-dot-left main0 start-class start-field))
  (define-values (main2 dot2) (add-dot-left/space main1 start-class start-field count))
  (define-values (main3 dot3) (add-dot-junction main2 dot2 cc-find finish-class ct-find))
  (connect-dots #t main3 dot1 dot2 dot3))

(define (right-left-reference main0
                              start-class
                              start-field
                              finish-class
                              finish-name
                              [offset (find-middle main0 start-class rc-find finish-class lc-find)]
                              #:connect-dots [connect-dots connect-dots])
  (define-values (main1 dot1) (add-dot-right main0 start-class start-field))
  (define-values (main2 dot2) (add-dot-right/offset main1 start-class start-field offset))
  (define-values (main3 dot3) (add-dot-left main2 finish-class finish-name))
  (define-values (main4 dot4) (add-dot-junction main3 dot2 dot3))
  (connect-dots #t main4 dot1 dot2 dot4 dot3))

(define (left-right-reference main0
                              start-class
                              start-field
                              finish-class
                              finish-name
                              [offset
                               (- (find-middle main0 start-class lc-find finish-class rc-find))]
                              #:connect-dots [connect-dots connect-dots])
  (define-values (main1 dot1) (add-dot-left main0 start-class start-field))
  (define-values (main2 dot2) (add-dot-left/offset main1 start-class start-field offset))
  (define-values (main3 dot3) (add-dot-right main2 finish-class finish-name))
  (define-values (main4 dot4) (add-dot-junction main3 dot2 dot3))
  (connect-dots #t main4 dot1 dot2 dot4 dot3))

(define (find-middle main p1 find1 p2 find2)
  (define-values (x1 y1) (find1 main p1))
  (define-values (x2 y2) (find2 main p2))
  (- (/ (+ x1 x2) 2) (min x1 x2)))

(define (right-top-reference main0
                             start-class
                             start-field
                             finish-class
                             [count 1]
                             #:connect-dots [connect-dots connect-dots])
  (define-values (main1 dot1) (add-dot-right main0 start-class start-field))
  (define-values (main2 dot2) (add-dot-right/space main1 start-class start-field count))
  (define-values (main3 dot3) (add-dot-junction main2 dot2 cc-find finish-class ct-find))
  (connect-dots #t main3 dot1 dot2 dot3))

(define connect-dots-contract (-> boolean? pict? pict? pict? ... (values pict?)))

(provide type-link-color)
(provide (contract-out
          [field-spec (->* ((or/c #f string?) string?) (string? #:default string?) pict?)]
          [class-name (->* (string?) (#:spacing-word string?) pict?)]
          [class-box (-> pict? (or/c #f (listof pict?)) (or/c #f (listof pict?)) pict?)]
          [hierarchy/layout
           (->* ((cons/c pict? (listof pict?)) (cons/c pict? (listof pict?)))
                (#:top-space integer?
                             #:bottom-space integer?
                             #:vertical-space integer?
                             #:every-other-space integer?)
                pict?)]
          [user-type-font (-> string? pict?)]
          [prim-type-font (-> string? pict?)]
          [var-font (-> string? pict?)]
          [normal-font (-> string? pict?)]
          [comment-font (-> string? pict?)]
          [hierarchy (-> pict? (cons/c pict? (listof pict?)) (cons/c pict? (listof pict?)) pict?)]
          [right-right-reference
           (->* (pict? pict? pict? pict? pict?)
                (number? #:connect-dots connect-dots-contract #:dot-delta number?)
                pict?)]
          [left-left-reference
           (->* (pict? pict? pict? pict? pict?)
                (number? #:connect-dots connect-dots-contract #:dot-delta number?)
                pict?)]
          [right-left-reference
           (->* (pict? pict? pict? pict? pict?) (number? #:connect-dots connect-dots-contract) pict?)]
          [left-right-reference
           (->* (pict? pict? pict? pict? pict?) (number? #:connect-dots connect-dots-contract) pict?)]
          [left-top-reference
           (->* (pict? pict? pict? pict?) (number? #:connect-dots connect-dots-contract) pict?)]
          [right-top-reference
           (->* (pict? pict? pict? pict?) (number? #:connect-dots connect-dots-contract) pict?)]
          [dot-edge-spacing number?]
          [connect-dots connect-dots-contract]
          [add-dot-right (-> pict? pict? pict? (values pict? pict?))]
          [add-dot-right/space (-> pict? pict? pict? (values pict? pict?))]
          [add-dot-left (-> pict? pict? pict? (values pict? pict?))]
          [add-dot-left/space (-> pict? pict? pict? (values pict? pict?))]
          [add-dot-junction
           (case-> (-> pict? pict? pict? (values pict? pict?))
                   (-> pict?
                       pict?
                       (-> pict? pict? (values number? number?))
                       pict?
                       (-> pict? pict? (values number? number?))
                       (values pict? pict?)))]
          [add-dot-offset (-> pict? pict? number? number? (values pict? pict?))]
          [add-dot (-> pict? number? number? (values pict? pict?))]
          [method-spec
           (->* (string? string?) (#:body (or/c #f pict?)) #:rest (listof (or/c #f string?)) pict?)]
          [java-this (-> pict?)]
          [field-arrowhead-size number?]))
