(local lume (require "lib.lume"))
(local {: int8-to-bytes : int16-to-bytes : int24-to-bytes : int32-to-bytes} (require "lib.util"))

(fn make-env [block parent]
  {:parent parent
   :prg (or parent.prg parent)
   :root-block (or parent.root-block block)
   :block block
   :lookup-org
   (fn [self name]
     (if (or (. self.block.symbols name) (. self.block.globals name))
       self.root-block.org
       (self.parent:lookup-org name)))
   :lookup-addr
   (fn [self name]
     (local ipdat (. self.block.symbols name))
     (local ipdat-global (. self.block.globals name))
;      (print "looking up" name "in" self)
     (if
       (and ipdat (= (type ipdat) :number) (> ipdat (length self.block.pdats)))
       (+ self.block.addr self.block.size)

       (and ipdat (= (type ipdat) :function)) (ipdat self)

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

       ipdat-global (: (make-env (. self.block.pdats ipdat-global) block) :lookup-addr name)

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

(fn program [prg-base ?processor]
  (local {: opcodes : op-pdat : parse-mode-arg : addr-parser} (require (.. :asm. (or ?processor :6502))))
  ; 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 [last-symbol] {:type :block :pdats [] :preserved {} :symbols {} :globals {} : last-symbol})

  (fn parse-dats [block dats]
    (each [_ dat (ipairs dats)]
      (if (= (type dat) "string")
          (do (set block.last-symbol dat)
              (tset block.symbols dat (+ (length block.pdats) 1))
              (when (= (dat:sub 1 2) "G-")
                (tset block.globals dat true)))

          (not= (type dat) :table)
          (error (.. "Invalid operation " dat))

        (let [opcode (. dat 1)
              parser (. dat-parser opcode)
              meta (getmetatable dat)
              pdat
              (if parser (parser dat block)
                  (. opcodes opcode) (dat-parser.op dat)
                  (error (.. "Unrecognized opcode " (fv opcode))))]
          (when pdat
            (when meta (set block.last-meta meta))
            (set pdat.meta block.last-meta)
            (set pdat.nearest-symbol block.last-symbol)
            (table.insert block.pdats pdat)
            (when pdat.globals
              (each [name _ (pairs pdat.globals)]
                (tset block.globals name (length block.pdats))))
            (when pdat.preserved
              (each [name pdat-preserved (pairs pdat.preserved)]
                (tset block.preserved name pdat-preserved)))))))
    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 block.last-symbol) 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.dl [dl]           {:type :var :init (. dl 2) :size 4})
  (fn dat-parser.bytes [[_ bytes]] {:type :raw :bytes (if (= (type bytes) :table) (string.char (table.unpack bytes)) bytes)})
  (fn dat-parser.ref [ref]         {:type :ref :target (. ref 2)})
  (fn dat-parser.flatten [flat block]
    (parse-dats block (lume.slice flat 2))
    nil)
  (fn dat-parser.export [label block]
    (tset block.globals (. label 2) true)
    nil)
  (fn dat-parser.computed [label block]
    (tset block.symbols (. label 2) (. label 3))
    nil)

  (fn dat-parser.align [pad] {:type :pad :align (. pad 2)})
  (fn dat-parser.hot-preserve [[_ label & dats] block]
    (let [preserve-block (new-block)]
      (tset block.preserved label preserve-block)
      (tset preserve-block.globals label true)
      (parse-dats preserve-block [label])
      (parse-dats preserve-block dats)
      preserve-block))

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

  (fn describe-pdat [pdat]
    (if pdat.meta (.. pdat.meta.filename "@" pdat.meta.line)
      (.. (or pdat.nearest-symbol "<start of block>") " @" (or pdat.addr "<no address>"))))

  (fn process-pdat [pdat process default ...]
    (fn complain [ok ...]
      (if ok (values ...)
          (do (error (.. process " failed in " pdat.type " near " (describe-pdat pdat) " - " ...)))))
    (local processor (. pdat-processor pdat.type process))
    (if processor (complain (pcall #(processor pdat $...) ...)) default))

  (fn pdat-processor.raw.size [raw] (length raw.bytes))
  (fn pdat-processor.var.size [d] d.size)
  (fn pdat-processor.ref.size [r] 2)
  (fn pdat-processor.pad.size [pad]
    (let [misalignment (% pad.addr pad.align)]
      (if (= misalignment 0) 0
        (- pad.align misalignment))))
  (fn pdat-processor.var.bytes [d env]
    (local init (match (type d.init)
                  :number d.init
                  :nil 0
                  :function (d.init env)))
    (match d.size
      1 (int8-to-bytes init)
      2 (int16-to-bytes init)
      3 (int24-to-bytes init)
      4 (int32-to-bytes init)
      n (string.rep "\0" n)))
  (fn pdat-processor.ref.bytes [ref env]
    (int16-to-bytes (env:lookup-addr ref.target)))

  (fn pdat-processor.pad.bytes [pad] (string.rep "\0" pad.size))
  (fn pdat-processor.block.symbols [block]
    (lume.concat (lume.keys block.symbols) (lume.keys block.globals)))

  (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)
       (set block.org org))
     {: 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 ((or addr-parser tonumber) name))
     (if addr addr (error (.. "Symbol '" name "' not found"))))
   :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)))
   :lookup-org
   (fn [self name]
     (or (self:env-lookup name :lookup-org) (self:parse-addr name)))
   :pass
   (fn [self passname]
     (print 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)
   :read-hotswap
   (fn [self machine]
     (let [addr-to-label {}
           addr-to-size {}]
       (each [_ block (pairs self.org-to-block)]
         (each [label pdat (pairs block.preserved)]
           (tset addr-to-label pdat.addr label)
           (tset addr-to-size pdat.addr pdat.size)))
       (collect [addr bytes (pairs (machine:read-batch addr-to-size))]
         (values (. addr-to-label addr) bytes))))
   :write-hotswap
   (fn [self machine hotswap]
     (machine:write-batch
       (collect [label bytes (pairs hotswap)]
         (values (self:lookup-addr label) bytes))))
   :upload
   (fn [self machine]
     (if machine.upload (machine:upload self)
       (each [org block (pairs self.org-to-block)]
         (machine:write block.addr block.bytes))))
  })

{:new program}