#!/usr/bin/env python3
# render.py - render a region of a memory image as reassemblable, annotated ca65 source.
#
# Inputs (all optional except image/region): a code-address set, a label map (addr->name), an
# annotation map (addr->{"name","line","block","routine"}) and a data-type map (addr->type).
import json, os, textwrap
from m6502 import decode, formatOperand, IMM, IMP, ACC, ZP, ZPX, ZPY, ABS, ABX, ABY, IND, IZX, IZY, REL, BRANCHES
from hwSymbols import c64Symbols, DRIVE

COMMENT_COL = 40


MARKERS = ("In:", "Out:", "Called", "*", "-", "$", "|")


def isStructured(line):
    """a line that carries its own layout (a field, a bullet, a hand made table row)"""
    stripped = line.lstrip()
    if line[:1].isspace():                       # indented: the author laid it out on purpose
        return True
    # three or more spaces in a row mean a hand made table; two are just sentence spacing
    return line.startswith(MARKERS) or stripped.startswith(("* ", "- ", "$")) or "   " in stripped


def wrapLines(lines, width=100):
    """Wrap comment text to width.  Consecutive prose lines are joined into one paragraph first, so a
    long line written by hand does not leave a two word orphan before the next line."""
    out = []
    paragraph = []

    def flush():
        if not paragraph:
            return
        text = " ".join(p.strip() for p in paragraph)
        out.extend(textwrap.wrap(text, width=width, break_long_words=False, break_on_hyphens=False) or [""])
        paragraph.clear()

    for line in lines:
        if isStructured(line) or not line.strip():
            flush()
            if len(line) <= width:
                out.append(line)
            else:
                indent = "    " if line.startswith(("In:", "Out:", "Called")) else ""
                out.extend(textwrap.wrap(line, width=width, subsequent_indent=indent,
                                         break_long_words=False, break_on_hyphens=False) or [""])
        else:
            paragraph.append(line)
    flush()
    return out


def isPrintable(b):
    return 0x20 <= b <= 0x7E


def asciiOf(bs):
    return "".join(chr(b) if isPrintable(b) else "." for b in bs)


def escapeStr(s):
    return s.replace("\\", "\\\\").replace('"', '\\"')


