#lang racket
; 4 stages:
; DEFINITION: at this stage, labels must be defined as pointing to zero page memory or not
; SIZING: given the definitions, how big is each piece?
; ALLOCATION: knowing how big each piece is now allows us to assign concrete addresses to labels
; GENERATION: with addresses calculated, we can now generate the correct machine language code

; env: symbol -> def
; DEFINITION: (label, area, sizer)
; SIZING: (sizer denv) -> (size, generator)
; ALLOCATION: (label, area, sizer) -> (label, address, generator)
; GENERATION: (generator aenv) -> bytes

; two kinds of labels:
; lexical (local) labels: not exported from the containing block, generally used by relative branches
; global labels: can be referred to from anywhere

(require racket/generic)
(define-generics address-def
  (zeropage? address-def))
(define-generics op
  (size op env)
  (gen-bytes op env addr))

(struct area (addrStart addrEnd) #:transparent
  #:methods gen:address-def
  [(define (zeropage? area) (< (area-addrEnd area) #x100))])
(struct addr (addr) #:transparent
  #:methods gen:address-def
  [(define (zeropage? addr) (< (addr-addr addr) #x100))])
(define zp (area #x00 #xff))
(define stack (area #x0100 #x01ff))
(define 300hole (area #x0300 #x03cf))
(define grpage1 (area #x0400 #x07ff))
(define grpage2 (area #x0800 #x0bff))
(define c00hole (area #xc00 #x1fff))
(define page1 (area #x2000 #x3fff))
(define page2 (area #x4000 #x5fff))
(define freespace (area #x6000 #x95ff))      

(struct op-singlebyte (byte) #:transparent
  #:methods gen:op
  [(define (size op env) 1)
   (define (gen-bytes op env addr) (bytes (op-singlebyte-byte op)))])

(define (get-addr env sym-or-addr)
    (if (addr? sym-or-addr) sym-or-addr
        (hash-ref env sym-or-addr)))

(define (mode-from-arg env arg)
  (define (if-zp addr zpmode nonzpmode)
    (let [(resolved (get-addr env addr))]
      (list
        (if (zeropage? resolved) zpmode nonzpmode)
        resolved)))
  (match arg
    [(list imm) #:when (number? imm) (list 'imm imm)]
    [(list addr 'x) (if-zp addr 'zp-x 'abs-x)]
    [(list addr 'y) (if-zp addr 'zp-y 'abs-y)]
    [(list addr) (if-zp addr 'zp 'abs)]
    [(list 'a) (list 'a #f)]
    [(list (list addr)) (list 'abs* (get-addr addr))]
    [(list (list addr) 'y) (if-zp addr 'zp*-y #f)]
    [(list (list addr 'x)) (if-zp addr 'zp-x* #f)]
    [else #f]))

(struct op-arg (base modes arg) #:transparent
  #:methods gen:op
  [(define (size op env)
     (match (mode-from-arg env (op-arg-arg op))
       [(list _ #f) 1]
       [(list mode _) #:when (string-prefix? (symbol->string mode) "abs") 3]
       [else 2]))
   (define (gen-bytes op env addr)
     (let* [(modes (op-arg-modes op))
            (mode-val (mode-from-arg env (op-arg-arg op)))
            (mode (car mode-val))
            (val (cadr mode-val))
            (byte (modes (op-arg-base op) mode))]
        (case (size op env)
          [(1) (bytes byte)]
          [(2) (bytes byte (bitwise-and val #xff))]
          [(3) (bytes byte (bitwise-and val #xff) (bitwise-and (arithmetic-shift val -8) #xff))])))])

(struct op-branch (byte target) #:transparent
  #:methods gen:op
  [(define (size op env) 2)
   (define (gen-bytes op env addr)
     (let* [(dst-addr (get-addr env (op-branch-target op)))
            (rel (- (+ addr 2) (addr-addr dst-addr)))]
        (bytes (op-branch-byte op) (bitwise-and rel #xff))))])

(define-syntax (make-ops stx)
  (syntax-case stx ()
    [(_ op-maker ops ...)
     #`(begin
         #,@(for/list [(op (syntax->list #'(ops ...)))]
              #`(op-maker #,@op)))]))

(define-syntax-rule (make-singlebyte-op op byte)
  (begin
    (define (op) (op-singlebyte byte))
    (provide op)))

(make-ops make-singlebyte-op 
  (php #x08) (plp #x28) (pha #x48) (pla #x68) (dey #x88) (tay #xa8) (iny #xc8) (inx #xe8)
  (clc #x18) (sec #x38) (cli #x58) (sei #x78) (tya #x98) (clv #xb8) (cld #xd8) (sed #xf8)
  (txa #x8a) (txs #x9a) (tax #xaa) (tsx #xba) (dex #xca) (nop #xea) (brk #x00) (rti #x40)
  (rts #x60))

(define-syntax-rule (make-arg-op op base modes)
  (begin
    (define (op . arg) (op-arg base modes arg))
    (provide op)))

(define (indexed-modes modelist)
  (lambda (base mode)
    (match (index-of modelist mode)
      [imode #:when (number? imode) (bitwise-ior base (arithmetic-shift imode 2))]
      [#f #f])))

(define (without-modes modes . ignored-modes)
  (lambda (base mode) (if (member mode ignored-modes) #f (modes base mode))))
(define (only-modes modes . mode-subset)
  (lambda (base mode) (if (member mode mode-subset) (modes base mode) #f)))
(define (override-modes lmodes rmodes)
  (lambda (base mode) (let [(lmode (lmodes base mode))] (if lmode lmode (rmodes base mode)))))
(define (make-base aaa cc)
  (bitwise-ior cc (arithmetic-shift aaa 5)))
(define-syntax (make-cc-ops stx)
  (syntax-case stx ()
    [(_ cc modes ops ...)
     #`(begin
         #,@(for/list [(op (syntax->list #'(ops ...)))]
              (syntax-case op ()
                [(opname base)
                 #'(make-arg-op opname (make-base base cc) modes)])))]))

(define cc1-modes (indexed-modes '(zp-x* zp imm abs zp*-y zp-x abs-y abs-x)))
(make-cc-ops 1 cc1-modes 
  (ora 0) (and 1) (eor 2) (adc 3) (lda 5) (cmp 6) (sbc 7))
(make-arg-op sta (make-base 4 1) (without-modes cc1-modes 'imm))

(define cc2-modes (indexed-modes '(imm zp a abs _ zp-x _ asb-x)))
(make-cc-ops 2 (without-modes cc2-modes 'imm)
  (asl 0) (rol 1) (lsr 2) (ror 3))
(make-arg-op stx (make-base 4 2) (indexed-modes '(_ zp _ abs _ _ zp-y)))
(make-arg-op ldx (make-base 5 2) (indexed-modes '(imm zp _ abs _ _ zp-y _ abs-y)))
(make-cc-ops 2 (without-modes cc2-modes 'imm 'a)
  (dec 6) (inc 7))

(define cc0-modes (indexed-modes '(imm zp _ abs _ zp-x _ abs-x)))
(make-arg-op bit (make-base 1 0) (only-modes cc0-modes 'zp 'abs))
(make-arg-op jmp (make-base 2 0) 
  (lambda (base mode)
    (case mode
      [('abs) #x4c]
      [('abs*) #x6c]
      [else #f])))
(make-arg-op sty (make-base 4 0) (only-modes cc0-modes 'zp 'abs 'zp-x))
(make-arg-op ldy (make-base 5 0) cc0-modes)
(make-cc-ops 0 (only-modes cc0-modes 'imm 'zp 'abs)
  (cpy 6) (cpx 7))

(define-syntax-rule (make-branch-op op byte)
  (begin
    (define (op target) (op-branch byte target))
    (provide op)))
(make-ops make-branch-op
  (bpl #x10) (bmi #x30) (bvc #x50) (bvs #x70) (bcc #x90) (bcs #xb0) (bne #xd0) (beq #xf0))

(struct block (oplist area labels exported) #:transparent
  #:methods gen:op
  [(define (size block env)
      (let [(new-env
              (for/fold [(new-env env)]
                        [(label (hash-keys (block-labels block)))]
                (hash-set new-env label (block-area block))))]
        (for/sum [(op (block-oplist block))] (size op))))
   (define (gen-bytes block env addr)
      (let-values 
        [([end-addr iop-to-addr cop]
          (for/fold [(curr-addr addr) (iop-to-addr #hash()) (iop 0)]
                    [(op (block-oplist block))]
            (let* [(opsize (size op env))
                   (next-addr (+ addr opsize))]
              (values next-addr (hash-set iop-to-addr iop curr-addr) (+ iop 1)))))]))])
        ; draw the rest of the fucking owl
        
(define defblock (area ops)
  (let [(oplist
          (for/fold [(i 0) (oplist)]))]))