#let colors = (
  address: rgb(228, 103, 103),
  label: rgb(119, 97, 204),
  op: rgb(235, 105, 40)
)

#let opcodes = (
  "lui":   0b0110111,
  "auipc": 0b0010111,
  "jal":   0b1101111,
  "jalr":  0b1100111,
  "beq":   0b1100011,
  "bne":   0b1100011,
  "blt":   0b1100011,
  "bge":   0b1100011,
  "bltu":  0b1100011,
  "bgeu":  0b1100011,
  "lb":    0b0000011,
  "lh":    0b0000011,
  "lw":    0b0000011,
  "lbu":   0b0000011,
  "lhu":   0b0000011,
  "sb":    0b0100011,
  "sh":    0b0100011,
  "sw":    0b0100011,
  "addi":  0b0010011,
  "slti":  0b0010011,
  "sltiu": 0b0010011,
  "xori":  0b0010011,
  "ori":   0b0010011,
  "andi":  0b0010011,
  "slli":  0b0010011,
  "srli":  0b0010011,
  "srai":  0b0010011,
  "add":   0b0110011,
  "sub":   0b0110011,
  "sll":   0b0110011,
  "slt":   0b0110011,
  "sltu":  0b0110011,
  "xor":   0b0110011,
  "srl":   0b0110011,
  "sra":   0b0110011,
  "or":    0b0110011,
  "and":   0b0110011,
)

#let types = (
  "lui":   "U",
  "auipc": "U",
  "jal":   "J",
  "jalr":  "I",
  "beq":   "B",
  "bne":   "B",
  "blt":   "B",
  "bge":   "B",
  "bltu":  "B",
  "bgeu":  "B",
  "lb":    "L",
  "lh":    "L",
  "lw":    "L",
  "lbu":   "L",
  "lhu":   "L",
  "sb":    "S",
  "sh":    "S",
  "sw":    "S",
  "addi":  "I",
  "slti":  "I",
  "sltiu": "I",
  "xori":  "I",
  "ori":   "I",
  "andi":  "I",
  "slli":  "I",
  "srli":  "I",
  "srai":  "I",
  "add":   "R",
  "sub":   "R",
  "sll":   "R",
  "slt":   "R",
  "sltu":  "R",
  "xor":   "R",
  "srl":   "R",
  "sra":   "R",
  "or":    "R",
  "and":   "R",
)

#let funct3 = (
  "jalr":  0b000,
  "beq":   0b000,
  "bne":   0b001,
  "blt":   0b100,
  "bge":   0b101,
  "bltu":  0b110,
  "bgeu":  0b111,
  "lb":    0b000,
  "lh":    0b001,
  "lw":    0b010,
  "lbu":   0b100,
  "lhu":   0b101,
  "sb":    0b000,
  "sh":    0b001,
  "sw":    0b010,
  "addi":  0b000,
  "slti":  0b010,
  "sltiu": 0b011,
  "xori":  0b100,
  "ori":   0b110,
  "andi":  0b111,
  "slli":  0b001,
  "srli":  0b101,
  "srai":  0b101,
  "add":   0b000,
  "sub":   0b000,
  "sll":   0b001,
  "slt":   0b010,
  "sltu":  0b011,
  "xor":   0b100,
  "srl":   0b101,
  "sra":   0b101,
  "or":    0b110,
  "and":   0b111,
)

#let funct7 = (
  "slli":  0b0000000,
  "srli":  0b0000000,
  "srai":  0b0100000,
  "add":   0b0000000,
  "sub":   0b0100000,
  "sll":   0b0000000,
  "slt":   0b0000000,
  "sltu":  0b0000000,
  "xor":   0b0000000,
  "srl":   0b0000000,
  "sra":   0b0100000,
  "or":    0b0000000,
  "and":   0b0000000,
)

#let regs = (
  "zero": 0, "ra": 1, "sp": 2, "gp": 3,
  "tp": 4, "t0": 5, "t1": 6, "t2": 7,
  "s0": 8, "fp": 8, "s1": 9, "a0": 10, "a1": 11,
  "a2": 12, "a3": 13, "a4": 14, "a5": 15,
  "a6": 16, "a7": 17, "s2": 18, "s3": 19,
  "s4": 20, "s5": 21, "s6": 22, "s7": 23,
  "s8": 24, "s9": 25, "s10": 26, "s11": 27,
  "t3": 28, "t4": 29, "t5": 30, "t6": 31
)

#let zfill(string, length, c: "0") = {
  let res = c * length + string
  return res.slice(-length)
}