class RendererT:
    def __init__(self, mem, codeSet, labels, notes, hwSyms, dataTypes=None, zpNames=None, useKernal=False):
        self.useKernal = useKernal
        self.mem       = mem
        self.code      = codeSet
        self.labels    = dict(labels)       # addr -> name (global, across files)
        self.notes     = notes              # addr(int) -> dict
        self.hw        = hwSyms
        self.dataTypes = dataTypes or {}    # addr -> ("text"|"addr"|"byte"|"word", length)
        self.zpNames   = zpNames or {}
        self.insnStarts = {}                # addr -> InsnT for code in rendered region
        self.altComments = {}               # addr -> extra comment when a referenced address has other variant names

    # ----- naming -----
    def nameFor(self, addr):
        if addr in self.labels:
            return self.labels[addr]
        if addr in self.hw:
            return self.hw[addr]
        return None

    def midInsn(self, addr):
        # (base, offset) if addr lies inside an instruction of the region being rendered
        for back in (1, 2):
            base = addr - back
            if base in self.insnStarts and self.insnStarts[base].length > back:
                return base, back
        return None

    def operandLabel(self, addr):
        # label for an operand address; in-file mid-instruction references become label+offset
        mid = self.midInsn(addr)
        if mid:
            bn = self.labels.get(mid[0])
            if bn:
                return f"{bn}+{mid[1]}"
        n = self.nameFor(addr)
        if n:
            return n
        if addr < 0x100:
            return self.zpNames.get(addr)
        return None

    # ----- rendering -----
    def renderRegion(self, start, end, title, header, externsOut):
        out = []
        mem = self.mem
        # pass 1: decode instructions in region
        a = start
        while a < end:
            if a in self.code:
                insn = decode(mem, a)
                if insn.isLegal() and a + insn.length <= end:
                    self.insnStarts[a] = insn
                    a += insn.length
                    continue
            a += 1
        # collect labels defined here and referenced externals
        used = set()
        defined = set()
        for da, dtSpec in self.dataTypes.items():
            kind, length = dtSpec[0], dtSpec[1]
            if kind == "addr" and start <= da < end:
                for i in range(0, length - 1, 2):
                    used.add(mem[da + i] | (mem[da + i + 1] << 8))
        for a, insn in self.insnStarts.items():
            if insn.mode in (ZP, ZPX, ZPY, ABS, ABX, ABY, IND, IZX, IZY, REL):
                t = insn.target
                if start <= t < end:
                    mid = self.midInsn(t)
                    if mid and mid[0] not in self.labels:
                        self.labels[mid[0]] = f"L_{mid[0]:04X}"
                used.add(t)
        out.append(f"; {'=' * 76}")
        for line in title.split("\n"):
            out.append(f"; {line}")
        out.append(f"; {'=' * 76}")
        for line in header:
            out.append(f"; {line}")
        out.append("")
        out.append('        .setcpu "6502"')
        out.append('        .include "c64.inc"')
        if self.useKernal:
            out.append('        .include "kernal.inc"')
        out.append('        .include "zeropage.inc"')
        out.append("")
        # externs: every referenced labelled address outside [start,end)
        externs = []
        for t in sorted(used):
            if start <= t < end or t < 0x100 or t in self.hw:
                continue
            n = self.labels.get(t)
            if not n:
                n = f"D_{t:04X}"
                self.labels[t] = n
            externs.append(f"{n:<24} := ${t:04X}")
        # also mid-instruction / mid-data references that resolve to label+off outside region
        if externs:
            out.append("; ---- references to code/data outside this file ----")
            out.extend(externs)
            out.append("")
        externsOut.extend(externs)
        # ---- table of contents: every labelled routine that carries a header comment
        toc = []
        for a in sorted(self.insnStarts):
            if not (start <= a < end):
                continue
            note = self.notes.get(a, {})
            head = note.get("routine")
            if not head or a not in self.labels:
                continue
            summary = head[0]
            if summary.startswith(self.labels[a]):
                summary = summary[len(self.labels[a]):].lstrip(" -")
            cut = summary.find(". ")                    # the table of contents shows one sentence
            if cut > 0:
                summary = summary[:cut + 1]
            toc.append((a, self.labels[a], summary))
        if toc:
            out.append("; Contents")
            out.append("; --------")
            for a, name, summary in toc:
                prefix = f"   ${a:04X}  {name:<28} "
                wrapped = wrapLines([prefix + summary], width=118)
                out.append("; " + wrapped[0])
                for cont in wrapped[1:]:
                    out.append("; " + " " * len(prefix) + cont.strip())
            out.append("")
        out.append(f"        .org ${start:04X}")
        out.append("")
        # pass 2: emit
        a = start
        pendingBytes = []
        def flushBytes(atAddr):
            nonlocal pendingBytes
            if not pendingBytes:
                return
            bs = pendingBytes
            base = atAddr - len(bs)
            i = 0
            while i < len(bs):
                chunk = bs[i:i+8]
                hexs = ",".join(f"${b:02X}" for b in chunk)
                line = f"        .byte   {hexs}"
                out.append(f"{line:<{COMMENT_COL}}; {base+i:04X}  {asciiOf(chunk)}")
                i += 8
            pendingBytes = []
        labelsHere = set()
        while a < end:
            note = self.notes.get(a, {})
            name = self.labels.get(a)
            if a in self.insnStarts:
                flushBytes(a)
                insn = self.insnStarts[a]
                if note.get("routine") or note.get("block"):
                    out.append("")
                    out.append(f"; {'-' * 70}")
                    for line in wrapLines((note.get("routine") or []) + (note.get("block") or [])):
                        out.append(f"; {line}")
                    out.append(f"; {'-' * 70}")
                if name:
                    out.append(f"{name}:")
                    labelsHere.add(a)
                opText = self.formatInsn(insn)
                lineText = f"        {insn.mnemonic:<7} {opText}".rstrip()
                comment = note.get("line", "")
                if insn.target is not None and insn.target in self.altComments and not (start <= insn.target < end):
                    comment = (comment + "  " if comment else "") + "(" + self.altComments[insn.target] + ")"
                hexBytes = " ".join(f"{mem[a+i]:02X}" for i in range(insn.length))
                out.append(f"{lineText:<{COMMENT_COL}}; {a:04X}  {comment}".rstrip())
                a += insn.length
                # blank line after unconditional flow end for readability
                if insn.mnemonic in ("rts", "rti", "jmp") and a < end:
                    flushBytes(a)
                    out.append("")
                continue
            # data
            dt = self.dataTypes.get(a)
            if name or note.get("block") or dt or note.get("routine"):
                flushBytes(a)
                if note.get("block") or note.get("routine"):
                    out.append("")
                    for line in wrapLines((note.get("routine") or []) + (note.get("block") or [])):
                        out.append(f"; {line}")
                if name:
                    out.append(f"{name}:")
                    labelsHere.add(a)
            if dt:
                kind, length = dt[0], dt[1]
                length = min(length, end - a)
                # never run a typed block across another label or annotated address - except that an
                # address table may contain a label at offset 1 (the high-byte reference), which becomes an equate
                for k in range(1, length):
                    if (a + k) in self.insnStarts or (a + k) in self.dataTypes or ((a + k) in self.notes and (a + k) not in self.labels):
                        length = k
                        break
                    if (a + k) in self.labels:
                        if kind == "addr" and k % 2 == 1 and name:
                            out.append(f"{self.labels[a + k]:<24} = {name}+{k}")
                            continue
                        length = k
                        break
                if kind == "text":
                    a = self.emitText(out, a, length, note.get("line", ""))
                    continue
                if kind in ("addr", "word") and length < 2:
                    kind = "byte"
                if kind == "addr":
                    for i in range(0, length - 1, 2):
                        t = mem[a+i] | (mem[a+i+1] << 8)
                        n = self.operandLabel(t) or f"${t:04X}"
                        line = f"        .addr   {n}"
                        out.append(f"{line:<{COMMENT_COL}}; {a+i:04X}  {note.get('line','') if i == 0 else ''}".rstrip())
                    a += length - (length % 2)
                    continue
                if kind == "word":
                    for i in range(0, length - 1, 2):
                        t = mem[a+i] | (mem[a+i+1] << 8)
                        line = f"        .word   ${t:04X}"
                        out.append(f"{line:<{COMMENT_COL}}; {a+i:04X}  {note.get('line','') if i == 0 else ''}".rstrip())
                    a += length - (length % 2)
                    continue
                if kind == "grid":
                    width = dt[2] if len(dt) > 2 else 8
                    for i in range(0, length, width):
                        chunk = mem[a+i:a+min(i+width, length)]
                        hexs = ",".join(f"${b:02X}" for b in chunk)
                        c = note.get("line", "") if i == 0 else ""
                        row = f"        .byte   {hexs}"
                        out.append(f"{row:<{COMMENT_COL}}; {a+i:04X}  {asciiOf(chunk)}  {c}".rstrip())
                    a += length
                    continue
                if kind in ("bitmap", "sprite"):
                    perRow  = 3 if kind == "sprite" else 1
                    rows    = 21 if kind == "sprite" else 8
                    unit    = perRow * rows
                    for base in range(0, length, unit):
                        if note.get("line") and base == 0:
                            out.append(f";   {note['line']}")
                        for i in range(0, min(unit, length - base), perRow):
                            # the block may have been clamped short by a label, so never emit past it
                            chunk = mem[a+base+i:a+base+min(i+perRow, length-base)]
                            hexs = ",".join(f"${b:02X}" for b in chunk)
                            art = "".join("#" if b & (0x80 >> k) else "." for b in chunk for k in range(8))
                            row = f"        .byte   {hexs}"
                            out.append(f"{row:<{COMMENT_COL}}; {a+base+i:04X}  {art}")
                        out.append("")
                    a += length
                    continue
                if kind == "byte":
                    # fixed-length row(s)
                    for i in range(0, length, 8):
                        chunk = mem[a+i:a+min(i+8, length)]
                        hexs = ",".join(f"${b:02X}" for b in chunk)
                        line = f"        .byte   {hexs}"
                        c = note.get("line", "") if i == 0 else ""
                        out.append(f"{line:<{COMMENT_COL}}; {a+i:04X}  {asciiOf(chunk)}  {c}".rstrip())
                    a += length
                    continue
            if note.get("line") and not dt:
                flushBytes(a)
                b = mem[a]
                line = f"        .byte   ${b:02X}"
                out.append(f"{line:<{COMMENT_COL}}; {a:04X}  {note['line']}")
                a += 1
                continue
            pendingBytes.append(mem[a])
            a += 1
            # break byte rows at labels / code boundaries
            if a in self.insnStarts or self.nameFor(a) or a in self.notes or a in self.dataTypes:
                flushBytes(a)
        flushBytes(a)
        return out, labelsHere

    def emitText(self, out, a, length, comment):
        mem = self.mem
        bs = mem[a:a+length]
        # split into printable runs and high-bit/ctrl bytes
        parts = []
        cur = ""
        for b in bs:
            if isPrintable(b) and b not in (0x22,):
                cur += chr(b)
            else:
                if cur:
                    parts.append(f'"{cur}"'); cur = ""
                if isPrintable(b & 0x7F) and (b & 0x80) and b != 0xA0 and b != 0xFF:
                    parts.append(f"'{chr(b & 0x7F)}'|$80")
                else:
                    parts.append(f"${b:02X}")
        if cur:
            parts.append(f'"{cur}"')
        # wrap to a few parts per line
        lineParts = []
        lines = []
        curLen = 0
        for p in parts:
            if curLen + len(p) > 60 and lineParts:
                lines.append(",".join(lineParts)); lineParts = []; curLen = 0
            lineParts.append(p); curLen += len(p) + 1
        if lineParts:
            lines.append(",".join(lineParts))
        for i, l in enumerate(lines):
            text = f"        .byte   {l}"
            out.append(f"{text:<{COMMENT_COL}}; {a:04X}  {comment if i == 0 else ''}".rstrip())
        return a + length

    def formatInsn(self, insn):
        m = insn.mode
        if m == IMP:
            return ""
        if m == ACC:
            return "a"
        if m == IMM:
            return f"#${insn.operand:02X}"
        t = insn.target
        lab = self.operandLabel(t)
        if m == REL:
            return lab or f"${t:04X}"
        if m in (ZP, ZPX, ZPY, IZX, IZY):
            s = lab or f"${t:02X}"
        else:
            s = lab or f"${t:04X}"
            if t < 0x100 and m in (ABS, ABX, ABY):
                s = "a:" + s
        if m == ZPX or m == ABX: return s + ",x"
        if m == ZPY or m == ABY: return s + ",y"
        if m == IND: return f"({s})"
        if m == IZX: return f"({s},x)"
        if m == IZY: return f"({s}),y"
        return s


def writeIncludes(outDir, zpNames, zpComments):
    from hwSymbols import KERNAL
    hw = c64Symbols()
    lines = ["; c64.inc - Commodore 64 hardware registers (auto-generated)", ""]
    for a in sorted(hw):
        if a in KERNAL: continue
        lines.append(f"{hw[a]:<24} := ${a:04X}")
    open(os.path.join(outDir, "c64.inc"), "w").write("\n".join(lines) + "\n")
    lines = ["; kernal.inc - C64 KERNAL jump table (only meaningful while the KERNAL ROM is banked in)", ""]
    for a in sorted(KERNAL):
        lines.append(f"{KERNAL[a]:<24} := ${a:04X}")
    open(os.path.join(outDir, "kernal.inc"), "w").write("\n".join(lines) + "\n")
    lines = ["; zeropage.inc - zero-page variables used by Modem Wars (auto-generated from annotations)", ""]
    for a in sorted(zpNames):
        c = zpComments.get(a, "")
        lines.append(f"{zpNames[a]:<24} := ${a:02X}" + (f"    ; {c}" if c else ""))
    open(os.path.join(outDir, "zeropage.inc"), "w").write("\n".join(lines) + "\n")
