(in-package #:user-side-compiler)

(defclass asm-object ()
  ((%source :accessor source :initarg :source :initform nil)
   (%next :accessor next :initform nil)
   (%address :accessor address :initform nil)))

(defclass asm-label (asm-object)
  ((%name :accessor name :initarg :name)))

(defclass asm-instruction (asm-object)
  ((%opcode :accessor opcode :initarg :opcode)
   (%operand :accessor operand :initarg :operand :initform nil)
   (%byte-length :accessor byte-length :initarg :byte-length)))

(defvar *asm-labels* (make-hash-table :test #'equal)
  "A hash table connecting identities (names) to ASM-LABEL objects.")

(defvar *asm-head* nil
  "The first assembly object, used for traversal.")
(defvar *asm-foot* nil
  "The (so far) last assembly object, used for construction.")

(defun emit-asm-object (obj)
  (if (null *asm-foot*)
      (setf *asm-head* obj)
      (setf (next *asm-foot*) obj))
  (setf *asm-foot* obj))

(defun emit-asm-instruction (&rest initargs)
  (let ((new-instruction (apply #'make-instance
                                'asm-instruction
                                initargs)))
    (emit-asm-object new-instruction)))

(defun emit-asm-label (identity)
  (let ((label (make-instance 'asm-label :name identity)))
    (setf (gethash identity *asm-labels*)
          label)
    (emit-asm-object label)))

(defvar *variable-allocations* (make-hash-table))
(defvar *last-instruction* :useless)

(defun varvec-index-in-zeropage-p (index)
  ;; TODO: Handle case of too many variables
  t)

(defun argvec-index-in-zeropage-p (index)
  ;; TODO: Handle case of too many arguments (? not realistically needed)
  t)

(defparameter +argvec-offset+ #x00)
(defparameter +varvec-offset+ #x08)

(defmacro with-variable-allocations (allocations &body body)
  `(let ((*variable-allocations* (make-hash-table)))
     (loop :for alloc :in ,allocations
           :do (setf (gethash (data alloc) *variable-allocations*) alloc))
     ,@body))

(defun allocation-details (data)
  (gethash data *variable-allocations*))

(defun data-reference (data)
  (ecase (strategy (allocation-details data))
    ((:named-variable :temporary-variable)
     (cons :address (+ (varvec-index (allocation-details data))
                       +varvec-offset+)))
    (:direct-to-argvec
     (unless (= (count data (inputs (first (users data)))) 1)
       (error "Tried to manually dereference reused function argument ~A." data))
     (cons :address (+ (position data (inputs (first (users data))))
                       +argvec-offset+)))
    (:constant
     (cons :immediate (ir-constant-value data)))
    (:accumulator
     (error "Accumulator-allocated data ~A being dereferenced. Ensuring this does not happen
is the responsibility of the pre-assembly compilation step."
            data))
    (:branch
     (error "Branch result ~A being dereferenced. This doesn't make much sense." data))
    (:not-saved
     (assert (null (users data)))
     (cons :not-saved nil))))

(defmacro define-normal-emitter (name immediate-opcode zeropage-opcode absolute-opcode)
  `(defun ,name (data-reference)
     (destructuring-bind (mode . value)
         data-reference
       (cond ((eql mode :immediate)
              (emit-asm-instruction :opcode ,immediate-opcode
                                    :operand (the (unsigned-byte 8) value)
                                    :byte-length 2))
             ((and (eql mode :address)
                   (< value #x100))
              (emit-asm-instruction :opcode ,zeropage-opcode
                                    :operand (the (unsigned-byte 8) value)
                                    :byte-length 2))
             ((eql mode :address)
              (emit-asm-instruction :opcode ,absolute-opcode
                                    :operand (the (unsigned-byte 16) value)
                                    :byte-length 3))
             (t
              (error "Invalid usage of ~A with arguments (~A ~A)" ',name mode value))))))

(define-normal-emitter emit-lda #xa9 #xa5 #xad)
(define-normal-emitter emit-sta (error "STA has no immediate mode.") #x85 #x8d)
(define-normal-emitter emit-adc #x69 #x65 #x6d)
(define-normal-emitter emit-sbc #xe9 #xe5 #xed)
(define-normal-emitter emit-cmp #xc9 #xc5 #xcd)

(defun emit-store-data (data)
  (if (member (strategy (allocation-details data))
              '(:not-saved :constant :branch))
      (setf *last-instruction* :useless)
      (progn
        (unless (eql (strategy (allocation-details data)) :accumulator)
          (if (and (eql (strategy (allocation-details data)) :direct-to-argvec)
                   (> (count data (inputs (first (users data)))) 1))
              (loop :with args := (inputs (first (users data)))
                    :for index := (position data args)
                      :then (position data args :start (1+ index))
                    :until (null index)
                    :do (emit-sta (cons :address (+ index +argvec-offset+))))
              (emit-sta (data-reference data))))
        (setf *last-instruction* (list :store data)))))

(defun emit-store-bool (data)
  "Stores the inverse of the zeroflag to DATA. Inverse so that non-0 is TRUE."
  ;; The "DATA is stored"-case
  (if (or (member (strategy (allocation-details data))
                  '(:constant :branch :not-saved)))
      (setf *last-instruction* :useless)
      (progn
        (emit-asm-instruction :opcode :php :byte-length 1)
        (emit-asm-instruction :opcode :pla :byte-length 1)
        (emit-asm-instruction :opcode :and :operand #b00000010 :byte-length 2)
        (emit-asm-instruction :opcode :lsr-a :byte-length 1)
        (emit-asm-instruction :opcode :not :operand #b00000001 :byte-length 2)
        (emit-sta :address (data-reference data))
        (setf *last-instruction* (list :store-zero-flag data)))))

(defun emit-load-data (data)
  (cond ((eql (strategy (allocation-details data)) :not-saved)
         (error "Tried to load unallocated data ~A." data))
        ((eql (strategy (allocation-details data)) :accumulator)
         (assert (eql (next (definition data)) (last-use data)))
         (setf *last-instruction* :useless))
        ((eql (strategy (allocation-details data)) :branch)
         (assert (typep (next (definition data)) 'ir-if)))
        ((eql (strategy (allocation-details data)) :direct-to-argvec)
         ;; TODO: Assert that it actually has been stored
         (setf *last-instruction* :useless))
        ((or (equal *last-instruction* (list :store data))
             (equal *last-instruction* (list :load data)))
         (setf *last-instruction* :useless))
        (t
         (emit-lda (data-reference data))
         (setf *last-instruction* (list :load data)))))

(defun emit-load-bool (data)
  (if (or (eql (strategy (allocation-details data)) :flags)
          (equal *last-instruction* (list :store-zero-flag data))
          (equal *last-instruction* (list :load-zero-flag data)))
      (setf *last-instruction* :useless)
      (progn
        (if (eql (strategy (allocation-details data)) :constant)
            (progn
              (emit-lda :immediate (ir-constant-value data))
              (emit-asm-instruction :opcode :and :operand 1 :byte-length 1)
              (emit-cmp :immediate 1))
            (progn
              (emit-lda :address (data-reference data))
              (emit-asm-instruction :opcode :and :operand 1 :byte-length 1)
              (emit-cmp :immediate 1)))
        (setf *last-instruction* (list :load-zero-flag data)))))

(defmethod compile-ir ((inst ir-inst))
  (warn "Skipped compiling ~A; no COMPILE-IR method" inst))

(defmethod compile-ir ((inst ir-return))
  (emit-asm-instruction :opcode #x60 :byte-length 1))

(defmethod compile-ir ((inst ir-plus))
  (unless (= (length (inputs inst)) 2)
    (error "During the final code generation step, IR-PLUS must have exactly 2 operands."))
  (emit-load-data (first (inputs inst)))
  (emit-asm-instruction :opcode #x18 :byte-length 1) ; Clear Carry
  (emit-adc (data-reference (second (inputs inst))))
  (emit-store-data (output inst)))

(defmethod compile-ir ((inst ir-minus))
  (unless (= (length (inputs inst)) 2)
    (error "During the final code generation step, IR-MINUS must have exactly 2 operands."))
  (emit-load-data (first (inputs inst)))
  (emit-asm-instruction :opcode #x38 :byte-length 1) ; Set Carry
  (emit-sbc (data-reference (second (inputs inst))))
  (emit-store-data (output inst)))

(defmethod compile-ir ((inst ir-assign))
  (emit-load-data (input inst))
  (emit-store-data (output inst))
  (setf *last-instruction* (list :store (output inst))))

(defmethod compile-ir ((inst ir-fetchvar))
  (emit-load-data (input inst))
  (emit-store-data (output inst))
  (setf *last-instruction* (list :store (output inst))))

(defmethod compile-ir ((inst ir-getconst))
  (emit-lda (cons :immediate (input inst)))
  (emit-store-data (output inst))
  (setf *last-instruction* (list :store (output inst))))

(defmethod compile-ir ((inst ir-call))
  (loop :for arg :in (inputs inst)
        :for arg-index :from 0
        :do (setf *last-instruction* :useless)
        :unless (eql (strategy (allocation-details arg)) :direct-to-argvec)
          :do (emit-lda (data-reference arg))
              (emit-sta (cons :address (+ arg-index +argvec-offset+)))
              (format t "~D. ~A~%" arg-index arg))
  (emit-asm-instruction :opcode #x20 :operand (callee inst) :byte-length 3)
  (emit-store-data (output inst)))

(defmethod compile-ir ((inst ir-jump))
  (unless (eql (next (iblock inst))
               (first (destinations inst)))
    (emit-asm-instruction :opcode #x4C
                          :operand (first (destinations inst))
                          :byte-length 3)
    (setf *last-instruction* (list :jump (first (destinations inst))))))

(defmethod compile-ir ((inst ir-if))
  (let ((next-iblock (next (iblock inst)))
        (then-iblock (first (destinations inst)))
        (else-iblock (second (destinations inst))))
    ;; With how the midstage is built, the true case's block comes immediately
    ;; after the IF, and therefore we should branch on false. This assert
    ;; ensures that this assumption become false without us noticing.
    ;; We implicitly fall through to THEN-BLOCK in the event of no branch.
    (assert (eql next-iblock then-iblock))
    (emit-load-data (input inst))
    (let ((prev-is-test-p (eql (strategy (allocation-details (input inst)))
                               :branch)))
      (if prev-is-test-p
          ;; If previous was a test, this instruction will be skipped if the
          ;; test succeeds and we fall through to the THEN case. If it fails,
          ;; this instruction executes, which jumps to the ELSE case.
          (emit-asm-instruction :opcode #x4C
                                :operand else-iblock
                                :byte-length 3)
          ;; If the input is a bool value, we have to actually do some work.
          (progn
            ;; BNE +3, if it's Not Equal (zero), then it's true, so we skip...
            (emit-asm-instruction :opcode #xD0 :operand 3 :byte-length 2)
            ;; ...the jump to the ELSE case.
            (emit-asm-instruction :opcode #x4C
                                  :operand else-iblock
                                  :byte-length 3))))
    (setf *last-instruction* :branch)))

(defun emit-branch-test-code (inputs output branch-opcode)
  (let ((branchp (eql (strategy (allocation-details output)) :branch)))
    (emit-load-data (first inputs))
    (unless branchp ; If we're *NOT* branching, we're storing a test result.
      ;; LDX #1. This value will go into A if the test succeeds.
      (emit-asm-instruction :opcode #xA2 :operand 1 :byte-length 2))
    (emit-cmp (data-reference (second inputs)))
    ;; The actual branch instruction for our test.
    (emit-asm-instruction :opcode branch-opcode :operand 3 :byte-length 2)
    (unless branchp
      ;; In the event of no branch -- we're storing the result -- we skip over
      ;; an LDX #0 instruction if it succeeded, run if if the test failed.
      (emit-asm-instruction :opcode #xA2 :operand 0 :byte-length 2)
      ;; And a NOP since the skip is three bytes, since in the event of a branch
      ;; it will be skipping over a whole JMP.
      (emit-asm-instruction :opcode #xEA :byte-length 1)
      ;; Then, regardless if we skipped or not, we put the result we have in X,
      ;; either 0 or 1, into A by emitting a TXA.
      (emit-asm-instruction :opcode #x8A :byte-length 1))
    ;; And regardless, we store the result, if applicable.
    (emit-store-data output)))

(defmethod compile-ir ((inst ir-test-equal))
  (emit-branch-test-code (inputs inst) (output inst) #xF0))

(defmacro do-asm-objects ((asm-obj start-asm-obj) &body body)
  `(loop :for ,asm-obj := ,start-asm-obj :then (next ,asm-obj)
         :until (null ,asm-obj)
         :do (progn
               ,@body)))

(defun link-compute-addresses (start-instruction origin-address)
  "First linking pass, computes the addresses of all labels."
  (let ((address origin-address))
    (loop :for asm-obj := start-instruction :then (next asm-obj)
          :until (null asm-obj)
          :do (setf (address asm-obj) address)
          :when (typep asm-obj 'asm-instruction)
            :do (incf address (byte-length asm-obj)))
    (values start-instruction address)))

(defun link-resolve-references (start-instruction)
  (flet ((resolve-iblock (asm-obj)
           (setf (operand asm-obj)
                 (multiple-value-bind (label existsp)
                     (gethash (name (operand asm-obj))
                              *asm-labels*)
                   (unless existsp
                     (error "Failed to resolve label ~A"
                            (name (operand asm-obj))))
                   label))))
    (loop :for asm-obj := start-instruction :then (next asm-obj)
          :until (null asm-obj)
          :when (typep asm-obj 'asm-instruction)
            :do (case (opcode asm-obj)
                  ((#x10 #x30 #x50 #x70 #x90 #xb0 #xd0 #xf0)
                   ;; Relative branches
                   (when (typep (operand asm-obj) 'iblock)
                     (resolve-iblock asm-obj))
                   ;; - 2 is to offset for the branch instruction's length
                   (unless (typep (operand asm-obj) '(unsigned-byte 8))
                     (let* ((offset (- (address (operand asm-obj)) (address asm-obj) 2)))
                       (setf (operand asm-obj) (ldb (byte 8 0) (the (signed-byte 8) offset))))))
                  (t
                   (when (typep (operand asm-obj) 'iblock)
                     (resolve-iblock asm-obj))
                   (when (typep (operand asm-obj) 'asm-label)
                     (setf (operand asm-obj) (address (operand asm-obj))))
                   (when (typep (operand asm-obj) 'asm-function)
                     (setf (operand asm-obj) (address (operand asm-obj)))))))))

(defun link-assembly (start-instruction origin-address)
  (link-compute-addresses start-instruction origin-address)
  ;; TODO: Branch correction
  (link-resolve-references start-instruction)
  start-instruction)

(defun compile-iblock (iblock)
  (emit-asm-label (name iblock))
  (do-instructions (inst iblock)
    (compile-ir inst)))

(defun compile-iblocks (start-iblock)
  (let ((*asm-head* nil)
        (*asm-foot* nil)
        (*asm-labels* (make-hash-table :test #'equal)))
    (do-iblocks (iblock start-iblock)
      (compile-iblock iblock))
    (link-assembly *asm-head* #xC000)
    *asm-head*))

(defun compiled-bytes (start-instruction)
  (declare (optimize (debug 3)))
  (check-type start-instruction asm-object)
  (loop :for asm-obj := start-instruction :then (next asm-obj)
        :do (loop :until (or (null asm-obj) (typep asm-obj 'asm-instruction))
                  :do (setf asm-obj (next asm-obj)))
        :until (null asm-obj)
        :append (list (opcode asm-obj))
        :when (> (byte-length asm-obj) 1)
          :append (ecase (byte-length asm-obj)
                    (2 (list (the (unsigned-byte 8) (operand asm-obj))))
                    (3 (list (ldb (byte 8 0)
                                  (the (unsigned-byte 16) (operand asm-obj)))
                             (ldb (byte 8 8)
                                  (the (unsigned-byte 16) (operand asm-obj))))))))

(defun quick-and-dirty-test-compile (text &key print-ir-p print-alloc-p make-asm-p)
  (with-input-from-string (source-stream text)
    (let ((*token-stream* (make-token-stream (tokenize source-stream))))
      (let ((rb (with-compilation-setup (root-block builder)
                  (compile-node (match-syntax program) builder)
                  root-block))
            (return-values '()))
        (do-iblocks (ib rb)
          (optim-prepare-direct-instructions ib)
          ;;(optim-commutative-constant-folding ib)
          (optim-reorder-arguments ib)
          (optim-call-duplicate-args ib)
          (optim-remove-unused ib))
        (push rb return-values)
        (let ((allocations (allocate-values rb)))
          (optim-reuse-temporary-slots rb allocations)
          (when print-ir-p
            (print-iblocks rb)
            (terpri))
          (when print-alloc-p
            (loop :for allocation :in allocations
                  :do (format t "~%~A - ~A~{ - ~A~}"
                              (data allocation)
                              (strategy allocation)
                              (unless (null (varvec-index allocation))
                                (list (varvec-index allocation)))))
            (terpri))
          (push allocations return-values)
          (when make-asm-p
            (with-variable-allocations allocations
              (push (compile-iblocks rb) return-values))))
        (values-list (nreverse return-values))))))