#!/usr/bin/env python3
# findStrings.py - locate the game's bit-7-terminated text strings inside regions that the
# disassembly still renders as raw bytes, and emit annotations/06_strings.json (labels + text
# data types + a comment holding the text).  Run after tools/build.py.
import os, re, json, glob

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT  = os.path.join(ROOT, "disassembly")
ANN  = os.path.join(ROOT, "annotations")

MIN_LEN = 4                     # shortest run of characters worth calling a string


def dataAddresses(unit, dataMap):
    """addresses the disassembler classified as data (from build/dataMap.json)"""
    addrs = set()
    for lo, hi in dataMap.get(unit, []):
        addrs.update(range(int(lo, 16), int(hi, 16) + 1))
    return addrs


# The game's text is uppercase PETSCII/ASCII; restricting the alphabet keeps graphics data
# (which is full of bytes that happen to fall in the printable range) from matching.
ALLOWED = set("ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,!?'\"-:;()/%$*+=&#")
SKIP_UNITS = ("boot/title_bitmap_A000", "boot/title_screen_9C00", "boot/title_colorram_9800")


def isText(b):
    return chr(b) in ALLOWED or b == 0x0D          # $0D separates the lines of a multi-line message


def looksLikeMessage(text):
    text = text.replace("\r", " ")
    letters = sum(1 for c in text if c.isalpha())
    if len(text) < 5 or letters < 3:
        return False
    if (letters + text.count(" ")) / len(text) < 0.7:
        return False
    for i in range(len(text) - 3):                 # runs like "OOOO" are graphics, not text
        if text[i].isalpha() and text[i] * 4 == text[i:i+4]:
            return False
    if len(text) < 8 and not all(c.isalpha() or c == " " for c in text):
        return False
    return True


def nameFor(text, used):
    words = re.findall(r"[A-Za-z0-9]+", text.title())
    base = "msg" + "".join(words)[:28]
    if len(base) < 6:
        base = "msgText"
    name = base
    n = 2
    while name in used:
        name = f"{base}{n}"
        n += 1
    used.add(name)
    return name


def main():
    index = {}
    for line in open(os.path.join(OUT, "INDEX.txt")):
        m = re.match(r"^(\S+)\s+\$([0-9A-F]{4})-\$([0-9A-F]{4})", line)
        if m:
            index[m.group(1)] = (int(m.group(2), 16), int(m.group(3), 16))

    dataMap = json.load(open(os.path.join(OUT, "build", "dataMap.json")))
    unitLabels, dataTypes, notes = {}, {}, {}
    used = set()
    total = 0
    for unit, (start, end) in sorted(index.items()):
        srcPath = os.path.join(OUT, unit + ".s")
        binPath = os.path.join(OUT, "build", unit.replace("/", "_") + ".orig.bin")
        if not (os.path.exists(srcPath) and os.path.exists(binPath)) or unit in SKIP_UNITS:
            continue
        data = open(binPath, "rb").read()
        dataSet = dataAddresses(unit, dataMap)
        a = start
        while a <= end:
            if a not in dataSet or not isText(data[a - start]):
                a += 1
                continue
            b = a
            text = ""
            terminated = False
            while b <= end and b in dataSet:
                v = data[b - start]
                if isText(v):
                    text += chr(v)
                    b += 1
                elif isText(v & 0x7F) and v & 0x80:
                    text += chr(v & 0x7F)
                    b += 1
                    terminated = True
                    break
                else:
                    break
            if terminated and looksLikeMessage(text):
                lines = [l for l in text.split("\r") if l != ""] or [text]
                name = nameFor(lines[0], used)
                unitLabels.setdefault(unit, {})[f"{a:04X}"] = name
                dataTypes.setdefault(unit, {})[f"{a:04X}"] = ["text", b - a]
                if len(lines) == 1:
                    block = [f'{name} - message text "{text}" (the last character carries bit 7 as the terminator).']
                else:
                    block = [f"{name} - message text, {len(lines)} lines separated by $0D; the last character carries bit 7:"]
                    block += [f'    "{l}"' for l in lines]
                notes[f"{a:04X}"] = {"unit": unit, "block": block}
                total += 1
                a = b
            else:
                a = max(b, a + 1)
    json.dump({"_comment": "generated by tools/findStrings.py - bit-7-terminated strings found in data areas",
               "unitLabels": unitLabels, "unitDataTypes": dataTypes, "notes": notes},
              open(os.path.join(ANN, "06_strings.json"), "w"), indent=1)
    print(f"{total} strings in {len(unitLabels)} units")
    for unit, m in sorted(unitLabels.items()):
        print(f"  {unit}: {len(m)}")


main()