(local lume (require "lib.lume"))
(local {: lo : hi} (require "lib.util"))
(local link (require :link))
(local asm (require :asm.asm))

(fn inc16-stk [l h]
  [:block
    [:inc l :x]
    [:bne :done]
    [:inc h :x]
   :done])

(fn inc16 [l h]
  [:block
    [:inc l]
    [:bne :done]
    [:inc h]
   :done
  ])

(fn dec16 [l h]
  [:block
    [:lda l]
    [:bne :declow]
    [:dec h]
   :declow
    [:dec l]])

(fn add16 [l h]
  [:block
    [:clc]
    [:adc l]
    [:sta l]
    [:bcc :go]
    [:inc h]
    :go
  ])

(fn install-vm-parser [prg]
  (fn prg.dat-parser.vm [bytecodes]
    (local block (prg:new-block))
    (each [_ bytecode (ipairs (lume.slice bytecodes 2))]
      (if
        (= (type bytecode) :number)
        (prg:parse-dats block [[:ref :lit] [:dw bytecode]])

        (= (type bytecode) :string)
        (prg:parse-dats block [[:ref bytecode]])

        (= (type bytecode) :table)
        (prg:parse-dats block [bytecode])

        (error (.. "VM can't parse " (fv bytecode)))))
    block))

