#!/usr/bin/env python3
"""
z80fmt - column formatter for Z80 assembly source (z88dk z80asm flavour).

Splits each line into  label | mnemonic | operands | ;comment  and pads the
first three into fixed-width columns with spaces. String and character
literals are parsed properly, so semicolons inside DEFM data are never
mistaken for comments, and EX AF,AF' does not open a quote.

Usage:
    z80fmt.py file.z80 [more.z80 ...]     # write formatted source to stdout
    z80fmt.py -i file.z80                 # rewrite in place
    z80fmt.py --check *.z80               # exit 1 if any file needs formatting
    cat file.z80 | z80fmt.py              # filter mode
"""

import argparse
import sys

MNEMONICS = {
    "adc", "add", "and", "bit", "call", "ccf", "cp", "cpd", "cpdr", "cpi",
    "cpir", "cpl", "daa", "dec", "di", "djnz", "ei", "ex", "exx", "halt",
    "im", "in", "inc", "ind", "indr", "ini", "inir", "jp", "jr", "ld",
    "ldd", "lddr", "ldi", "ldir", "neg", "nop", "or", "otdr", "otir", "out",
    "outd", "outi", "pop", "push", "res", "ret", "reti", "retn", "rl",
    "rla", "rlc", "rlca", "rld", "rr", "rra", "rrc", "rrca", "rrd", "rst",
    "sbc", "scf", "set", "sla", "sll", "sli", "sra", "srl", "sub", "xor",
    # Z180 extras
    "in0", "out0", "mlt", "tst", "tstio", "slp", "otim", "otdm", "otimr",
    "otdmr",
}

DIRECTIVES = {
    "org", "defb", "defw", "defm", "defs", "defc", "equ", "include",
    "binary", "incbin", "public", "extern", "global", "module", "section",
    "align", "defvars", "if", "ifdef", "ifndef", "else", "elif", "endif",
    "end", "db", "dw", "ds", "dm", "dc", "byte", "word", "ascii", "asciz",
    "asciiz", "macro", "endm", "rept", "endr", "local", "title", "line",
    "defgroup", "defarray",
}


def split_comment(line):
    """Return (code, comment). Quote-aware so ';' inside literals is kept."""
    quote = None
    i, n = 0, len(line)
    while i < n:
        c = line[i]
        if quote:
            if c == "\\" and quote == '"' and i + 1 < n:
                i += 2
                continue
            if c == quote:
                quote = None
        elif c == '"':
            quote = '"'
        elif c == "'":
            # A quote right after an identifier char is a register prime
            # (EX AF,AF' / HL'), not the start of a character literal.
            prev = line[i - 1] if i else ""
            if not (prev.isalnum() or prev == "_"):
                quote = "'"
        elif c == ";":
            return line[:i], line[i:]
        i += 1
    return line, ""


def is_opcode(token):
    base = token.rstrip(":").lstrip(".").lower()
    return base in MNEMONICS or base in DIRECTIVES


def split_code(code):
    """Return (label, mnemonic, operands) from the comment-stripped code."""
    if not code.strip():
        return "", "", ""

    label = ""
    if code[0] not in " \t":
        first = code.split(None, 1)[0]
        rest = code.split(None, 1)[1] if len(code.split(None, 1)) > 1 else ""
        if first.endswith(":") or not is_opcode(first):
            label, code = first, rest

    rest = code.strip()
    if not rest:
        return label, "", ""
    parts = rest.split(None, 1)
    mnemonic = parts[0]
    operands = parts[1].strip() if len(parts) > 1 else ""
    return label, mnemonic, operands


