(local lume (require "lib.lume"))
(local {: int8-to-bytes : int16-to-bytes} (require "lib.util"))
(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 (lume.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] ? (or (= (type imm) "number") (= (type imm) "function"))) [:imm imm]
    ([opcode addr] ? (and (= (type addr) "string") (= (opcode:sub 1 1) "b"))) [:rel addr] ; branch
    [_ addr :x] [:addr-x addr]
    [_ [addr] :y] [:zp*-y addr]
    [_ addr :y] [:addr-y addr]
    [_ [addr :x]] [:zp-x* addr]
    ([_ addr] ? (= (type addr) "string")) [:addr addr]
    [_ [addr]] [:abs* addr]
    [_] [nil nil]
    _ (error (.. "Unrecognized syntax" (fv op)))))

(fn make-env [block parent]
  {:parent parent
   :block block
   :is-zp? (fn [self name] (self.parent:is-zp? name))
   :lookup-addr
   (fn [self name]
     (local ipdat (. self.block.symbols name))
;      (print "looking up" name "in" self)
     (if
       (and ipdat (> ipdat (length self.block.pdats)))
       (+ self.block.addr self.block.size)

       ipdat (. self.block.pdats ipdat :addr)

       (self.parent:lookup-addr name)))})

(fn program [prg-base]
  ; dat - anything that takes up space in the assembled output (op, dw, db, etc)
  ;   takes the form [:op args]
  ; pdat - a parsed dat; takes the form {:type type :addr addr ...}
  (local dat-parser {})
  (fn new-block [] {:type :block :pdats [] :symbols {}})

  (fn parse-dats [block dats]
    (each [_ dat (ipairs dats)]
      (if (= (type dat) "string")
        (tset block.symbols dat (+ (length block.pdats) 1))
        (let [opcode (. dat 1)
              parser (. dat-parser opcode)
              pdat
              (if
                parser (parser dat block)
                (. opcodes opcode) (dat-parser.op dat)
                (error (.. "Unrecognized opcode " (fv opcode))))]
          (table.insert block.pdats pdat))))
    block)

  (fn dat-parser.op [op]
    (let [[mode arg] (parse-mode-arg op)]
      {:type :op :opcode (. op 1) : mode : arg}))

  (fn dat-parser.block [block]
    (let [dats (lume.clone block)]
      (table.remove dats 1)
      (parse-dats (new-block) dats)))

  (fn dat-parser.db [db]       {:type :var :init (. db 2) :size 1})
  (fn dat-parser.dw [dw]       {:type :var :init (. dw 2) :size 2})
  (fn dat-parser.bytes [bytes] {:type :raw :bytes (. bytes 2)})
  (fn dat-parser.ref [ref]     {:type :ref :target (. ref 2)})
  (fn dat-parser.flatten [flat block]
    (parse-dats block (lume.slice flat 2))
    nil)

  (local pdat-processor {
    :op {}
    :var {}
    :ref {}
    :raw {}
    :block {}
  })

  (fn process-pdat [pdat process default ...]
    (pp pdat)
    (local processor (. pdat-processor pdat.type process))
    (if processor (processor pdat ...) default))

  (fn pdat-processor.op.patch [op env]
    (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)))))

  (fn pdat-processor.raw.size [raw] (length raw.bytes))
  (fn pdat-processor.op.size [op] (size op.mode))
  (fn pdat-processor.var.size [d] d.size)
  (fn pdat-processor.ref.size [r] 2)

  (fn pdat-processor.op.bytes [op env]
    (local bytegen (. opcodes op.opcode))
  ;   (pp op)
    (if bytegen
      (let [opbyte (bytegen op.mode)
            argbytes
            (if
              (and (= op.mode :imm) (= (type op.arg) "function"))
              (int8-to-bytes (op.arg env))

              (= op.mode :imm) (int8-to-bytes op.arg)
              (= op.mode :rel)
                (int8-to-bytes (- (env:lookup-addr op.arg) (+ op.addr 2)))
              (= (size op.mode) 2) (int8-to-bytes (env:lookup-addr op.arg))
              (= (size op.mode) 3) (int16-to-bytes (env:lookup-addr op.arg))
              "")]
        (if opbyte
          (.. (int8-to-bytes opbyte) argbytes)
          (error (.. op.opcode " doesn't support mode " op.mode))))
      ""))
  (fn pdat-processor.var.bytes [d env]
    (match d.size
      1 (int8-to-bytes (or d.init 0))
      2 (int16-to-bytes (or d.init 0))
      n (string.rep "\0" n)))
  (fn pdat-processor.ref.bytes [ref env]
    (int16-to-bytes (env:lookup-addr ref.target)))

  (fn pdat-processor.block.symbols [block]
    (lume.keys block.symbols))

  (fn pdat-processor.block.patch [block env]
    (local block-env (make-env block env))
    (each [_ pdat (ipairs block.pdats)]
      (process-pdat pdat :patch nil block-env)))

  (fn pdat-processor.block.allocate [block addr base-org]
    (var size 0)
    (local baseaddr (if base-org (+ base-org.addr base-org.size) addr))
    (set block.addr baseaddr)
    (each [_ pdat (ipairs block.pdats)]
      (set pdat.addr (+ baseaddr size))
      (process-pdat pdat :allocate nil pdat.addr)
      (local pdatsize (process-pdat pdat :size pdat.size))
      (set pdat.size pdatsize)
      (set pdat.addr (+ baseaddr size))
      (set size (+ size pdatsize)))
    (set block.size size))

  (fn pdat-processor.block.generate [block env]
    (local block-env (make-env block env))
    (var bytes "")
    (each [_ pdat (ipairs block.pdats)]
      (process-pdat pdat :generate nil block-env)
      (local pdatbytes (process-pdat pdat :bytes pdat.bytes block-env))
      (assert (= (type pdatbytes) :string) (.. "failed to generate bytes: " (fv pdat)))
      (set pdat.bytes pdatbytes)
      (set bytes (.. bytes pdatbytes)))
    (set block.bytes bytes))

  {:type :program
   :org-to-block {}
   :symbol-to-org {}
   :start-symbol :main
   : dat-parser
   : pdat-processor
   : new-block
   : prg-base
   :parse-dats (fn [self block dats] (parse-dats block dats))
   :dbg
   (fn [self ...]
     (when self.dbgfile
       (for [i 1 (select :# ...)]
         (when (not= i 1) (self.dbgfile:write " "))
         (self.dbgfile:write (fv (select i ...))))
       (self.dbgfile:write "\n")))
   :debug-to
   (fn [self filename]
     (set self.dbgfile (io.open filename :w)))
   :org
   (fn [self org]
     (var block (. self.org-to-block org))
     (when (not block)
       (set block (new-block))
       (tset self.org-to-block org block))
     {: block
      : org
      :prg self
      :ptr (fn [self] (tostring self.org))
      :append (fn [self ...] (self.prg:dbg self.org ...) (parse-dats self.block [...]) self)})
   :parse-addr
   (fn [self name]
     (local addr (tonumber name))
     (if addr addr (error (.. "Symbol '" name "' not found"))))
   :is-zp?
   (fn [self name]
     (local org (. self.symbol-to-org name))
     (if org (< org 0x100)
         self.prg-base (self.prg-base:is-zp? name)
       (< (self:parse-addr name) 0x100)))
   :env-lookup
   (fn [self name lookup ...]
     (local org (. self.symbol-to-org name))
     (if org (: (make-env (. self.org-to-block org) self) lookup name ...)
         self.prg-base (self.prg-base:env-lookup name lookup ...)
         nil))
;   :lookup-symbol (fn [self addr]) TODO
   :lookup-addr
   (fn [self name]
;       (print "looking up" name "in" self)
      (or (self:env-lookup name :lookup-addr) (self:parse-addr name)))
   :pass
   (fn [self passname]
     (each [org block (pairs self.org-to-block)]
       (: self passname org block (if self.prg-base (. self.prg-base.org-to-block org) nil))))
   :gather-symbols
   (fn [self org block]
     (each [_ name (ipairs (process-pdat block :symbols []))]
       (tset self.symbol-to-org name org)))
   :patch    (fn [self org block] (process-pdat block :patch nil self))
   :allocate (fn [self org block base-org] (process-pdat block :allocate nil org base-org))
   :generate (fn [self org block] (process-pdat block :generate nil self))
   :debug-pass (fn [self org block] (self:dbg org block))
   :assemble
   (fn [self]
     (self:pass :gather-symbols)
     (self:pass :patch)
     (self:pass :allocate)
     (self:pass :debug-pass)
     (self:pass :generate)
     (when self.dbgfile
       (self.dbgfile:close)
       (set self.dbgfile nil))
     self)
   :upload
   (fn [self machine only-self]
     (when (and (not only-self) self.prg-base) (self.prg-base:upload machine))
     (each [org block (pairs self.org-to-block)]
       (machine:write block.addr block.bytes)))
  })

{:new program}