#!/usr/bin/env python3
"""极简 RISC-V 汇编器：仅支持本项目测试所需的指令子集。

用法: python tools/rv_asm.py <input.asm> <output.hex>
同时生成 <output.hex>.lst 反汇编/编码清单，便于核对。
支持：add/sub/and/or/xor/sll/srl/sra/slt/sltu、
      mul/mulh/mulhsu/mulhu、
      addi/slti/sltiu/xori/ori/andi、slli/srli/srai、
      lw/sw/sb、beq/bne/blt/bge/bltu/bgeu、lui、jal/jalr、ecall、标签、# 注释
伪指令：print（= lui t6,LED_ADDR>>12 + sw a0,0(t6)，向 MMIO 的 LED 寄存器写入 a0）、
halt（= addi a7,zero,10 + ecall），各占两个字。
示例：python tools/rv_asm.py testbench/programs/isa_test.asm out.hex
"""

import re
import sys

REGS = {
    "zero": 0,
    "ra": 1,
    "sp": 2,
    "gp": 3,
    "tp": 4,
    "t0": 5,
    "t1": 6,
    "t2": 7,
    "s0": 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,
}


def reg(t):
    t = t.strip()
    if t in REGS:
        return REGS[t]
    m = re.fullmatch(r"x(\d+)", t)
    if m:
        return int(m.group(1))
    raise SystemExit("未知寄存器: " + t)


def num(t):
    return int(t.strip(), 0)


def enc_r(f7, rs2, rs1, f3, rd):
    return (f7 << 25) | (rs2 << 20) | (rs1 << 15) | (f3 << 12) | (rd << 7) | 0x33


def enc_i(imm, rs1, f3, rd, op=0x13):
    return ((imm & 0xFFF) << 20) | (rs1 << 15) | (f3 << 12) | (rd << 7) | op


def enc_s(imm, rs2, rs1, f3):
    imm &= 0xFFF
    return (
        ((imm >> 5) << 25)
        | (rs2 << 20)
        | (rs1 << 15)
        | (f3 << 12)
        | ((imm & 0x1F) << 7)
        | 0x23
    )


def enc_b(imm, rs2, rs1, f3):
    imm &= 0x1FFF
    return (
        (((imm >> 12) & 1) << 31)
        | (((imm >> 5) & 0x3F) << 25)
        | (rs2 << 20)
        | (rs1 << 15)
        | (f3 << 12)
        | (((imm >> 1) & 0xF) << 8)
        | (((imm >> 11) & 1) << 7)
        | 0x63
    )


def enc_u(imm, rd):
    return ((imm & 0xFFFFF) << 12) | (rd << 7) | 0x37


def enc_j(imm, rd):
    imm &= 0x1FFFFF
    return (
        (((imm >> 20) & 1) << 31)
        | (((imm >> 1) & 0x3FF) << 21)
        | (((imm >> 11) & 1) << 20)
        | (((imm >> 12) & 0xFF) << 12)
        | (rd << 7)
        | 0x6F
    )


RTYPE = {
    "add": (0x00, 0),
    "sub": (0x20, 0),
    "sll": (0x00, 1),
    "slt": (0x00, 2),
    "sltu": (0x00, 3),
    "xor": (0x00, 4),
    "srl": (0x00, 5),
    "sra": (0x20, 5),
    "or": (0x00, 6),
    "and": (0x00, 7),
    "mul": (0x01, 0),
    "mulh": (0x01, 1),
    "mulhsu": (0x01, 2),
    "mulhu": (0x01, 3),
}
ITYPE = {"addi": 0, "slti": 2, "sltiu": 3, "xori": 4, "ori": 6, "andi": 7}
SHIFT = {"slli": 1, "srli": 5, "srai": 5}
BTYPE = {"beq": 0, "bne": 1, "blt": 4, "bge": 5, "bltu": 6, "bgeu": 7}
# MMIO：LED 数据寄存器地址，需与 RTL 顶层 parameter LED_ADDR 保持一致
LED_ADDR = 0x1000