(fn mk-vm [prg options]
  (local code1 (prg:org 0xc00))
  (install-vm-parser prg)
  (local vm {
    :IP   :0x60
    :IPH  :0x61
    :W    :0x62
    :WH   :0x63
    :ROFF :0x64
    :TOP  :0x80
    :TOPH :0x81
    :ST1  :0x7e
    :ST1H :0x7f
    :ST2  :0x7c
    :ST2H :0x7d
    :RSTACK :0x6000
    :true 0xffff
    :false 0
    :code code1
    :nextsymid 1
    :ret (fn [self] [:jmp :next])
    :reserve (fn [self] [:block [:inx] [:inx]])
    :push
    (fn [self v]
      (local l (bit.band v 0xff))
      (local h (bit.band (bit.rshift v 8) 0xff))
      [:block (self:reserve)
        [:lda l]
        [:sta self.TOP :x]
        [:lda h]
        [:sta self.TOPH :x]
      ])
    :drop (fn [self] [:block [:dex] [:dex]])
    :def
    (fn [self name ...]
      (self.code:append name [:flatten ...] (self:ret)))
    :word
    (fn [self name ...]
      (self.code:append name [:jsr :subroutine] [:vm ...] [:vm :ret]))
    :inline
    (fn [self ...]
      [:block [:jsr :subroutine] [:vm ...] [:vm :restore]])
    :asm
    (fn [self ...]
      [:block [:vm :native] [:block ...] [:jsr :interpret]])
    :hotswap-sync
    (fn [self ...]
      (if link.machine.stub
        (do
          (link.machine:stub self.code :next
            [:lda #(lo ($1:lookup-addr :G-POST-HOTSWAP-RESET))]
            [:sta self.IP]
            [:lda #(hi ($1:lookup-addr :G-POST-HOTSWAP-RESET))]
            [:sta self.IPH]
            [:jmp :next])
          [:block [:vm :debug-stub :jmp :no-hotswap-reset [:block :G-POST-HOTSWAP-RESET] ...] :no-hotswap-reset])
        [:block]))
  })
  (set prg.vm vm)

  (vm.code:append :next
    [:ldy 0]
    [:lda [vm.IP] :y]  [:sta vm.W]
    (inc16 vm.IP vm.IPH)
    [:lda [vm.IP] :y] [:sta vm.WH]
    (inc16 vm.IP vm.IPH)

    [:jmp [vm.W]])

  (vm.code:append :reset
    [:lda #(lo ($1:lookup-addr :quit))]
    [:sta vm.IP]
    [:lda #(hi ($1:lookup-addr :quit))]
    [:sta vm.IPH]
    [:lda 0]
    [:sta vm.ROFF]
    [:ldx 0xfe]
    [:rts])

  (vm:def
   :subroutine ; usage: [jsr :subroutine] followed by bytecode
    [:ldy vm.ROFF]
    [:lda vm.IP] [:sta vm.RSTACK :y] [:iny]
    [:lda vm.IPH] [:sta vm.RSTACK :y] [:iny]
    [:sty vm.ROFF]
   :interpret ; usage: [jsr :interpret] followed by bytecode
    [:pla] [:sta vm.IP] [:pla] [:sta vm.IPH]
    (inc16 vm.IP vm.IPH))

  (vm:def :ret
    [:ldy vm.ROFF]
    [:dey] [:lda vm.RSTACK :y] [:sta vm.IPH]
    [:dey] [:lda vm.RSTACK :y] [:sta vm.IP]
    [:sty vm.ROFF])

  (vm.code:append :native [:jmp [vm.IP]])
  (vm.code:append :quit [:rts])
  (vm.code:append :restore
    [:lda vm.IP]  [:sta vm.W]
    [:lda vm.IPH] [:sta vm.WH]
    [:ldy vm.ROFF]
    [:dey] [:lda vm.RSTACK :y] [:sta vm.IPH]
    [:dey] [:lda vm.RSTACK :y] [:sta vm.IP]
    [:sty vm.ROFF]
    [:jmp [vm.W]])

  (vm.code:append :execute
    [:lda vm.TOP :x] [:sta vm.W]
    [:lda vm.TOPH :x] [:sta vm.WH]
    (vm:drop)
    [:jmp [vm.W]])

  (vm:def :drop (vm:drop))

  (vm:def :dup
    (vm:reserve)
    [:lda vm.ST1H :x]
    [:sta vm.TOPH :x]
    [:lda vm.ST1 :x]
    [:sta vm.TOP :x])

  (vm:def :swap
    [:lda vm.TOP :x]
    [:ldy vm.ST1 :x]
    [:sty vm.TOP :x]
    [:sta vm.ST1 :x]
    [:lda vm.TOPH :x]
    [:ldy vm.ST1H :x]
    [:sty vm.TOPH :x]
    [:sta vm.ST1H :x])

  (vm:def :over
    (vm:reserve)
    [:lda vm.ST2H :x]
    [:sta vm.TOPH :x]
    [:lda vm.ST2 :x]
    [:sta vm.TOP :x])
  (vm:word :2dup :over :over)

  (vm:def :>rot ; (a b c -- c a b)
    [:lda vm.TOP :x] ; a: c (a b c)
    [:ldy vm.ST2 :x] ; y: a (a b c)
    [:sta vm.ST2 :x] ; a: c (c b c)
    [:lda vm.ST1 :x] ; a: b (c b c)
    [:sta vm.TOP :x] ; a: b (c b b)
    [:sty vm.ST1 :x] ; y: a (c a b)
    [:lda vm.TOPH :x] ; a: c (a b c)
    [:ldy vm.ST2H :x] ; y: a (a b c)
    [:sta vm.ST2H :x] ; a: c (c b c)
    [:lda vm.ST1H :x] ; a: b (c b c)
    [:sta vm.TOPH :x] ; a: b (c b b)
    [:sty vm.ST1H :x] ; y: a (c a b)
  )
  (vm:def :<rot ; (a b c -- b c a)
    [:lda vm.TOP :x] ; a: c (a b c)
    [:ldy vm.ST1 :x] ; y: b (a b c)
    [:sta vm.ST1 :x] ; a: c (a c c)
    [:lda vm.ST2 :x] ; a: a (a c c)
    [:sta vm.TOP :x] ; a: a (a c a)
    [:sty vm.ST2 :x] ; y: b (b c a)
    [:lda vm.TOPH :x] ; a: c (a b c)
    [:ldy vm.ST1H :x] ; y: b (a b c)
    [:sta vm.ST1H :x] ; a: c (a c c)
    [:lda vm.ST2H :x] ; a: a (a c c)
    [:sta vm.TOPH :x] ; a: a (a c a)
    [:sty vm.ST2H :x] ; y: b (b c a)
  )

  (vm:def :get
    [:lda [vm.TOP :x]]
    [:tay]
    (inc16-stk vm.TOP vm.TOPH)
    [:lda [vm.TOP :x]]
    [:sta vm.TOPH :x]
    [:sty vm.TOP :x])

  (vm:def :set ; v p --
    [:lda vm.ST1 :x]
    [:sta [vm.TOP :x]]
    (inc16-stk vm.TOP vm.TOPH)
    [:lda vm.ST1H :x]
    [:sta [vm.TOP :x]]
    (vm:drop) (vm:drop))

  (vm:def :bget ; p -- b
    [:lda [vm.TOP :x]] [:sta vm.TOP :x]
    [:lda 0] [:sta vm.TOPH :x])

  (vm:def :bset ; b p --
    [:lda vm.ST1 :x] [:sta [vm.TOP :x]]
    (vm:drop) (vm:drop))

  (vm:def :lit
    (vm:reserve) [:ldy 0]
    [:lda [vm.IP] :y] [:sta vm.TOP :x]
    (inc16 vm.IP vm.IPH)
    [:lda [vm.IP] :y] [:sta vm.TOPH :x]
    (inc16 vm.IP vm.IPH))

  (vm:def :>r ; v --
    [:ldy vm.ROFF]
    [:lda vm.TOP :x]  [:sta vm.RSTACK :y] [:iny]
    [:lda vm.TOPH :x] [:sta vm.RSTACK :y] [:iny]
    [:sty vm.ROFF]
    (vm:drop))

  (vm:def :r> ; -- v
    (vm:reserve)
    [:ldy vm.ROFF]
    [:dey] [:lda vm.RSTACK :y] [:sta vm.TOPH :x]
    [:dey] [:lda vm.RSTACK :y] [:sta vm.TOP :x]
    [:sty vm.ROFF])

  (vm:def :rtop ; -- v
    (vm:reserve)
    [:ldy vm.ROFF]
    [:dey] [:lda vm.RSTACK :y] [:sta vm.TOPH :x]
    [:dey] [:lda vm.RSTACK :y] [:sta vm.TOP :x])

  (vm:def :rdrop
    [:ldy vm.ROFF] [:dey] [:dey] [:sty vm.ROFF])

  (vm:def :bz ; f --
    [:block
      [:lda vm.TOP :x]
      [:bne :skip]
      [:lda vm.TOPH :x]
      [:beq :dojmp]
     :skip
      (vm:drop)
      [:lda 2] (add16 vm.IP vm.IPH) (vm:ret)
     :dojmp (vm:drop)]
   :jmp
    ; ugh I don't have enough registers for this; a (one-byte?) relative jump would maybe be better
    [:ldy 0] [:lda [vm.IP] :y] [:sta vm.W]
    [:iny] [:lda [vm.IP] :y] [:sta vm.IPH]
    [:lda vm.W] [:sta vm.IP])

  (fn vm.while [self preamble ...]
    [:block
      :_start
        [:vm (table.unpack preamble)]
        [:ref :bz] [:ref :_end]
        [:vm ...]
        [:ref :jmp] [:ref :_start]
      :_end])

  (fn vm.until [self ...]
    [:block :_start [:vm ...] [:ref :bz] [:ref :_start]])

  (fn vm.forever [self ...]
    [:block :_start [:vm ...] [:vm :jmp :_start]])

  (fn vm.for [self ...]
    [:vm :>r (vm:while [:rtop] :r> :dec :>r ...) :rdrop])

  (fn vm.when [self ...]
    [:block [:vm :bz :_end ...] :_end])

  (fn vm.if [self iftrue iffalse]
    [:block [:vm :bz :_else (table.unpack iftrue)] [:vm :jmp :_end] :_else [:vm (table.unpack iffalse)] :_end])

  (fn vm.case [self ...]
    (local block [:block])
    (local cases [...])
    (each [icase [case & action] (ipairs cases)]
      (table.insert block (.. :_case icase))
      (if (< icase (length cases))
          (do (table.insert block [:vm :dup case := :bz (.. :_case (+ 1 icase)) :drop (table.unpack action)])
              (table.insert block [:vm :jmp :_end]))
          (do (table.insert block [:vm :drop (table.unpack action)])
              (table.insert block :_end))))
    block)

  (fn vm.gensym [self]
    (local sym (.. "G-GEN-SYM-" self.nextsymid))
    (set self.nextsymid (+ self.nextsymid 1))
    sym)

  (fn vm.anon [self ...]
    (local sym (self:gensym))
    (self.code:append sym ...)
    [:vm :lit sym])

  (fn vm.str [self str] (self:anon [:bytes str] [:db 0]))

  (vm:def :$dovar ; usage: [jsr :$dovar] followed by reserved space
    (vm:reserve)
    [:pla] [:sta vm.TOP :x] [:pla] [:sta vm.TOPH :x]
    (inc16-stk vm.TOP vm.TOPH))

  (fn vm.var [self name init]
    (self.code:append name [:jsr :$dovar]
      (if (= (type init) :table) init
          [:dw init])))

  (vm:def :$doconst ; usage: [jsr :$doconst] followed by two bytes
    (vm:reserve)
    [:pla] [:sta vm.W] [:pla] [:sta vm.WH]
    [:ldy 1] [:lda [vm.W] :y] [:sta vm.TOP]
    [:iny] [:lda [vm.W] :y] [:sta vm.TOPH])

  (fn vm.const [self name val]
    (self.code:append name [:jsr :$doconst]
      (if (= (type val) :number) [:dw val] [:ref val])))

  (fn vm.defer [self name ?init]
    (self.code:append name [:jmp (or ?init :next)]))

  (vm:word :redefine ; xp deferredxp --
    1 :+ :set)

  (fn vm.override-deferred-parent [self prg name redirect]
    (local addr (prg.prg-base:lookup-addr name))
    (local org (prg:org (+ addr 1)))
    (org:append [:ref redirect]))

  (vm:defer :interactive-eval-checkpoint)

  (fn vm.gen-eval-prg [self vmcode]
    (local eval-prg (asm.new self.code.prg))
    (install-vm-parser eval-prg)
    (local org (eval-prg:org self.code.org))
    (org:append :code-to-eval
      [:jsr :subroutine]
      [:vm :lit :next :lit :interactive-eval-checkpoint :redefine]
      vmcode
      [:vm :ret])
    (self:override-deferred-parent eval-prg :interactive-eval-checkpoint :code-to-eval)
    (eval-prg:assemble)
    eval-prg)

  (vm:def :+ ; a b -- c
    [:clc]
    [:lda vm.ST1 :x] [:adc vm.TOP :x] [:sta vm.ST1 :x]
    [:lda vm.ST1H :x] [:adc vm.TOPH :x] [:sta vm.ST1H :x]
    (vm:drop))

  (vm:def :- ; a b -- c
    [:sec]
    [:lda vm.ST1 :x] [:sbc vm.TOP :x] [:sta vm.ST1 :x]
    [:lda vm.ST1H :x] [:sbc vm.TOPH :x] [:sta vm.ST1H :x]
    (vm:drop))

  (vm:def :inc ; a -- a+1
    (inc16-stk vm.TOP vm.TOPH))

  (vm:def :dec ; a -- a-1
    [:block
      [:lda vm.TOP :x]
      [:bne :declow]
      [:dec vm.TOPH :x]
     :declow
      [:dec vm.TOP :x]])

  (vm:def :not ; f - !f
    [:block
      [:lda vm.TOP :x]
      [:bne :zero]
      [:lda vm.TOPH :x]
      [:bne :zero]
      [:lda 0xff] [:sta vm.TOP :x] [:sta vm.TOPH :x] (vm:ret)
     :zero
      [:lda 0] [:sta vm.TOP :x] [:sta vm.TOPH :x]])

  (vm:def := ; a b -- f
    [:block
      [:lda vm.ST1 :x]
      [:cmp vm.TOP :x]
      [:bne :noteq]
      [:lda vm.ST1H :x]
      [:cmp vm.TOPH :x]
      [:bne :noteq]
      [:lda 0xff] (vm:drop) [:sta vm.TOP :x] [:sta vm.TOPH :x] (vm:ret)
     :noteq
      [:lda 0] (vm:drop) [:sta vm.TOP :x] [:sta vm.TOPH :x]])

  (vm:def :& ; a b -- c
    [:lda vm.TOP :x] [:and vm.ST1 :x] [:sta vm.ST1 :x]
    [:lda vm.TOPH :x] [:and vm.ST1H :x] [:sta vm.ST1H :x]
    (vm:drop))

  (vm:def :| ; a b -- c
    [:lda vm.TOP :x] [:ora vm.ST1 :x] [:sta vm.ST1 :x]
    [:lda vm.TOPH :x] [:ora vm.ST1H :x] [:sta vm.ST1H :x]
    (vm:drop))

  vm)

{:new mk-vm}