def normalize_commas(operands, sep):
    """Rewrite comma spacing outside of literals."""
    out, quote, i, n = [], None, 0, len(operands)
    while i < n:
        c = operands[i]
        if quote:
            out.append(c)
            if c == "\\" and quote == '"' and i + 1 < n:
                out.append(operands[i + 1])
                i += 2
                continue
            if c == quote:
                quote = None
        elif c == '"':
            quote = '"'
            out.append(c)
        elif c == "'":
            prev = operands[i - 1] if i else ""
            if not (prev.isalnum() or prev == "_"):
                quote = "'"
            out.append(c)
        elif c == ",":
            while out and out[-1] == " ":
                out.pop()
            out.append("," + sep)
            j = i + 1
            while j < n and operands[j] in " \t":
                j += 1
            i = j
            continue
        else:
            out.append(c)
        i += 1
    return "".join(out).rstrip()


def pad(text, width, gap=1):
    return text.ljust(width) if len(text) < width else text + " " * gap


def format_line(line, opt):
    raw = line.rstrip("\n").rstrip()
    if not raw.strip():
        return [""]

    stripped = raw.lstrip()

    # Preprocessor lines pass through untouched.
    if stripped.startswith("#"):
        return [stripped]

    # Whole-line comments: keep column-0 banners at column 0, indent the
    # rest to the mnemonic column.
    if stripped.startswith(";"):
        if raw[0] in " \t":
            return [" " * opt.label_width + stripped]
        return [stripped]

    code, comment = split_comment(raw)
    label, mnemonic, operands = split_code(code)

    if opt.case == "upper":
        mnemonic = mnemonic.upper()
    elif opt.case == "lower":
        mnemonic = mnemonic.lower()

    if opt.comma_space and operands:
        operands = normalize_commas(operands, " ")
    elif operands:
        operands = normalize_commas(operands, "")

    lines = []
    has_code = bool(mnemonic or operands or comment)

    if label and has_code and opt.wrap_long_labels and len(label) >= opt.label_width:
        lines.append(label)
        label = ""

    if not has_code:
        lines.append(label)
        return lines

    out = pad(label, opt.label_width) if label else " " * opt.label_width
    out += pad(mnemonic, opt.mnemonic_width) if mnemonic else " " * opt.mnemonic_width

    if comment:
        out += pad(operands, opt.operand_width) if operands else " " * opt.operand_width
        out += comment
    else:
        out += operands

    lines.append(out.rstrip())
    return lines


def format_text(text, opt):
    out = []
    for line in text.splitlines():
        out.extend(format_line(line, opt))
    return "\n".join(out) + ("\n" if text.endswith("\n") or text else "")


def main():
    p = argparse.ArgumentParser(description="Format Z80 assembly into columns.")
    p.add_argument("files", nargs="*", help="source files (default: stdin)")
    p.add_argument("-i", "--in-place", action="store_true")
    p.add_argument("--check", action="store_true",
                   help="exit 1 if any file would change; write nothing")
    p.add_argument("--label-width", type=int, default=16)
    p.add_argument("--mnemonic-width", type=int, default=8)
    p.add_argument("--operand-width", type=int, default=24)
    p.add_argument("--comma-space", action="store_true",
                   help="put one space after operand commas (default: none)")
    p.add_argument("--case", choices=["keep", "upper", "lower"], default="keep",
                   help="case of mnemonics only; symbols are left alone")
    p.add_argument("--wrap-long-labels", action="store_true",
                   help="put over-long labels on their own line")
    opt = p.parse_args()

    if not opt.files:
        sys.stdout.write(format_text(sys.stdin.read(), opt))
        return 0

    changed = False
    for path in opt.files:
        with open(path, "r", encoding="utf-8", errors="surrogateescape") as fh:
            original = fh.read()
        formatted = format_text(original, opt)
        if formatted != original:
            changed = True
            if opt.check:
                print(f"would reformat: {path}", file=sys.stderr)
        if opt.check:
            continue
        if opt.in_place:
            if formatted != original:
                with open(path, "w", encoding="utf-8",
                          errors="surrogateescape") as fh:
                    fh.write(formatted)
        else:
            sys.stdout.write(formatted)

    return 1 if (opt.check and changed) else 0


if __name__ == "__main__":
    sys.exit(main())