; join [-a 1|2] [-v 1|2] [-e s] [-o list] [-t c] [-1 f] [-2 f] file1 file2:
; the lines of two files sorted on a field that agree on it, as one line:
; the field, then the rest of file1's, then the rest of file2's. Fields are
; split by blanks (leading ones ignored, a space between them out), or by
; c with -t; -1 and -2 name the field (1); -a n also prints the lines of
; file n that pair with none, -v n only those; -o the fields to print
; (0, or n.m, comma or blank separated), -e what a missing one prints.
; Keys compare byte by byte (LC_ALL=C); a key many lines share in both
; files gives every pair.

(def usage "usage: join [-a 1|2] [-v 1|2] [-e s] [-o list] [-t c] [-1 f] [-2 f] file1 file2")
(def opt (fn (opts letter default) (fold (fn (v o) (letv (l val) o (if (= l letter) val v))) default opts)))
(def all-of (fn (opts letter) (map (fn (o) (letv (l val) o val)) (filter (fn (o) (letv (l val) o (= l letter))) opts))))
(def is-stdin (fn (src) (and (= (type-of src) "string") (= src "-"))))
(def next-line (fn (src) (if (is-stdin src) (in-line) (file-line src))))
(def strip (fn (l) (if (and (> (byte-len l) 0) (= (byte-at l (- (byte-len l) 1)) 10)) (byte-sub l 0 (- (byte-len l) 1)) l)))
(def opened (fn (name) (if (= name "-") "-" (file-open name))))
(def whole (re-compile "^[0-9][0-9]*$"))

(def fields (fn (l sep)
  (if (= sep "")
    (filter (fn (f) (> (byte-len f) 0)) (str-split " " (str-replace "\t" " " l)))
    (str-split sep l))))
; field n (from 1) of fs, nil when there is none
(def field (fn (fs n) (if (and (>= n 1) (<= n (length fs))) (nth fs (- n 1)) (list))))
(def key-of (fn (l f sep) (let ((k (field (fields l sep) f))) (if (is-nil k) "" k))))

; (key lines next): the lines sharing first's key, and the line after them
(def read-group (fn (src first f sep)
  (let ((key (key-of (strip first) f sep)))
    (letv (lines next)
          (iterate (fn (st)
                     (letv (lines next) st
                       (if (or (is-nil next) (not (= (key-of (strip next) f sep) key))) (list)
                         (tuple (list-append lines (strip next)) (next-line src)))))
                   (tuple (list (strip first)) (next-line src)))
      (tuple key lines next)))))
(def group-at (fn (src line f sep) (if (is-nil line) (list) (read-group src line f sep))))

; the output fields of a pair (l1 or l2 nil for a line that pairs with none)
(def out-fields (fn (key l1 l2 f1 f2 sep spec empty)
  (let ((fs1 (if (is-nil l1) (list) (fields l1 sep))) (fs2 (if (is-nil l2) (list) (fields l2 sep))))
    (if (= (length spec) 0)
      (list-concat (list key)
                   (list-concat (map (fn (i) (nth fs1 i)) (filter (fn (i) (not (= i (- f1 1)))) (range 0 (length fs1))))
                                (map (fn (i) (nth fs2 i)) (filter (fn (i) (not (= i (- f2 1)))) (range 0 (length fs2))))))
      (map (fn (item)
             (letv (file n) item
               (let ((v (cond ((= file 0) key) ((= file 1) (field fs1 n)) (else (field fs2 n)))))
                 (if (is-nil v) empty v))))
           spec)))))

(def emit (fn (fs sep) (out-write (str-join (if (= sep "") " " sep) fs) "\n")))

; -o's list: (file n) tuples, file 0 for the join field
(def parse-spec (fn (s)
  (map (fn (p)
         (cond ((= p "0") (tuple 0 0))
               ((and (> (byte-len p) 2) (or (= (byte-sub p 0 2) "1.") (= (byte-sub p 0 2) "2."))
                     (not (is-nil (re-match whole (byte-sub p 2)))))
                (tuple (number (byte-sub p 0 1)) (number (byte-sub p 2))))
               (else (tuple -1 0))))
       (filter (fn (p) (> (byte-len p) 0)) (str-split " " (str-replace "," " " s))))))

