(local {: stream : kvstream : one} (require "stream"))

(local opcodes {})

; op mode arg
; single-byte ops
(let [ops
      {:php 0x08 :plp 0x28 :pha 0x48 :pla 0x68 :dey 0x88 :tay 0xa8 :iny 0xc8 :inx 0xe8
       :clc 0x18 :sec 0x38 :cli 0x58 :sei 0x78 :tya 0x98 :clv 0xb8 :cld 0xd8 :sed 0xf8
       :txa 0x8a :txs 0x9a :tax 0xaa :tsx 0xba :dex 0xca :nop 0xea :brk 0x00 :rti 0x40
       :rts 0x60}]
  (each [opcode byte (pairs ops)]
    (tset opcodes opcode (fn [mode] (if mode nil byte)))))

; branch ops
(let [ops {:bpl 0x10 :bmi 0x30 :bvc 0x50 :bvs 0x70 :bcc 0x90 :bcs 0xb0 :bne 0xd0 :beq 0xf0}]
  (each [opcode byte (pairs ops)]
    (tset opcodes opcode (fn [mode] (if (= mode :rel) byte nil)))))
(set opcodes.jsr (fn [mode] (if (= mode :abs) 0x20 nil)))

; aaabbbcc ops
(fn aaabbbcc [aaa cc modemap]
  (local base (bit.bor cc (bit.lshift aaa 5)))
  (fn [mode]
    (local bbb (. modemap mode))
      (if bbb (bit.bor base (bit.lshift bbb 2)) nil)))

(fn indexed-modes [...]
  (let [modemap {}]
    (each [index mode (pairs [...])]
      (tset modemap mode (- index 1)))
    modemap))

(fn without-modes [modemap ...]
  (let [newmodemap (table.clone modemap)]
    (each [_ mode (pairs [...])]
      (tset newmodemap mode nil))
    newmodemap))

(fn only-modes [modemap ...]
  (let [newmodemap {}]
    (each [_ mode (pairs [...])]
      (tset newmodemap mode (. modemap mode)))
    newmodemap))

; cc=1 ops
(let [cc1-modes (indexed-modes :zp-x* :zp :imm :abs :zp*-y :zp-x :abs-y :abs-x)
      ops {:ora 0 :and 1 :eor 2 :adc 3 :lda 5 :cmp 6 :sbc 7}]
  (each [opcode aaa (pairs ops)]
    (tset opcodes opcode (aaabbbcc aaa 1 cc1-modes))
  (tset opcodes :sta (aaabbbcc 4 1 (without-modes cc1-modes :imm)))))

; cc=2 ops
(let [cc2-modes (indexed-modes nil :zp :a :abs nil :zp-x nil :abs-x)]
  (each [opcode aaa (pairs {:asl 0 :rol 1 :lsr 2 :ror 3})]
    (tset opcodes opcode (aaabbbcc aaa 2 cc2-modes))
  (each [opcode aaa (pairs {:dec 6 :inc 7})]
    (tset opcodes opcode (aaabbbcc aaa 2 (without-modes cc2-modes :a))))))
(tset opcodes :stx (aaabbbcc 4 2 (indexed-modes nil :zp nil :abs nil nil :zp-y)))
(tset opcodes :ldx (aaabbbcc 5 2 (indexed-modes :imm :zp nil :abs nil nil :zp-y nil :abs-y)))

; cc=0 ops
(let [cc0-modes (indexed-modes :imm :zp nil :abs nil :zp-x nil :abs-x)]
  (tset opcodes :bit (aaabbbcc 1 0 (only-modes cc0-modes :zp :abs)))
  (tset opcodes :sty (aaabbbcc 4 0 (only-modes cc0-modes :zp :abs :zp-x)))
  (tset opcodes :ldy (aaabbbcc 5 0 cc0-modes))
  (each [opcode aaa (pairs {:cpy 6 :cpx 7})]
    (tset opcodes opcode (aaabbbcc aaa 0 (only-modes cc0-modes :imm :zp :abs)))))
(tset opcodes :jmp (fn [mode] (match mode :abs 0x4c :abs* 0x6c _ nil)))

(fn size [mode]
  (if
    (or (= mode nil) (= mode :a)) 1
    (= (mode:sub 1 3) :abs) 3
    2))
(fn opsize [op] (if (= op.opcode :block) 0 (size op.mode)))

(fn parse-mode-arg [op]
  (match op
    [_ :a] [:a nil]
    ([_ imm] ? (= (type imm) "number")) [:imm imm]
    ([op addr] ? (and (= (type addr) "string") (= (op:sub 1 1) "b"))) [:rel addr] ; branch
    ([_ addr] ? (= (type addr) "string")) [:addr addr]
    [_ addr :x] [:addr-x addr]
    [_ addr :y] [:addr-y addr]
    [_ [addr] :y] [:zp*-y addr]
    [_ [addr :x]] [:zp-x* addr]
    [_ [addr]] [:abs* addr]
    [_] [nil nil]
    _ (error (.. "Unrecognized syntax" (fv op)))))