#let int-to-bin(value, bits) = {
  if value < 0 {
    let m = (1).bit-lshift(bits) - 1
    let v = str((-value).bit-xor(m) + 1, base: 2)
    return zfill(v, bits, c: "1")
  }
  let v = str(value, base: 2)
  return zfill(v, bits)
}

#let parse-reg(value) = {
  let number = value.last()
  if not value.starts-with("x") {
    number = regs.at(value)
  }
  return int(number)
}

#let parse-imm(value) = {
  let imm = eval(value)
  assert(type(imm) == int)
  return imm
}

#let parse-addr(value) = {
  let matches = value.match(regex("(.*)\((.*)\)"))
  let imm = matches.captures.first()
  let reg = matches.captures.last()
  imm = parse-imm(imm)
  reg = parse-reg(reg)
  return (reg, imm)
}

#let parse-args(instr, instr-type, args) = {
  let parsed = (:)
  let r1 = parse-reg(args.first())

  if instr-type in ("R", "I", "L", "U", "J") {
    parsed.insert("rd", r1)
  } else if instr-type == "B" {
    parsed.insert("rs1", r1)
  } else if instr-type == "S" {
    parsed.insert("rs2", r1)
  }

  let r2 = none
  let r3 = none
  let imm = none
  if instr-type == "R" {
    r2 = parse-reg(args.at(1))
    r3 = parse-reg(args.at(2))
  } else if instr-type in ("I", "B") {
    r2 = parse-reg(args.at(1))
    imm = parse-imm(args.at(2))
  } else if instr-type in ("U", "J") {
    imm = parse-imm(args.at(1))
  } else if instr-type in ("L", "S") {
    (r2, imm) = parse-addr(args.at(1))
  }

  if instr-type == "R" {
    parsed.insert("rs1", r2)
    parsed.insert("rs2", r3)
  } else if instr-type in ("I", "L", "S") {
    parsed.insert("rs1", r2)
    parsed.insert("imm", imm)
  } else if instr-type == "B" {
    parsed.insert("rs2", r2)
    parsed.insert("imm", imm)
  } else if instr-type in ("U", "J") {
    parsed.insert("imm", imm)
  }
  
  return parsed
}

#let format-imm(imm, sections) = {
  let value = 0
  for section in sections {
    if type(section) == int {
      let bit = imm.bit-rshift(section).bit-and(1)
      value = value.bit-lshift(1).bit-or(bit)
    } else {
      let (start, end) = section
      if start > end {
        (start, end) = (end, start)
      }
      let span = end - start + 1
      let mask = (1).bit-lshift(span) - 1
      let bits = imm.bit-rshift(start).bit-and(mask)
      value = value.bit-lshift(span).bit-or(bits)
    }
  }
  return value
}

#let instr-to-bin(instruction) = {
  let (instr, ..args) = instruction.split(" ")
  args = args.join("").split(",").map(a => a.trim())
  instr = lower(instr)
  let parts = ()
  let parts-names = ()
  let opcode = opcodes.at(instr)
  let instr-type = types.at(instr)
  let code = 0

  let parsed-args = parse-args(instr, instr-type, args)
  let imm = parsed-args.at("imm", default: none)

  let add-part(code, parts, part, bits) = {
    code = code.bit-lshift(bits).bit-or(part)
    parts.push(int-to-bin(part, bits))
    return (code, parts)
  }

  if instr in ("slli", "srli", "srai") {
    let f7 = funct7.at(instr)
    (code, parts) = add-part(code, parts, f7, 7)
    parts-names.push("funct7")
    let v = format-imm(imm, ((4,0),))
    (code, parts) = add-part(code, parts, v, 5)
    parts-names.push("shamt")
    
  } else if instr-type == "R" {
    let v = funct7.at(instr)
    parts-names.push("funct7")
    (code, parts) = add-part(code, parts, v, 7)
  } else if instr-type in ("I", "L") {
    let v = format-imm(imm, ((11,0),))
    (code, parts) = add-part(code, parts, v, 12)
    parts-names.push("imm[11:0]")
  } else if instr-type == "S" {
    let v = format-imm(imm, ((11, 5),))
    (code, parts) = add-part(code, parts, v, 7)
    parts-names.push("imm[11:5]")
  } else if instr-type == "B" {
    let v = format-imm(imm, (12, (10, 5)))
    (code, parts) = add-part(code, parts, v, 7)
    parts-names.push("imm[12|10:5]")
  } else if instr-type == "U" {
    let v = format-imm(imm, ((19, 0),))
    (code, parts) = add-part(code, parts, v, 20)
    parts-names.push("imm[31:12]")
  } else if instr-type == "J" {
    let v = format-imm(imm, (20, (10, 1), 11, (19, 12)))
    (code, parts) = add-part(code, parts, v, 20)
    parts-names.push("imm[20|10:1|11|19:12]")
  }
  
  if instr-type in ("R", "S", "B") {
    (code, parts) = add-part(code, parts, parsed-args.rs2, 5)
    parts-names.push("rs2")
  }
  if instr-type in ("R", "I", "L", "S", "B") {
    (code, parts) = add-part(code, parts, parsed-args.rs1, 5)
    parts-names.push("rs1")
    (code, parts) = add-part(code, parts, funct3.at(instr), 3)
    parts-names.push("funct3")
  }
  
  if instr-type in ("R", "I", "L", "U", "J") {
    (code, parts) = add-part(code, parts, parsed-args.rd, 5)
    parts-names.push("rd")
  } else if instr-type == "S" {
    let v = format-imm(imm, ((4, 0),))
    (code, parts) = add-part(code, parts, v, 5)
    parts-names.push("imm[4:0]")
  } else if instr-type == "B" {
    let v = format-imm(imm, ((4, 1), 11))
    (code, parts) = add-part(code, parts, v, 5)
    parts-names.push("imm[4:1|11]")
  }

  (code, parts) = add-part(code, parts, opcode, 7)
  parts-names.push("opcode")

  return (
    instruction: instruction,
    parts: parts,
    names: parts-names,
    code: code
  )
}