(letv (opts files bad) (getopt ARGS "a:v:e:o:t:1:2:")
  (let ((sep (opt opts "t" ""))
        (f1s (opt opts "1" "1")) (f2s (opt opts "2" "1"))
        (as (all-of opts "a")) (vs (all-of opts "v"))
        (empty (opt opts "e" ""))
        (spec (parse-spec (str-join " " (all-of opts "o")))))
    (cond ((not (= bad "")) (do (err-write "join: bad option -" bad "\n" usage "\n") (exit-status 1)))
          ((not (= (length files) 2)) (do (err-write usage "\n") (exit-status 1)))
          ((or (is-nil (re-match whole f1s)) (is-nil (re-match whole f2s)) (= (number f1s) 0) (= (number f2s) 0))
           (do (err-write "join: -1 and -2 take a field number from 1\n") (exit-status 1)))
          ((> (length (filter (fn (x) (not (or (= x "1") (= x "2")))) (list-concat as vs))) 0)
           (do (err-write "join: -a and -v take 1 or 2\n") (exit-status 1)))
          ((> (length (filter (fn (it) (letv (file n) it (< file 0))) spec)) 0)
           (do (err-write "join: -o takes 0 and file.field items\n") (exit-status 1)))
          ((> (byte-len sep) 1) (do (err-write "join: -t takes one character\n") (exit-status 1)))
          (else
            (let ((a (opened (nth files 0))) (b (opened (nth files 1)))
                  (f1 (number f1s)) (f2 (number f2s))
                  (only (> (length vs) 0))
                  (show1 (> (length (filter (fn (x) (= x "1")) (list-concat as vs))) 0))
                  (show2 (> (length (filter (fn (x) (= x "2")) (list-concat as vs))) 0)))
              (cond ((and (= (type-of a) "string") (not (is-stdin a))) (do (err-write "join: " (nth files 0) ": " a "\n") (exit-status 1)))
                    ((and (= (type-of b) "string") (not (is-stdin b))) (do (err-write "join: " (nth files 1) ": " b "\n") (exit-status 1)))
                    (else
                      (iterate
                        (fn (st)
                          (letv (g1 g2) st
                            (cond ((and (is-nil g1) (is-nil g2)) (list))
                                  ((or (is-nil g2) (and (not (is-nil g1)) (letv (k1 ls1 n1) g1 (letv (k2 ls2 n2) g2 (< (byte-cmp k1 k2) 0)))))
                                   (letv (k1 ls1 n1) g1
                                     (do (if show1 (map (fn (l) (emit (out-fields k1 l (list) f1 f2 sep spec empty) sep)) ls1) #f)
                                         (tuple (group-at a n1 f1 sep) g2))))
                                  ((or (is-nil g1) (letv (k1 ls1 n1) g1 (letv (k2 ls2 n2) g2 (> (byte-cmp k1 k2) 0))))
                                   (letv (k2 ls2 n2) g2
                                     (do (if show2 (map (fn (l) (emit (out-fields k2 (list) l f1 f2 sep spec empty) sep)) ls2) #f)
                                         (tuple g1 (group-at b n2 f2 sep)))))
                                  (else
                                    (letv (k1 ls1 n1) g1
                                      (letv (k2 ls2 n2) g2
                                        (do (if only #f
                                              (map (fn (x) (map (fn (y) (emit (out-fields k1 x y f1 f2 sep spec empty) sep)) ls2)) ls1))
                                            (tuple (group-at a n1 f1 sep) (group-at b n2 f2 sep)))))))))
                        (tuple (group-at a (next-line a) f1 sep) (group-at b (next-line b) f2 sep))))))))))