#lang racket


;;;; REGISTERS ;;;;
(define stack-pointer   20)
(define program-counter 0)
(define frame-pointer   0)
(define io-in           0)
(define cmp-res         0)

;;;;;; STACK ;;;;;;

(define word 8)
(define stack (make-vector 1024 #\nul))

(define (push value)
    (let ([bytes (to-bytes value)])
        (let loop ([start (+ stack-pointer word)]
                   [byte-index 0])
            (if (< byte-index (vector-length bytes))
                (begin
                    (vector-set! stack (- start byte-index) (vector-ref bytes byte-index))
                    (loop start (+ byte-index 1)))
                (begin
                    (set! stack-pointer start))))
    ))

(define (pop)
    (let loop ([value (make-vector word #\nul)]
               [stack-index stack-pointer]
               [value-index 0])
            (if (< value-index word)
                (begin
                    (vector-set! value value-index (vector-ref stack stack-index))
                    (loop value (- stack-index 1) (+ value-index 1)))
                (begin
                    (set! stack-pointer stack-index)

                    value))))

(define (to-bytes sth)
    (cond
        [(vector? sth) sth]
        [(string? sth) (list->vector (map char->integer (string->list sth)))]
        [(number? sth) (let loop ([value sth]
                                  [remainders '()])
                            (if (> value 0)
                                (loop (quotient value 2) (cons (remainder value 2) remainders))
                                (let loop ([to-pad (- 64 (length remainders))]
                                           [result remainders])
                                        (if (> to-pad 0)
                                            (loop (- to-pad 1) (cons 0 result))
                                            (let loop ([to-split result]
                                                       [splitted '()])
                                                    (if (not (empty? to-split))
                                                        (loop (drop to-split 8) (cons (take to-split 8) splitted))
                                                        (list->vector (map
                                                            (lambda (list-of-bits)
                                                                (let loop ([index 0]
                                                                           [res 0])
                                                                    (if (< index (length list-of-bits))
                                                                        (loop (+ index 1) (+ res (* (list-ref (reverse list-of-bits) index) (expt 2 index))))
                                                                        res)))
                                                            (reverse splitted)))))))))]))

(define (bytes->number bytes)
    (foldl + 0 (vector->list bytes))
)


;;;;;; PROGRAM ;;;;;;

(define program (make-vector 1024 #\nul))

;;;;;; PROCEDURES ;;;;;;

(define (io-read string)
    (display string)
    (let ([value (list->vector (map char->integer (string->list (symbol->string (read)))))])
        (mov "io-in" value)))

(define (io-write string)
    (display string))

(define (jump address)
    (mov "program-counter" address)
    #f
)

(define (mov register value)
    (cond
        [(eq? register "stack-pointer") (set! stack-pointer value)]
        [(eq? register "program-counter") (set! program-counter value)]
        [(eq? register "frame-pointer") (set! frame-pointer value)]
        [(eq? register "io-in") (set! io-in value)]
        [(eq? register "cmp-res") (set! cmp-res value)]
    )
)

(define (push-register register)
    (cond
        [(eq? register "stack-pointer") (push stack-pointer)]
        [(eq? register "program-counter") (push program-counter)]
        [(eq? register "frame-pointer") (push frame-pointer)]
        [(eq? register "io-in") (push io-in)]
        [(eq? register "cmp-res") (push cmp-res)]
    )
)

(define (call address)
    (push program-counter)
    (set! frame-pointer stack-pointer)
    (jump address)
)

(define (pwd-cmp)
    (let ([given-pwd (pop)])
        (if (eq? given-pwd (flag))
            (mov "cmp-res" 1)
            (mov "cmp-res" 0)
        )
    )
)

(define (ifz address)
    (when (eq? cmp-res 0)
        (jump address)
    )
)

(define (ret)
    (let loop ()
        (unless (= stack-pointer frame-pointer)
            (pop)
            (loop)
        )
    )
    (let ([restore-program-counter (bytes->number (pop))])
        (set! program-counter restore-program-counter)
    )
)

(define (instruction-call cmd . arg)
    (cond
        [(eq? cmd "io-read") (io-read (list-ref arg 0))]
        [(eq? cmd "io-write") (io-write (list-ref arg 0))]
        [(eq? cmd "push") (push (list-ref arg 0))]
        [(eq? cmd "push-register") (push-register (list-ref arg 0))]
        [(eq? cmd "pop")  (pop)]
        [(eq? cmd "mov")  (mov (list-ref arg 0) (list-ref arg 1))]
        [(eq? cmd "jump") (jump (list-ref arg 0))]
        [(eq? cmd "call") (call (list-ref arg 0))]
        [(eq? cmd "ret")  (ret)]
        [(eq? cmd "pwd-cmp") (pwd-cmp)]
        [(eq? cmd "ifz") (ifz (list-ref arg 0))]
    )
)


(define (flag)
    (if (getenv "FLAG") (getenv "FLAG") "NO-FLAG-DEFINED"))

(define (run)
    (let ((value (vector-ref program program-counter)))
        (unless (eq? value #\nul)
            (unless (eq? (apply instruction-call value) #f)
                (set! program-counter (+ program-counter 1))
            )
            (run)
        )
    )
)

(define (debug)
    (displayln stack)
    (display "stack-pointer: ")(displayln stack-pointer)
    (display "program-counter: ")(displayln program-counter)
    (display "frame-pointer:")(displayln frame-pointer)
    (display "io-in: ")(displayln io-in)
    (display "cmp-res: ")(displayln cmp-res)
    (displayln "-------------------")
)

(define (write idx cmd . arg)
    (vector-set! program idx (cons cmd arg))
)

;;;;;;;PROGRAM DEFINITION;;;;;;;

;;;main
(write 0 "call" 10)


(write 10 "io-read" "Enter password:")
(write 11 "push-register" "io-in")
(write 12 "pwd-cmp")
(write 13 "ifz" 16)
(write 14 "io-write" "Successfully logged in")
(write 15 "ret")

(write 16 "io-write" "Wrong password...")
(write 17 "ret")

(write 50 "io-write" (flag))





(run)