#lang racket

(require racket/async-channel)

;; First some helpers

(define (op-mode opcode idx)
  (match (modulo (floor (/ opcode (expt 10 (+ idx 2)))) 10)
    [0 'pos]
    [1 'imm]
    [2 'rel]
    [_ (error "invalid mode")]))

(define (op-op opcode)
  (modulo opcode 100))

(define (tape-expand tape len)
  (cond [(> len (length tape)) (append tape (for/list ([i (in-range (- len (length tape)))]) 0))]
        [else tape]))

(define (tape-ref tape pos)
  (define new-tape (tape-expand tape (add1 pos)))
  (list-ref new-tape pos))

(define (tape-ref/mode tape pos mode)
  (match mode
    ['pos (tape-ref tape (tape-ref tape pos))]
    ['imm (tape-ref tape pos)]
    ['rel (tape-ref tape (+ (intcode-rel-base) (tape-ref tape pos)))]
    [_ (error "invalid mode")]))

(define (tape-set tape where what)
  (define new-tape (tape-expand tape (add1 where)))
  (list-set new-tape where what))

(define (tape-set/mode tape where what mode)
  (match mode
    ['pos (tape-set tape (tape-ref tape where) what)]
    ['rel (tape-set tape (+ (intcode-rel-base) (tape-ref tape where)) what)]
    [_ (error "invalid mode")]))

;; Generic abstracted opcode interpreter function
;; Calls handler for opcode-specific handling, which should call callback to continue or return to
;; exit.
(define (intcode-handle-opcode callback spec handler tape pos)
  (let* ([opcode (list-ref tape pos)]
         [args (for/list ([spec-item spec] [idx (in-range (length spec))])
                 (match spec-item
                   ['in (tape-ref/mode tape (+ pos idx 1) (op-mode opcode idx))]
                   ['out
                    (lambda (what)
                      (tape-set/mode tape (+ pos idx 1) what (op-mode opcode idx)))]))])
    (apply handler callback tape (+ pos (length spec) 1) args)))

;; Create a system to register intcode ops
(struct intcode-op [code spec handler])
(define opcode-table (make-hash))

;; Evil mutation incoming!!!!
(define (intcode-register code spec handler)
  (hash-set! opcode-table code (intcode-op code spec handler)))

;; Main eval function
(define (intcode-eval tape [pos 0])
  (let* ([opcode (list-ref tape pos)]
         [handler-entry (hash-ref opcode-table (op-op opcode))])
    (intcode-handle-opcode intcode-eval
                           (intcode-op-spec handler-entry)
                           (intcode-op-handler handler-entry)
                           tape pos)))

;; Now it's time to define all the operations we can do

;; Binary operations!!!
(define spec-binop '(in in out))
(define ((intcode-generic-binop op) cb tape pos lhs rhs dst)
    (cb (dst (op lhs rhs)) pos))

(intcode-register 1 spec-binop (intcode-generic-binop +))
(intcode-register 2 spec-binop (intcode-generic-binop *))
(define (bool->number b) (if b 1 0))
(intcode-register 7 spec-binop (intcode-generic-binop (compose bool->number <)))
(intcode-register 8 spec-binop (intcode-generic-binop (compose bool->number =)))

;; I/O time
(define current-input-channel (make-parameter (make-async-channel)))
(define current-output-channel (make-parameter (make-async-channel)))
(define instance (make-parameter '()))
(define (instance-log what)
  (if (empty? (instance)) (void) (displayln (format "~a ~a" (instance) what))))
(define (intcode-input cb tape pos dst)
  (define inp (async-channel-get (current-input-channel)))
  (instance-log (list "input" inp))
  (cb (dst inp) pos))
(intcode-register 3 '(out) intcode-input)

(define (intcode-output cb tape pos src)
  (instance-log (list "output" src))
  (async-channel-put (current-output-channel) src)
  (cb tape pos))
(intcode-register 4 '(in) intcode-output)

;; Jumps
(define ((intcode-generic-jump which) cb tape pos arg dst)
  (cb tape (if (which arg) dst pos)))
(intcode-register 5 '(in in) (intcode-generic-jump (compose not zero?)))
(intcode-register 6 '(in in) (intcode-generic-jump zero?))

;; Relative base
(define intcode-rel-base (make-parameter 0))
(define (intcode-set-rel-base cb tape pos arg)
  (define new-base (+ (intcode-rel-base) arg))
  (instance-log (list "rel base" (intcode-rel-base) "+" arg "=" new-base))
  (parameterize ([intcode-rel-base new-base])
    (cb tape pos)))
(intcode-register 9 '(in) intcode-set-rel-base)

;; Exit
(define (intcode-exit cb tape pos)
  (instance-log "exit")
  (void))
(intcode-register 99 '() intcode-exit)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (parse-input input)
  (map string->number (string-split (string-trim input) ",")))

(define input (parse-input (file->string "input.9.txt")))

(async-channel-put (current-input-channel) 1)
(intcode-eval input)

(define (get-all chan)
  (let ([val (async-channel-try-get chan)])
    (if (equal? val #f) (void)
        (begin (displayln val) (get-all chan)))))
(get-all (current-output-channel))

;; Part 2
(async-channel-put (current-input-channel) 2)
(intcode-eval input)
(get-all (current-output-channel))