"""Tiny assembler that writes Power C object files (.o/.obj) directly.

Syntax (one statement per line, ';' comments):
    label:                 local label
    .def NAME              export NAME at this offset (def kind 1)
    .byte 1,2,$ff          raw bytes
    lda #$12   lda $14   lda $0123,x   lda ($14),y   lda ($14,x)
    bne label  jmp label  jsr label     (jmp/jsr/abs operands on local labels are relocated)
    jsr ext:NAME           reference to an external symbol (absolute)
    lda ext:NAME+1         external symbol + 1
    ldx #<ext:NAME         low byte of an external symbol   (#> for the high byte)
    lda label,x            absolute operand on a local label (relocated)
"""
import re, sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from objfmt import Obj
from pcopt import ENC, SIZE, BRANCH


def num(t):
    t = t.strip()
    if t.startswith('$'):
        return int(t[1:], 16)
    return int(t, 0)


def parse_operand(s):
    """returns (mode, value-or-label, ext) ; ext = (name, kind) for external symbols"""
    s = s.strip()
    if not s:
        return 'imp', None, None
    if s == 'a':
        return 'acc', None, None
    if s.startswith('#'):
        t = s[1:]
        part = None
        if t.startswith('<') or t.startswith('>'):
            part = t[0]
            t = t[1:]
        if t.startswith('ext:'):
            return 'imm', 0, (t[4:].encode(), 2 if part == '<' else 1)
        v = num(t)
        if part == '<':
            v &= 255
        elif part == '>':
            v = v >> 8 & 255
        return 'imm', v, None
    m = re.match(r'^\((.+),x\)$', s)
    if m:
        return 'izx', num(m.group(1)), None
    m = re.match(r'^\((.+)\),y$', s)
    if m:
        return 'izy', num(m.group(1)), None
    idx = ''
    m = re.match(r'^(.*),([xy])$', s)
    if m:
        s, idx = m.group(1), m.group(2)
    if s.startswith('ext:'):
        t = s[4:]
        kind = 0
        m = re.match(r'^(.*)\+1$', t)
        if m:
            t, kind = m.group(1), 4
        return ('abs' + idx if idx else 'abs'), 0, (t.encode(), kind)
    try:
        v = num(s)
        md = 'zp' if v < 256 else 'abs'
        return (md + idx if idx else md), v, None
    except ValueError:
        return ('abs' + idx if idx else 'abs'), s, None          # local label


def assemble(text, base=0):
    lines = []
    for ln in text.splitlines():
        ln = ln.split(';')[0].strip()
        m = re.match(r'^([A-Za-z_]\w*):\s*(.*)$', ln)
        while m:
            lines.append(m.group(1) + ':')
            ln = m.group(2).strip()
            m = re.match(r'^([A-Za-z_]\w*):\s*(.*)$', ln)
        if ln:
            lines.append(ln)
    # pass 1: sizes
    items = []     # ('label',name) ('def',name) ('bytes',[...]) ('ins',mn,md,val,ext)
    for ln in lines:
        if ln.endswith(':'):
            items.append(('label', ln[:-1]))
            continue
        if ln.startswith('.def'):
            items.append(('def', ln.split()[1]))
            continue
        if ln.startswith('.byte'):
            items.append(('bytes', [num(x) for x in ln[5:].split(',')]))
            continue
        parts = ln.split(None, 1)
        mn = parts[0].lower()
        md, val, ext = parse_operand(parts[1] if len(parts) > 1 else '')
        if mn in BRANCH:
            md = 'rel'
        elif mn in ('jmp', 'jsr'):
            md = 'abs'
        elif md == 'zp' and (mn, 'zp') not in ENC and (mn, 'abs') in ENC:
            md = 'abs'
        elif md == 'zpx' and (mn, 'zpx') not in ENC and (mn, 'abx') in ENC:
            md = 'abx'
        elif md == 'zpy' and (mn, 'zpy') not in ENC and (mn, 'aby') in ENC:
            md = 'aby'
        if (mn, md) not in ENC:
            raise ValueError('bad instruction: ' + ln)
        items.append(('ins', mn, md, val, ext))
    labels, defs = {}, []
    pos = base
    for it in items:
        if it[0] == 'label':
            labels[it[1]] = pos
        elif it[0] == 'def':
            defs.append([it[1].encode(), 1, pos])
        elif it[0] == 'bytes':
            pos += len(it[1])
        else:
            pos += SIZE[it[2]]
    code = bytearray()
    reloc, refs = [], []
    for it in items:
        if it[0] == 'bytes':
            code += bytes(it[1])
        elif it[0] == 'ins':
            _, mn, md, val, ext = it
            at = base + len(code)
            code.append(ENC[(mn, md)])
            n = SIZE[md]
            if md == 'rel':
                t = labels[val]
                rel = t - (at + 2)
                if not -128 <= rel <= 127:
                    raise ValueError('branch out of range at %x' % at)
                code.append(rel & 255)
            elif n == 2:
                if ext:
                    refs.append([ext[0], ext[1], at])
                    val = ENC[(mn, md)]
                code.append(val & 255)
            elif n == 3:
                if ext:
                    refs.append([ext[0], ext[1], at])
                    v = ENC[(mn, md)] * 0x101
                elif isinstance(val, str):
                    v = labels[val]
                    reloc.append(at)
                else:
                    v = val
                code += bytes((v & 255, v >> 8 & 255))
    o = Obj()
    o.code = code
    o.reloc = reloc
    o.defs = defs
    o.refs = sorted(refs, key=lambda r: r[2])
    o.res = []
    o.trailing = b''
    return o, labels