#let color-instruction(line) = {
  let res = ""
  if line.starts-with("0x") {
    let (addr, ..rest) = line.split(" ")
    line = rest.join(" ").trim()
    res += text(raw(addr), fill: colors.address) + " "
  }
  
  if ":" in line {
    let label
    (label, line) = line.split(":")
    line = line.trim()
    res += text(raw(label), fill: colors.label) + ": "
    if line.len() == 0 {
      return res
    }
  }
  let (op, ..args) = line.split(" ")
  args = args.join("").split(",").map(a => a.trim())
  args = args.map(a => raw(a))
  
  res += text(raw(op), fill: colors.op) + " "
  res += args.join(", ")

  return res
}

#let explain-instruction-encoding(instruction) = {
  let encoded = instr-to-bin(instruction)
  let parts = encoded.parts
  let names = encoded.names
  let code = encoded.code
  parts.zip(names).map(((v, n)) => $underbrace(#raw(v), #n)$).join(" ")
  [ \= #raw(instruction) \= 0x#raw(zfill(str(code, base: 16), 8))]
}

#let replace-label-refs(instr, labels, addr) = {
  let (op, ..args) = instr.split(" ")
  args = args.join("").split(",").map(a => a.trim())
  if op in ("beq", "bne", "blt", "bge", "bltu", "bgeu", "jal") {
    let label = args.last()
    let target = labels.at(label)
    let offset = target - addr
    args.last() = str(offset)
  }
  return op + " " + args.join(", ")
}

#let riscv-to-bin(riscv, start-address: 0) = {
  let src = riscv
  if type(riscv) == raw or type(riscv) == content {
    src = riscv.text
  }

  let lines = src.split("\n")
                 .map(l => l.split("#").first().trim())
                 .filter(l => l.len() != 0)
  let labels = (:)
  let addr = start-address

  let lines2 = ()
  let addresses = ()
  for line in lines {
    let set-addr = none
    if line.starts-with("0x") {
      let rest
      (set-addr, ..rest) = line.split(" ")
      line = rest.join(" ")
      addr = eval(set-addr)
    }
    
    if ":" in line {
      let label
      (label, line) = line.split(":")
      labels.insert(label, addr)
      line = line.trim()
      if line.len() == 0 {
        continue
      }
    }
    lines2.push(line.trim())
    addresses.push(addr)
    if set-addr == none {
      addr += 4
    }
  }

  let cells = ()
  for (line, addr) in lines2.zip(addresses) {
    let instr = replace-label-refs(line, labels, addr)
    let encoded = instr-to-bin(instr)
    let addr-txt = "0x" + zfill(str(addr, base: 16), 8)
    let code-txt = "0x" + zfill(str(encoded.code, base: 16), 8)
    addr-txt = raw(addr-txt)
    addr-txt = text(addr-txt, fill: colors.address)
    

    let label = ""
    if addr in labels.values() {
      label = labels.pairs().filter(p => p.last() == addr).first().first()
      label += ":"
    }
    label = text(raw(label), fill: colors.label)
    cells.push(addr-txt)
    cells.push(label)
    cells.push(raw(code-txt))
    cells.push(color-instruction(line))
  }
  table(
    align: left + horizon,
    stroke: none,
    columns: 4,
    inset: 5pt,
    ..cells
  )
}