use bio
use std

const strdup = {s
    var buf = std.slalloc(s.len)
    std.slcp(buf, s)
    -> buf
}

const upcase = {s
    var i
    var buf = std.mksb()
    for i = 0; i < s.len; i++
        std.sbputc(buf, std.toupper((s[i] : char)))
    ;;
    -> std.sbfin(buf)
}

const readline = {line
    var cmt = std.strfind(line, ";")

    match cmt
    | `std.Some (pos): line = line[:pos]
    | `std.None:
    ;;

    if line.len == 0
        -> (`std.None, [][:])
    ;;

    var parts = std.strtok(line)
    if std.isspace((line[0] : char))
        -> (`std.None, parts)
    else
        -> (`std.Some(parts[0]), parts[1:])
    ;;
}

type literal = union
    `Number int
    `Register int
    `Label byte[:]
    `DerefAddr
    `Keypress
    `AddressRegister
    `DelayTimer
    `SoundTimer
    `Invalid
;;

const getnum = {s
    -> std.intparsebase((s[1:] : byte[:]), 16)
}

const parselit = {s
    if s.len < 1
        -> `Invalid
    elif s[0] == ('#' : byte)
        match getnum(s)
        |`std.Some(n): -> `Number(n)
        |`std.None: -> `Invalid
        ;;
    elif s[0] == ('v' : byte) || s[0] == ('V' : byte)
        match getnum(s)
        |`std.Some(n): -> `Register(n)
        |`std.None: -> `Invalid
        ;;
    elif std.isdigit((s[0] : char))
        match std.intparsebase((s : byte[:]), 10)
        |`std.Some(n): -> `Number(n)
        |`std.None: -> `Label(s)
        ;;
    elif s[0] == ('[' : byte)
        -> `DerefAddr
    elif s[0] == ('K' : byte)
        -> `Keypress
    elif s[0] == ('I' : byte)
        -> `AddressRegister
    elif s[0] == ('D' : byte)
        -> `DelayTimer
    elif s[0] == ('S' : byte)
        -> `SoundTimer
    else
        -> `Label(s)
    ;;
}

const emitinstr = {p
    var b = `std.None
    var mark = `std.None

    if p.len < 1
        -> `std.Ok(b, mark)
    ;;

    match upcase(p[0])
    | "CLS":
        b = `std.Some([0x00, 0xe0])
    | "RET":
        b = `std.Some([0x00, 0xee])
    | "SYS":
        if p.len < 2
            -> `std.Err("SYS missing subroutine address")
        ;;
        match parselit(p[1])
        | `Number(n): b = `std.Some([0x00 | (n & 0x0f00), (n & 0xff)])
        | `Label(l): b = `std.Some([0x00, 0x00])
            mark = `std.Some(strdup(l))
        | _:
            -> `std.Err(std.fmt("invalid argument to SYS: {}", p[1]))
        ;;
    | "CALL":
        if p.len < 2
            -> `std.Err("CALL missing subroutine address")
        ;;
        match parselit(p[1])
        | `Number(n): b = `std.Some([0x20 | (n & 0x0f00), (n & 0xff)])
        | `Label(l): b = `std.Some([0x20, 0x00])
            mark = `std.Some(strdup(l))
        | _:
            -> `std.Err(std.fmt("invalid argument to CALL: {}", p[1]))
        ;;
    | "JP":
        if p.len < 2
            -> `std.Err("JP missing address")
        elif p.len == 2
            match parselit(p[1])
            | `Number(n): b = `std.Some([0x10 | (n & 0x0f00), (n & 0xff)])
            | `Label(l): b = `std.Some([0x10, 0x00])
                mark = `std.Some(strdup(l))
            | _:
                -> `std.Err(std.fmt("invalid argument to JP: {}", p[1]))
            ;;
        else
            match (parselit(p[1]), parselit(p[2]))
            | (`Register(0), `Number(n)):
                b = `std.Some([0xb0 | (n & 0x0f00), (n & 0xff)])
            | (`Register(0), `Label(l)): b = `std.Some([0xb0, 0x00])
                mark = `std.Some(strdup(l))
            | _:
                -> `std.Err(std.fmt("invalid arguments to JP: {}, {}", p[1], p[2]))
            ;;
        ;;
    | "SE":
        if p.len < 3
            -> `std.Err("SE missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`Register(m), `Number(n)):
            b = `std.Some([0x30 | (m & 0x0f), (n & 0xff)])
        | (`Register(m), `Register(n)):
            b = `std.Some([0x50 | (m & 0x0f), (n & 0x0f) << 1])
        | _:
            -> `std.Err(std.fmt("invalid arguments to SE: {}, {}", p[1], p[2]))
        ;;
    | "SNE":
        if p.len < 3
            -> `std.Err("SNE missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`Register(m), `Number(n)):
            b = `std.Some([0x40 | (m & 0x0f), (n & 0xff)])
        | (`Register(m), `Register(n)):
            b = `std.Some([0x90 | (m & 0x0f), (n & 0x0f) << 1])
        | _:
            -> `std.Err(std.fmt("invalid arguments to SNE: {}, {}", p[1], p[2]))
        ;;
    | "SKP":
        if p.len < 2
            -> `std.Err("SKP missing address")
        ;;
        match parselit(p[1])
        | `Register(n): b = `std.Some([0xe0 | (n & 0x0f), 0x9e])
        | _:
            -> `std.Err(std.fmt("invalid argument to SKP: {}", p[1]))
        ;;
    | "SKNP":
        if p.len < 2
            -> `std.Err("SKNP missing address")
        ;;
        match parselit(p[1])
        | `Register(n): b = `std.Some([0xe0 | (n & 0x0f), 0xa1])
        | _:
            -> `std.Err(std.fmt("invalid argument to SKNP: {}", p[1]))
        ;;
    | "LD":
        if p.len < 3
            -> `std.Err("LD missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`Register(n), `Keypress): b = `std.Some([0xf0 | (n & 0x0f), 0x0a])
        | (`Register(m), `Number(n)): b = `std.Some([0x60 | (m & 0x0f), (n & 0xff)])
        | (`Register(m), `Label(l)): b = `std.Some([0x60 | (m & 0x0f), 00])
            mark = `std.Some(strdup(l))
        | (`Register(m), `Register(n)): b = `std.Some([0x80 | (n & 0x0f), (m & 0x0f) << 1])
        | (`Register(n), `DelayTimer): b = `std.Some([0xf0 | (n & 0x0f), 0x07])
        | (`DelayTimer, `Register(n)): b = `std.Some([0xf0 | (n & 0x0f), 0x15])
        | (`SoundTimer, `Register(n)): b = `std.Some([0xf0 | (n & 0x0f), 0x18])
        | (`AddressRegister, `Number(n)): b = `std.Some([0xa0 | (n & 0x0f), (n & 0xff)])
        | (`AddressRegister, `Label(l)): b = `std.Some([0xa0, 0x00])
            mark = `std.Some(strdup(l))
        | (`DerefAddr, `Register(n)): b = `std.Some([0xf0 | (n & 0x0f), 0x55])
        | (`Register(n), `DerefAddr): b = `std.Some([0xf0 | (n & 0x0f), 0x65])
        | _:
            -> `std.Err(std.fmt("invalid argument to LD: {}", p[1]))
        ;;
    | "ADD":
        if p.len < 3
            -> `std.Err("ADD missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`AddressRegister, `Register(n)): b = `std.Some([0xf0 | (n & 0x0f), 0x1e])
        | (`Register(m), `Number(n)): b = `std.Some([0x70 | (m & 0x0f), n & 0xff])
        | (`Register(m), `Label(l)): b = `std.Some([0x70 | (m & 0x0f), 0x00])
            mark = `std.Some(strdup(l))
        | (`Register(m), `Register(n)): b = `std.Some([0x80 | (m & 0x0f), ((n & 0x0f) << 1) | 0x04])
        | _:
            -> `std.Err(std.fmt("invalid argument to ADD: {}", p[1]))
        ;;
    | "SUB":
        if p.len < 3
            -> `std.Err("SUB missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`Register(m), `Register(n)): b = `std.Some([0x80 | (m & 0x0f), ((n & 0x0f) << 1) | 0x05])
        | _:
            -> `std.Err(std.fmt("invalid argument to SUB: {}", p[1]))
        ;;
    | "SUBN":
        if p.len < 3
            -> `std.Err("SUBN missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`Register(m), `Register(n)): b = `std.Some([0x80 | (m & 0x0f), ((n & 0x0f) << 1) | 0x07])
        | _:
            -> `std.Err(std.fmt("invalid argument to SUBN: {}", p[1]))
        ;;
    | "OR":
        if p.len < 3
            -> `std.Err("OR missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`Register(m), `Register(n)): b = `std.Some([0x80 | (m & 0x0f), ((n & 0x0f) << 1) | 0x01])
        | _:
            -> `std.Err(std.fmt("invalid argument to OR: {}", p[1]))
        ;;
    | "AND":
        if p.len < 3
            -> `std.Err("AND missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`Register(m), `Register(n)): b = `std.Some([0x80 | (m & 0x0f), ((n & 0x0f) << 1) | 0x02])
        | _:
            -> `std.Err(std.fmt("invalid argument to AND: {}", p[1]))
        ;;
    | "XOR":
        if p.len < 3
            -> `std.Err("XOR missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`Register(m), `Register(n)): b = `std.Some([0x80 | (m & 0x0f), ((n & 0x0f) << 1) | 0x03])
        | _:
            -> `std.Err(std.fmt("invalid argument to XOR: {}", p[1]))
        ;;
    | "SHR":
        if p.len < 2
            -> `std.Err("SHR missing arguments")
        ;;
        match parselit(p[1])
        | `Register(n): `std.Some([0x80 | (n & 0x0f), 0x06])
        | _:
            -> `std.Err(std.fmt("invalid argument to SHR: {}", p[1]))
        ;;
    | "SHL":
        if p.len < 2
            -> `std.Err("SHL missing arguments")
        ;;
        match parselit(p[1])
        | `Register(n): `std.Some([0x80 | (n & 0x0f), 0x0e])
        | _:
            -> `std.Err(std.fmt("invalid argument to SHL: {}", p[1]))
        ;;
    | "BCD":
        if p.len < 2
            -> `std.Err("BCD missing arguments")
        ;;
        match parselit(p[1])
        | `Register(n): `std.Some([0xf0 | (n & 0x0f), 0x33])
        | _:
            -> `std.Err(std.fmt("invalid argument to BCD: {}", p[1]))
        ;;
    | "RND":
        if p.len < 3
            -> `std.Err("RND missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]))
        | (`Register(m), `Number(n)): b = `std.Some([0xc0 | (m & 0x0f), n & 0xff])
        | _:
            -> `std.Err(std.fmt("invalid argument to RND: {}", p[1]))
        ;;
    | "DRW":
        if p.len < 4
            -> `std.Err("DRW missing arguments")
        ;;
        match (parselit(p[1]), parselit(p[2]), parselit(p[3]))
        | (`Register(m), `Register(n), `Number(q)): b = `std.Some([0xc0 | (m & 0x0f), (n & 0x0f) << 1 | (q & 0x0f)])
        | _:
            -> `std.Err(std.fmt("invalid argument to DRW: {}", p[1]))
        ;;
    | _:
        -> `std.Err(std.fmt("unrecognized instruction {}", p[0]))
    ;;

    -> `std.Ok((b, mark))
}

const die_emit = {msg, lc
    std.die(std.fmt("{} on line {}\n", msg, lc))
}

const die_rewrite = {l, lc
    const s = "ref to undefined label '{}' on line {}\n"
    std.die(std.fmt(s, l, lc))
}


/* convert `arr` to a byte[:] slice */
const itob = {arr
    var i
    var ret = std.slalloc(arr.len)
    for i = 0; i < arr.len; i++;
        ret[i] = ((arr[i] & 0xff) : byte)
    ;;
    -> ret
}

const main = {
    var fd = bio.mkfile(std.In, bio.Rd)
    var src = bio.byline(fd)
    var lc = 0

    var offset = 0
    var buf = [][:]
    var q = [][:]
    var labels = std.mkht()
    var res

    for line : src ;
        match readline(line)
        | (`std.None, i): res = emitinstr(i)
        | (`std.Some (label), i): res = emitinstr(i)
            std.htput(labels, strdup(label), offset)
        ;;

        match res
        | `std.Err(msg): die_emit(msg, lc)
        | `std.Ok((bytes, marker)):
            match marker
            | `std.Some(label):
                q = std.slpush(&q, (label, offset, lc))
            | `std.None:
            ;;
            match bytes
            | `std.Some(arr):
                buf = std.sljoin(&buf, arr[:])
                offset += 2
            | `std.None:
            ;;
        ;;

        lc++
    ;;

    for p : q ;
        match p
        | (l, buf_pos, linum):
            match std.htget(labels, l)
            | `std.None: die_rewrite(l, linum)
            | `std.Some(addr):
                buf[buf_pos] |= addr & 0x0f00
                buf[buf_pos + 1] |= addr & 0xff
            ;;
        ;;
    ;;

    std.write(std.Out, itob(buf))

    bio.free(fd)
}