#lang racket/base

(provide merge-preserves
         merge-preserves*
         merge2-preserves)

(require racket/dict)
(require racket/match)
(require racket/set)

(require "annotation.rkt")
(require "embedded.rkt")
(require "record.rkt")

(define (merge-preserves merge-embedded item . items)
  (merge-preserves* merge-embedded (cons item items)))

(define (merge-preserves* merge-embedded items)
  (match items
    [(cons item items)
     (for/fold [(item item)] [(other (in-list items))] (merge2-preserves merge-embedded item other))]
    ['() (error 'merge-preserves* "At least one value is required")]))

(define (unannotate x)
  (match x
    [(annotated annotations srcloc item) (values annotations srcloc item)]
    [_ (values '() #f x)]))

(define (merge-seqs merge-embedded as bs)
  (let loop ((as as) (bs bs))
    (match* (as bs)
      [['() bs] bs]
      [[as '()] as]
      [[(cons a as) (cons b bs)] (cons (merge2-preserves merge-embedded a b)
                                       (loop as bs))])))

(define (merge2-preserves merge-embedded a0 b0)
  (let merge ((a0 a0) (b0 b0))
    (define-values (a-annotations a-srcloc a) (unannotate a0))
    (define-values (b-annotations b-srcloc b) (unannotate b0))
    (define merged
      (match* (a b)
        [[(embedded a) (embedded b)] (embedded (merge-embedded a b))]
        [[(record al afs) (record bl bfs)] (record (merge al bl) (merge-seqs merge-embedded afs bfs))]
        [[(? list? as) (? list? bs)] (merge-seqs merge-embedded as bs)]
        [[(? set? as) (? set? bs)] (error 'merge-preserves "Set merge not supported")]
        [[(? dict? as) (? dict? bs)]
         (let* ((d (for/fold [(d (hash))] [((ak av) (in-dict as))]
                     (define bv (hash-ref bs ak void))
                     (hash-set d ak (if (void? bv) av (merge av bv)))))
                (d (for/fold [(d d)] [((bk bv) (in-dict bs))]
                     (if (hash-has-key? d bk)
                         d
                         (hash-set d bk bv)))))
           d)]
        [[_ _] (if (equal? a b) a (error 'merge-preserves "Cannot merge"))]))
    (define srcloc (or a-srcloc b-srcloc))
    (define annotations (append a-annotations b-annotations))
    (if (or (pair? annotations) srcloc)
        (annotated annotations srcloc merged)
        merged)))

(module+ test
  (require rackunit)
  (check-equal? (merge-preserves list '(1 2) '(1 2 3)) '(1 2 3))
  (check-equal? (merge-preserves list (hash 'a 1) (hash 'b 1)) (hash 'a 1 'b 1))
  (check-equal? (merge-preserves list (hash 'a (embedded 3)) (hash 'b 4)) (hash 'a (embedded 3) 'b 4))
  (check-equal? (merge-preserves list (hash 'a '()) (hash 'a '(4))) (hash 'a '(4)))
  (check-equal? (merge-preserves list (hash 'a (record 'a '(1 2))) (hash 'a (record 'a '(1 2 3))))
                (hash 'a (record 'a '(1 2 3))))
  (check-exn #px"Cannot merge" (lambda () (merge-preserves list 'a 'b)))
  (check-exn #px"Cannot merge" (lambda () (merge-preserves list '(a) '(b a))))
  (check-exn #px"Cannot merge" (lambda () (merge-preserves list (hash 'a '(1 2)) (hash 'a '(2 1)))))
  (check-equal? (merge-preserves list (hash 'a '(1 2)) (hash 'b '(2 1))) (hash 'a '(1 2) 'b '(2 1)))
  (check-equal? (merge-preserves + (embedded 1) (embedded 2)) (embedded 3))
  (check-equal? (strip-annotations (merge-preserves list (annotate 1 'a) 1)) 1)
  (check-equal? (annotations (merge-preserves list (annotate 1 'a) 1)) '(a))
  (check-equal? (strip-annotations (merge-preserves list 1 (annotate 1 'a))) 1)
  (check-equal? (annotations (merge-preserves list 1 (annotate 1 'a))) '(a))
  (check-equal? (strip-annotations (merge-preserves list (annotate 1 'b) (annotate 1 'a))) 1)
  (check-equal? (annotations (merge-preserves list (annotate 1 'b) (annotate 1 'a))) '(b a))
  )