(fn parse-ops [block ops]
  (var index 1)
  (each [_ op (ipairs ops)]
    (if (= (type op) "string")
      (tset block.symbols op index)
      (let [opcode (. op 1)]
        (if 
          (. opcodes opcode)
          (let [[mode arg] (parse-mode-arg op)]
            (table.insert block.code {: opcode : mode : arg}))

          (= opcode :block)
          (let [ops (table.clone op)]
            (table.remove ops 1)
            (table.insert block.code (parse-ops {: opcode :code [] :symbols {}} ops)))
          
          (error (.. "Unrecognized opcode " (fv opcode))))
        (set index (+ index 1)))))
  block)

(fn block [org ...] (parse-ops {:type :code : org :code [] :symbols {}} [...]))
(fn db [org init] {:type :var : org : init :size 1})
(fn dw [org init] {:type :var : org : init :size 2})
(fn allot [org size] {:type :var : org : size})

(fn make-env [blocks]
  {: blocks
   :push (fn [self block] 
     (make-env 
       (-> (stream self.blocks)
         (: :concat (one block))
         (: :tolist))))
   :lookup
   (fn [self name]
     (-> (stream self.blocks)
       (: :reverse)
       (: :map
         (fn [block] 
           (let [symbol (. block.symbols name)]
             (match (type symbol)
               "number" (. block.code symbol)
               "table" symbol
               _ nil))))
       (: :filter (fn [symbol] symbol))
       (: :first)))
   :lookup-block
   (fn [self name]
     (if
       (-> (stream self.blocks)
         (: :skip 1)
         (: :filter (fn [block] (. block.symbols name)))
         (: :next))
       (. self.blocks 2)
       (. (. self.blocks 1) name)))
   :is-zp? (fn [self name] (< (. (self:lookup-block name) :org) 0x100))})

(fn op-stream [env block]
  (-> (stream block.code)
    (: :map
      (fn [op]
        (if (= op.opcode :block)
          (: (one op env)
            :concat (op-stream (env:push op) op))
          (one op env))))
    (: :flatten)))

(fn int8-to-bytes [i]
  (string.char (bit.band i 0xff)))
(fn int16-to-bytes [i]
  (string.char (bit.band i 0xff) (bit.band (bit.rshift i 8) 0xff)))

(fn program []
  {:symbols {}
   :start-symbol :main
   :add (fn [self name elem] (tset self.symbols name elem))
   :block (fn [self name org ...] (self:add name (block org ...)))
   :db (fn [self name org init] (self:add name (db org init)))
   :dw (fn [self name org init] (self:add name (dw org init)))
   :allot (fn [self name org size] (self:add name (allot org size)))
   :op-stream
   (fn [self block] (op-stream (make-env [self block]) block))
   :patch-addr-modes
   (fn [self]
     (each [name block (pairs self.symbols)]
       (each [op env (: (self:op-stream block) :iter)]
         (when (and op.mode (= (op.mode:sub 1 4) :addr))
           (let [zp-mode (.. :zp (op.mode:sub 5))
                 abs-mode (.. :abs (op.mode:sub 5))
                 is-zp (and ((. opcodes op.opcode) zp-mode) (env:is-zp? op.arg))]
             (set op.mode (if is-zp zp-mode abs-mode)))))))
    :allocate
    (fn [self org size org-addr]
      (var addr (. org-addr org))
      (when (= nil addr)
        (tset org-addr org org)
        (set addr org))
      (tset org-addr org (+ addr size))
      addr)
    :allocate-var-block
    (fn [self block org-addr]
      (set block.addr (self:allocate block.org block.size org-addr)))
    :allocate-code-block
    (fn [self block org-addr]
      (set block.addr (self:allocate block.org 0 org-addr))
      (each [op env (: (self:op-stream block) :iter)]
        (set op.addr (self:allocate block.org (opsize op) org-addr))))
    :allocate-addresses
    (fn [self]
      (let [org-addr {}]
        (each [name block (pairs self.symbols)]
          (match block.type
            :var (self:allocate-var-block block org-addr)
            :code (self:allocate-code-block block org-addr)))))
    :generate-var-block
    (fn [self block]
      (match block.size
        1 (int8-to-bytes (or block.init 0))
        2 (int16-to-bytes (or block.init 0))
        n (string.rep "\0" n)))
    :lookup-addr
    (fn [self sym env]
      (local op (env:lookup sym))
      (if op
        op.addr
        (tonumber sym)))
    :generate-op-arg
    (fn [self op env]
      (if
        (= op.mode :imm) (int8-to-bytes op.arg)
        (= op.mode :rel) 
          (int8-to-bytes (- (self:lookup-addr op.arg env) (+ op.addr 2)))
        (= (size op.mode) 2) (int8-to-bytes (self:lookup-addr op.arg env))
        (= (size op.mode) 3) (int16-to-bytes (self:lookup-addr op.arg env))
        ""))
    :generate-op
    (fn [self op env]
      (let [bytegen (. opcodes op.opcode)]
        (if bytegen
          (.. (int8-to-bytes (bytegen op.mode)) (self:generate-op-arg op env))
          "")))
    :generate-code-block
    (fn [self block]
      (var bytes "")
      (each [op env (: (self:op-stream block) :iter)]
        (set bytes (.. bytes (self:generate-op op env))))
      (set block.bytes bytes))
    :generate-bytes
    (fn [self]
      (each [name block (pairs self.symbols)]
        (match block.type
          :var (self:generate-var-block block)
          :code (self:generate-code-block block))))
    :assemble
    (fn [self]
      (self:patch-addr-modes)
      (self:allocate-addresses)
      (self:generate-bytes))})

program