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

(local opcodes {})
; http://www.z80.info/decoding.htm
(fn argmatch [matcher arg]
  (case (type matcher)
    :function (matcher arg)
    :table    (when (= (type arg) :table)
                (accumulate [result {} i child (ipairs matcher) &until (= result nil)]
                  (case (argmatch child (. arg i))
                    argresult (lume.extend result argresult))))
    _         (when (= matcher arg) {})))

(fn comp-matchers [m1 m2] (fn [arg] (or (argmatch m1 arg) (argmatch m2 arg))))
(fn rekey [matcher k knew]
  (fn [arg] (match (matcher arg) {k val} {knew val})))

(fn try-parse-op [op matchers prefixgen]
  (when (= (length matchers) (- (length op) 1))
    (let [params
          (accumulate [result {} i matcher (ipairs matchers) &until (= result nil)]
            (let [arg (. op (+ i 1))
                  argresult (argmatch matcher arg)]
              (when (not= argresult nil)
                (lume.extend result argresult))))]
      (when (not= params nil)
        (case (prefixgen params)
          prefix (lume.extend {: prefix} params))))))

(fn chain-op [opcode f]
  (let [prev (or (. opcodes opcode) #nil)]
    (tset opcodes opcode
      (fn [op] (case (prev op)
                 result result
                 nil    (f op))))))

(fn opform [opcode matchers prefixgen]
  (chain-op opcode #(try-parse-op $1 matchers prefixgen)))

(fn table-matcher [tbl key]
  (let [lookup (collect [i val (ipairs tbl)] val (- i 1))]
    (fn [param] (case (. lookup param) octet {key octet}))))

(local cc (table-matcher [:nz :z :nc :c :po :pe :p :m] :cc))
(local reg (comp-matchers (table-matcher [:b :c :d :e :h :l :*hl :a] :reg)
                        #(when (argmatch [:hl] $1) {:reg 6})))
(local rp (table-matcher [:bc :de :hl :sp] :rp))
(local rp2 (table-matcher [:bc :de :hl :af] :rp))

(fn is-addr? [param] (and (= (type param) :string)
                          (not= param :ix) (not= param :iy)
                          (= (reg param) nil) (= (rp param) nil) (= (rp2 param) nil)))
(fn is-computed? [param] (= (type param) :function))
(fn is-number? [param] (= (type param) :number))
(fn rel-addr [param] (when (is-addr? param) {:rel8 param}))
(fn num [param] (when (or (is-number? param) (is-computed? param)) {:num param}))
(fn imm16 [param] (when (or (is-number? param) (is-computed? param) (is-addr? param)) {:imm16 param}))
(fn imm8 [param] (when (or (is-number? param) (is-computed? param)) {:imm8 param}))
(local addr imm16)

(fn im [arg] (match arg
                0 {:im 0}
                1 {:im 2}
                2 {:im 3}
                _ nil))

(fn ix [arg]
  (case arg
    :ix {:ixprefix "\xdd"}
    :iy {:ixprefix "\xfd"}))

(fn def-alu [f]
  (each [i opcode (ipairs [:add :adc :sub :sbc :and :xor :or :cp])]
    (f opcode (- i 1))))
(fn def-rot [f]
  (each [i opcode (ipairs [:rlc :rrc :rl :rr :sla :sra :sll :srl])]
    (f opcode (- i 1))))

(fn xyz [x y z] (int8-to-bytes (bit.bor (bit.lshift x 6) (bit.lshift y 3) z)))
(fn xpqz [x p q z] (int8-to-bytes (bit.bor (bit.lshift x 6) (bit.lshift p 4) (bit.lshift q 3) z)))

(opform :nop [] #(xyz 0 0 0))
(opform :ex [:af :af_] #(xyz 0 1 0))
(opform :djnz [rel-addr] #(xyz 0 2 0))
(opform :jr [rel-addr] #(xyz 0 3 0))
(opform :jr [cc rel-addr] #(when (< $1.cc 4) (xyz 0 (+ $1.cc 4) 0)))
(opform :ld [rp imm16] #(xpqz 0 $1.rp 0 1))
(opform :add [:hl rp] #(xpqz 0 $1.rp 1 1))
(opform :ld [[:bc] :a] #(xpqz 0 0 0 2))
(opform :ld [[:de] :a] #(xpqz 0 1 0 2))
(opform :ld [[addr] :hl] #(xpqz 0 2 0 2))
(opform :ld [[addr] :a] #(xpqz 0 3 0 2))
(opform :ld [:a [:bc]] #(xpqz 0 0 1 2))
(opform :ld [:a [:de]] #(xpqz 0 1 1 2))
(opform :ld [:hl [addr]] #(xpqz 0 2 1 2))
(opform :ld [:a [addr]] #(xpqz 0 3 1 2))
(opform :inc [rp] #(xpqz 0 $1.rp 0 3))
(opform :dec [rp] #(xpqz 0 $1.rp 1 3))
(opform :inc [reg] #(xyz 0 $1.reg 4))
(opform :dec [reg] #(xyz 0 $1.reg 5))
(opform :ld [reg imm8] #(xyz 0 $1.reg 6))
(opform :rlca [] #(xyz 0 0 7))
(opform :rrca [] #(xyz 1 0 7))
(opform :rla [] #(xyz 2 0 7))
(opform :rra [] #(xyz 3 0 7))
(opform :daa [] #(xyz 4 0 7))
(opform :cpl [] #(xyz 5 0 7))
(opform :scf [] #(xyz 6 0 7))
(opform :ccf [] #(xyz 7 0 7))
(opform :ld [reg (rekey reg :reg :reg2)] #(when (or (not= $1.reg 6) (not= $1.reg 6))
                                            (xyz 1 $1.reg $1.reg2)))
(opform :halt [] #(xyz 1 6 6))
(def-alu (fn [opcode alu] (opform opcode [reg] #(xyz 2 alu $1.reg))))
(opform :ret [cc] #(xyz 3 $1.cc 0))
(opform :pop [rp2] #(xpqz 3 $1.rp 0 1))
(opform :ret [] #(xpqz 3 0 1 1))
(opform :jp [:hl] #(xpqz 3 1 1 1))
(opform :exx [] #(xpqz 3 2 1 1))
(opform :ld [:sp :hl] #(xpqz 3 3 1 1))
(opform :jp [cc addr] #(xyz 3 $1.cc 2))
(opform :jp [addr] #(xyz 3 0 3))
(opform :out [[imm8] :a] #(xyz 3 2 3))
(opform :in [:a [imm8]] #(xyz 3 3 3))
(opform :ex [[:sp] :hl] #(xyz 3 4 3))
(opform :ex [:de :hl] #(xyz 3 5 3))
(opform :di [] #(xyz 3 6 3))
(opform :ei [] #(xyz 3 7 3))
(opform :call [cc addr] #(xyz 3 $1.cc 4))
(opform :push [rp2] #(xpqz 3 $1.rp 0 5))
(opform :call [addr] #(xpqz 3 0 1 5))
(def-alu (fn [opcode alu] (opform opcode [imm8] #(xyz 3 alu 6))))
(opform :rst [num] #(xyz 3 (/ $1.num 8) 7))

; DD / FD prefix
(each [opcode prev (pairs opcodes)]
  (tset opcodes opcode
    (fn [op]
      (case op
        [:ex :de :ix] (error "EX DI, IX does not exist")
        [:ex :de :iy] (error "EX DI, IY does not exist"))
      (var prefix nil)
      (var rel8 nil)
      (fn rewrite [new-prefix new-val ?rel8]
        (if (= prefix nil) (set prefix new-prefix)
            (not= prefix new-prefix) (error "Can't mix IX and IY in one op"))
        (if (and ?rel8 rel8) (error "Only one displacement is allowed")
            ?rel8 (set rel8 ?rel8))
        new-val)
      (let [op-new (icollect [_ arg (ipairs op)]
                      (case arg
                        :ix (rewrite "\xdd" :hl)
                        :iy (rewrite "\xfd" :hl)
                        :ixl (rewrite "\xdd" :l)
                        :iyl (rewrite "\xfd" :l)
                        :ixh (rewrite "\xdd" :h)
                        :iyh (rewrite "\xfd" :h)
                        [:ix rel8] (rewrite "\xdd" [:hl] rel8)
                        [:iy rel8] (rewrite "\xfd" [:hl] rel8)
                        _ arg))
            result (prev op-new)]
        (if (= prefix nil) result
            (= result nil) nil
          (lume.extend result {: rel8 :prefix (.. prefix result.prefix)}))))))

; CB prefix
(def-rot (fn [opcode rot] (opform opcode [reg] #(.. "\xcb" (xyz 0 rot $1.reg)))))
(opform :bit [num reg] #(.. "\xcb" (xyz 1 $1.num $1.reg)))
(opform :res [num reg] #(.. "\xcb" (xyz 2 $1.num $1.reg)))
(opform :set [num reg] #(.. "\xcb" (xyz 3 $1.num $1.reg)))
; ED prefix
(opform :in [reg [:c]] #(when (not= $1.reg 6) (.. "\xed" (xyz 1 $1.reg 0))))
(opform :in [[:c]] #(.. "\xed" (xyz 1 6 0)))
(opform :out [[:c] reg] #(when (not= $1.reg 6) (.. "\xed" (xyz 1 $1.reg 1))))
(opform :out [[:c]] #(.. "\xed" (xyz 1 6 1)))
(opform :sbc [:hl rp] #(.. "\xed" (xpqz 1 $1.rp 0 2)))
(opform :adc [:hl rp] #(.. "\xed" (xpqz 1 $1.rp 1 2)))
(opform :ld [[addr] rp] #(.. "\xed" (xpqz 1 $1.rp 0 3)))
(opform :ld [rp [addr]] #(.. "\xed" (xpqz 1 $1.rp 1 3)))
(opform :neg [] #(.. "\xed" (xyz 1 0 4)))
(opform :retn [] #(.. "\xed" (xyz 1 0 5)))
(opform :reti [] #(.. "\xed" (xyz 1 1 5)))
(opform :im [im] #(.. "\xed" (xyz 1 $1.im 6)))
(opform :ld [:i :a] #(.. "\xed" (xyz 1 0 7)))
(opform :ld [:r :a] #(.. "\xed" (xyz 1 1 7)))
(opform :ld [:a :i] #(.. "\xed" (xyz 1 2 7)))
(opform :ld [:a :r] #(.. "\xed" (xyz 1 3 7)))
(opform :rrd [] #(.. "\xed" (xyz 1 4 7)))
(opform :rld [] #(.. "\xed" (xyz 1 5 7)))
(opform :ldi [] #(.. "\xed" (xyz 2 4 0)))
(opform :cpi [] #(.. "\xed" (xyz 2 4 1)))
(opform :ini [] #(.. "\xed" (xyz 2 4 2)))
(opform :outi [] #(.. "\xed" (xyz 2 4 3)))
(opform :ldd [] #(.. "\xed" (xyz 2 5 0)))
(opform :cpd [] #(.. "\xed" (xyz 2 5 1)))
(opform :ind [] #(.. "\xed" (xyz 2 5 2)))
(opform :outd [] #(.. "\xed" (xyz 2 5 3)))
(opform :ldir [] #(.. "\xed" (xyz 2 6 0)))
(opform :cpir [] #(.. "\xed" (xyz 2 6 1)))
(opform :inir [] #(.. "\xed" (xyz 2 6 2)))
(opform :otir [] #(.. "\xed" (xyz 2 6 3)))
(opform :lddr [] #(.. "\xed" (xyz 2 7 0)))
(opform :cpdr [] #(.. "\xed" (xyz 2 7 1)))
(opform :indr [] #(.. "\xed" (xyz 2 7 2)))
(opform :otdr [] #(.. "\xed" (xyz 2 7 3)))

; DDCB / FDCB prefix
(def-rot (fn [opcode rot]
  (opform :ld [reg opcode [ix rel-addr]]
          #(when (not= $1.reg 6) (.. $1.ixprefix "\xcb" (xyz 0 rot $1.reg))))
  (opform opcode [[ix rel-addr]] #(.. $1.ixprefix "\xcb" (xyz 0 rot 6)))))
(opform :bit [num [ix rel-addr]] #(.. $1.ixprefix "\xcb" (xyz 1 $1.num 0)))
(opform :ld [reg :res num [ix rel-addr]]
        #(when (not= $1.reg 6) (.. $1.ixprefix "\xcb" (xyz 2 $1.num $1.reg))))
(opform :res [num [ix rel-addr]] #(.. $1.ixprefix "\xcb" (xyz 2 $1.num 6)))
(opform :ld [reg :set num [ix rel-addr]]
        #(when (not= $1.reg 6) (.. $1.ixprefix "\xcb" (xyz 3 $1.num $1.reg))))
(opform :set [num [ix rel-addr]] #(.. $1.ixprefix "\xcb" (xyz 3 $1.num 6)))

(fn parse-op [[opcode &as op]]
  (let [result ((. opcodes opcode) op)]
    (if (= result nil) (error (.. "no such opcode " (fennel.view op)))
        result)))

(local op-pdat {})
(fn op-pdat.size [op env]
  (+ (length op.prefix)
     (case op
       {: rel8}  1
       {: imm16} 2
       {: imm8}  1
       _         0)))

(fn decode-number [param env]
  (case (type param)
    :number param
    :string (env:lookup-addr param)
    :function (param (setmetatable {} {:__index #(env:lookup-addr $2)}) env)))

(fn op-pdat.bytes [op env]
  (.. op.prefix
      (case op
        {: rel8}  (int8-to-bytes (- (env:lookup-addr rel8) (+ op.addr 2)))
        {: imm16} (int16-to-bytes (decode-number imm16 env))
        {: imm8}  (int8-to-bytes (decode-number imm8 env))
        _         "")))

{: opcodes : parse-op : op-pdat : try-parse-op}