(import (chicken sort))

(include-relative "egg")
(include-relative "list")

(define (lay:dag:resolve-aliases dag aliases)
  (define (latest egg)
    (let* ((name (egg-name egg))
           (alias (alist-ref name aliases eq?)))
      (cond (alias
              (if (version>=? (egg-version alias) (egg-version egg))
                alias
                (error "egg version not satisfied" egg alias)))
            ((egg-to-skip? egg) #f)
            (else (error "egg not found" egg)))))
  (let loop ((dag dag)
             (seen '())
             (out '()))
    (if (null? dag)
      out
      (let* ((rest (cdr dag))
             (rule (car dag))
             (head (car rule))
             (deps (cdr rule))
             (head-name (egg-name head))
             (head-vers (egg-version head))
             (head-seen (alist-ref head-name seen eq? "")))
        (if (version>=? head-vers head-seen)
          (loop
            rest
            (cons (cons head-name head-vers) seen)
            (cons
              (cons head (lay:list:filter-map latest deps))
              out))
          (loop
            rest
            seen
            out))))))

(define (lay:dag:generate index)
  (define (make-alias name version)
    (cons name (egg-name+version (list name version))))

  (let ((version (car index)))
    (unless (string=? version "2")
      (error "unhandled eggs index version" version)))

  (let loop ((rest (cdr index))
             (last-name #f)
             (last-version #f)
             (aliases '())
             (dag '()))
    (if (null? rest)
      (let ((aliases (cons (make-alias last-name last-version) aliases)))
        (cons aliases (lay:dag:resolve-aliases dag aliases)))
      (let-values (((name version size checksum deps test-deps) (apply values (car rest))))
        (let ((new-name (and last-name (not (eq? last-name name))))
              (new-rule (cons (egg-name+version (list name version))
                              (map egg-name+version deps))))
          ; We assume that all versions of an egg are listed in a contiguous
          ; block of lines in the eggs index file: whenever we encounter a
          ; new egg name, we know that we have encountered all versions of
          ; the previous egg. It's easy enough to compare versions, that we
          ; don't need to assume that the versions are listed in order.
          ; This is so we can alias egg-name to egg-name:latest-version.
          (loop (cdr rest)
                name
                (if (or (not last-name)
                        new-name
                        (version>=? version last-version))
                  version
                  last-version)
                (if new-name
                  (cons (make-alias last-name last-version) aliases)
                  aliases)
                (cons new-rule dag)))))))

(define (lay:dag:prune dag-data eggs)
  (let ((aliases (car dag-data))
        (full-dag (cdr dag-data)))
   (let loop ((eggs (lay:list:filter (compose not egg-to-skip?) eggs))
               (dag '())
               (seen '()))
      (if (null? eggs)
        dag
        (let ((egg (or (alist-ref (car eggs) aliases eq?) (car eggs))))
          (cond
            ((memq egg seen)
             (loop (cdr eggs) dag seen))
            (else
              (let* ((rule (assq egg full-dag))
                     (_ (unless rule (error "no rule found for egg" egg)))
                     (deps (lay:list:filter
                             (lambda (node) (memq (car node) rule))
                             full-dag)))
                (loop
                  (flatten deps (cdr eggs))
                  (append deps dag)
                  (append rule seen))))))))))

(define (lay:dag:to-levels dag)
  (define (depends-on? elem lvl)
    (let ((adj-list (alist-ref elem dag eq?)))
      (and adj-list (lay:list:any (lambda (x) (memq x adj-list)) lvl))))
  (let loop ((topo (reverse (topological-sort dag eq?))) (lvls '()))
    (cond ((null? topo) (reverse lvls))
          ((null? lvls) (loop (cdr topo) (list (list (car topo)))))
          (else
            (let ((curr (car topo))
                  (next (cdr topo))
                  (x (car lvls))
                  (y (cdr lvls)))
              (let lvl-loop ((seen-lvl '())
                             (cand-lvl '())
                             (curr-lvl x)
                             (next-lvl y))
                (if (depends-on? curr curr-lvl)
                  (loop next
                        (append seen-lvl
                                (cons (cons curr cand-lvl)
                                      (cons curr-lvl next-lvl))))
                  (if (null? next-lvl)
                    (loop next
                          (append seen-lvl
                                  (list cand-lvl)
                                  (list (cons curr curr-lvl))))

                    (lvl-loop
                      (append seen-lvl (if (null? cand-lvl) '() (list cand-lvl)))
                      curr-lvl
                      (car next-lvl)
                      (cdr next-lvl))))))))))