(ns day9
  (:require [clojure.string :as string]
            [helpers :refer [get-input]]))

(def disk-map "2333133121414131402")

(defn blocks [dm]
  (loop [file? true
         file-id 0
         dm (vec dm)
         bs []]
    (if (empty? dm)
      bs
      (let [cur (parse-long (str (first dm)))]
        (if file?
          (recur false (inc file-id) (rest dm) (apply conj bs (take cur (repeat file-id))))
          (recur true file-id (rest dm) (apply conj bs (take cur (repeat \.)))))))))

(defn index-of
  "Returns the index of item. If start is given indexes prior to
  start are skipped."
  ([coll item] (.indexOf coll item))
  ([coll item start]
   (let [unadjusted-index (.indexOf (drop start coll) item)]
     (if (= -1 unadjusted-index)
       unadjusted-index
       (+ unadjusted-index start)))))

(defn defrag [bs]
  (let [free-pos (index-of bs \.)]
    (if (neg? free-pos)
      bs
      (let [last-block (peek bs)]
        (if (= \. last-block)
          (recur (pop bs))
          (recur (assoc (pop bs) free-pos last-block)))))))

(defn fs-checksum [bs]
  (->>
   bs
   (map-indexed
    (fn [i x]
      (if (number? x) (* i x) 0)))
   (reduce +)))

(time
 (println "filesystem checksum:"
          (-> 9 get-input string/trim blocks defrag fs-checksum)))


(defn write-blocks [bs range val]
  (reduce (fn [bs idx] (assoc bs idx val)) bs range))

(defn write-file [bs file new-pos old-pos]
  (let [f-len (count file)
        new-range (range new-pos (+ new-pos f-len))
        old-range (range old-pos (+ old-pos f-len))]
    (-> bs
        (write-blocks new-range (peek file))
        (write-blocks old-range \.))))

(defn free-space-len [bs start-pos]
  (loop [len 0
         idx start-pos]
    (if (not= \. (bs idx))
      len
      (recur (inc len) (inc idx)))))

(defn move-file [bs file f-pos]
  ;;(println "moving file:" file)
  (loop [bs bs search-pos (index-of bs \.)]
    ;;(println "looking for free space starting at position" search-pos (bs search-pos))
    (if (or (neg? search-pos) (>= search-pos f-pos))
      ;; looked through whole disk
      bs
      ;; check if file will fit
      (let [free-space-available (free-space-len bs search-pos)]
        (if (>= free-space-available (count file))
          (write-file bs file search-pos f-pos)
          (recur bs (index-of bs \. (+ search-pos free-space-available))))))))

(defn defrag2 [bs]
  (loop [bs bs
         block-positions (vec (range (count bs)))
         cur-file []]
    ;;(println (apply str bs))
    (if (empty? block-positions)
      bs ;; defragmentation done
      (let [cur-pos (peek block-positions)
            cur-block (bs cur-pos)]
        ;;(println "processing block" cur-block "at index" cur-pos)
        (if (and (seq cur-file) ;; current file is not empty
                 (or (= \. cur-block) ;; read whole file, free space started
                     ;; read whole file, new file started
                     (not= cur-block (peek cur-file))))
          ;; defrag read file and continue
          (recur (move-file bs cur-file (inc cur-pos)) block-positions [])
          (if (= cur-block \.)
            ;; going through free space
            (recur bs (pop block-positions) [])
            ;; new file started
            (recur bs (pop block-positions) (conj cur-file cur-block))))))))

(time
 (println "fs checksum (part 2):"
          (-> 9 get-input string/trim blocks defrag2 fs-checksum)))