# 双字伪指令：print 向 MMIO 的 LED 写寄存器写入 a0；halt 停机
PSEUDO_SIZE = {"print": 2, "halt": 2}


def assemble(lines):
    labels, items, addr = {}, [], 0
    for raw in lines:
        line = raw.split("#")[0].strip()
        if not line:
            continue
        while True:
            m = re.match(r"^([A-Za-z_.$][\w.$]*)\s*:\s*(.*)$", line)
            if not m:
                break
            labels[m.group(1)] = addr
            line = m.group(2).strip()
            if not line:
                break
        if not line:
            continue
        items.append((addr, line))
        addr += 4 * PSEUDO_SIZE.get(line.split()[0].lower(), 1)

    out = []
    for addr, line in items:
        parts = re.split(r"[\s,]+", line)
        mn, a = parts[0].lower(), parts[1:]
        if mn in RTYPE:
            f7, f3 = RTYPE[mn]
            w = enc_r(f7, reg(a[2]), reg(a[1]), f3, reg(a[0]))
        elif mn in ITYPE:
            im = num(a[2])
            if not (-2048 <= im <= 2047):
                raise SystemExit("I 型立即数超出 12 位有符号范围: " + line)
            w = enc_i(im, reg(a[1]), ITYPE[mn], reg(a[0]))
        elif mn in SHIFT:
            w = enc_i(num(a[2]) & 0x1F, reg(a[1]), SHIFT[mn], reg(a[0]))
        elif mn == "lui":
            w = enc_u(num(a[1]), reg(a[0]))
        elif mn in ("lw", "sw", "sb"):
            m = re.fullmatch(r"(-?\w+)\((\w+)\)", a[1])
            if not m:
                raise SystemExit("内存操作数格式应为 imm(rs1): " + line)
            im = num(m.group(1))
            if not (-2048 <= im <= 2047):
                raise SystemExit("访存偏移超出 12 位有符号范围: " + line)
            if mn == "lw":
                w = enc_i(im, reg(m.group(2)), 2, reg(a[0]), 0x03)
            elif mn == "sw":
                w = enc_s(im, reg(a[0]), reg(m.group(2)), 2)
            else:  # sb
                w = enc_s(im, reg(a[0]), reg(m.group(2)), 0)
        elif mn in BTYPE:
            tgt = labels[a[2]] if a[2] in labels else num(a[2])
            w = enc_b(tgt - addr, reg(a[1]), reg(a[0]), BTYPE[mn])
        elif mn == "jal":
            tgt = labels[a[1]] if a[1] in labels else num(a[1])
            w = enc_j(tgt - addr, reg(a[0]))
        elif mn == "jalr":
            w = enc_i(num(a[2]), reg(a[1]), 0, reg(a[0]), 0x67)
        elif mn == "ecall":
            w = 0x00000073
        elif mn == "print":
            # 写 MMIO 的 LED 寄存器：lui t6, LED_ADDR>>12 ; sw a0, 0(t6)
            w = [
                enc_u(LED_ADDR >> 12, reg("t6")),
                enc_s(0, reg("a0"), reg("t6"), 2),
            ]
        elif mn == "halt":
            w = [enc_i(10, 0, 0, 17), 0x00000073]
        else:
            raise SystemExit("未知助记符: " + mn)
        out.append((addr, line, w if isinstance(w, list) else [w]))
    return out


def main():
    src, dst = sys.argv[1], sys.argv[2]
    with open(src, encoding="utf-8") as f:
        prog = assemble(f.readlines())
    with open(dst, "w", encoding="ascii") as f:
        for _, _, ws in prog:
            for w in ws:
                f.write("%08X\n" % w)
    with open(dst + ".lst", "w", encoding="utf-8") as f:
        for addr, line, ws in prog:
            f.write("%04X: %s  %s\n" % (addr, " ".join("%08X" % w for w in ws), line))
    print("已生成 %s（%d 个字）" % (dst, sum(len(ws) for _, _, ws in prog)))


main()