(ns adventofcode2019.intcode
  [:require [adventofcode2019.lib :refer :all]])

;; 0: function args&mem -> [mem (ctr -> ctr)]
;; 1: number of args
(defn- decode-op [opcode] 
  (let [str-code (format "%05d" opcode)
        flags (reverse (take 3 str-code))
        op (vec (drop 3 str-code))
        apply-flag (fn [flag arg]
                     (case flag
                       ;; ORs avoid returning nil
                       \0 (fn ([S] (or (get-in S [:memory arg]) 0)) 
                              ([_ _] arg))
                       \1 (constantly arg)
                       \2 (fn ([S] (or (get-in S [:memory (+ arg (:relctr S))]) 0)) 
                              ([S _] (+ arg (:relctr S))))))
        with-flags (fn [f]
                     (fn [S & args]
                       (apply f S (map apply-flag flags args))))]
    (with-flags 
      (case op
        [\0 \1] (fn [S a b c] ; ADD
                  (-> S 
                      (assoc-in [:memory (c S true)] (+' (a S) (b S)))
                      (update :ctr + 4)))
        [\0 \2] (fn [S a b c] ; MULT
                  (-> S
                      (assoc-in [:memory (c S true)] (*' (a S) (b S)))
                      (update :ctr + 4)))
        [\0 \3] (fn [S a _ _] ; IN
                  (-> S
                      (assoc-in [:memory (a S true)] (first (:input S)))
                      (update :input subvec 1)
                      (update :ctr + 2)))
        [\0 \4] (fn [S a _ _] ; OUT
                  (-> S
                      (update :output conj (or (get-in S [:memory (a S true)]) 0))
                      (update :ctr + 2))) 
        [\0 \5] (fn [S a b _] ; BNEQ
                   (update S :ctr (if (not= (a S) 0) (constantly (b S)) #(+ % 3))))
        [\0 \6] (fn [S a b _] ; BEQ
                   (update S :ctr (if (= (a S) 0) (constantly (b S)) #(+ % 3))))
        [\0 \7] (fn [S a b c] ; SLT
                  (-> S
                      (assoc-in [:memory (c S true)] (if (< (a S) (b S)) 1 0))
                      (update :ctr + 4)))
        [\0 \8] (fn [S a b c] ; SEQ
                  (-> S
                      (assoc-in [:memory (c S true)] (if (= (a S) (b S)) 1 0))
                      (update :ctr + 4)))
        [\0 \9] (fn [S a _ _] ; SREL
                   (-> S
                       (update :relctr + (a S))
                       (update :ctr + 2)))
        [\9 \9] (fn [S _ _ _] ; EXIT
                   (assoc S :exit 1))))))

(defn- perform-operation [{:keys [memory ctr] :as state}]
  (let [operation (decode-op (memory ctr))
        args (map memory [(+ 1 ctr) (+ 2 ctr) (+ 3 ctr)])]
    (apply operation state args)))

(defn build-state 
  ([program]
   (let [memory (into {} (map-indexed #(vector %1 %2) program))]
     {:memory memory :ctr 0 :input [] :output [] :relctr 0}))
  ([program settings]
   (merge (build-state program) settings)))

(defn intcode [{:as state :keys [memory output]}]
   (cond ; quit if :exit, step and return state if :step, else loop
       (get state :exit) {:memory memory :output output}
       (get state :step) (perform-operation state)
       :else (recur (perform-operation state))))