#lang racket/base
(require pict)
(require racket/gui/base)
(require racket/format)
(require racket/function)
(require racket/match)
(require racket/contract)
(require "amb.rkt")

(define CELL-SIZE 30)
(define ARROW-LENGTH 30)
(define ARROW-SIZE 10)

(define (square n)
  (rectangle n n))

(define (pict-cell x)
  (cc-superimpose (square CELL-SIZE)
                  (scale (filled-ellipse CELL-SIZE CELL-SIZE) 1/4)))

(define (pict-cons-cell pair)
  (hc-append (pict-cell (car pair))
             (pict-cell (cdr pair))))

(define (cdr-find _ p)
  (values (- (pict-width p) (/ CELL-SIZE 2.0))
          (/ CELL-SIZE 2.0)))

(define (car-find p _)
  (values (/ CELL-SIZE 2.0)
          (/ CELL-SIZE 2.0)))

(define (left-find p p1)
  (values (- (pict-width p) (pict-width p1))
          (/ CELL-SIZE 2.0)))

(define (top-find p p1)
  (values (/ CELL-SIZE 2.0)
          (- (pict-height p) (pict-height p1))))

(define (horizontal-combined p1 p2)
  (ht-append ARROW-LENGTH p1 p2))

(define (vertical-combined n p1 p2)
  (vl-append (- (* (+ CELL-SIZE ARROW-LENGTH) n)
                (pict-height p1))
             p1 p2))

(define (arrow-from-cdr p1 p2)
  (pin-arrow-line ARROW-SIZE
                  (horizontal-combined p1 p2)
                  p1 cdr-find
                  p2 left-find))

(define (arrow-from-car n p1 p2)
  (pin-arrow-line ARROW-SIZE
                  (vertical-combined n p1 p2)
                  p1 car-find
                  p2 top-find))

(define (pict-atom x)
  (cc-superimpose (cellophane (filled-rectangle CELL-SIZE CELL-SIZE) 0)
                  (text (~s x))))

(define (cdr* lst)
  (if (null? lst)
      '()
      (cdr lst)))

(define (%pict-cons sexp height-list)
  (when (and (pair? height-list)
             (<= 0 (car height-list)))
    (amb))
  (cond
    [(pair? sexp)
     (let*-values ([(p) (pict-cons-cell sexp)]
                   [(p-1 height-list)
                    (let-values ([(pict-cdr cdr-height-list)
                                  (%pict-cons (cdr sexp) (cdr* height-list))])
                      (values (arrow-from-cdr p pict-cdr)
                              (cons 0 cdr-height-list)))]
                   [(p-2 height-list)
                    (match-let ([(cons len height-list)
                                 (let retry ([len 0]
                                             [height-list height-list])
                                   (amb (cons len height-list)
                                        (retry (add1 len) (map sub1 height-list))))])
                      (let-values ([(pict-car height-list)
                                    (%pict-cons (car sexp) height-list)])
                        (values (arrow-from-car len p pict-car)
                                (map (λ (x) (+ x len)) height-list))))])
       (values (lt-superimpose p-1 p-2) height-list))]
    [else
     (values (pict-atom sexp) (cons 0 (cdr* height-list)))]))

(define (pict-cons sexp)
  (call-with-amb
   (thunk
    (define-values (result height-list)
      (%pict-cons sexp '()))
    result)))
(provide/contract [pict-cons (any/c . -> . pict?)])