import sys
import struct

# sorry not sorry
P = struct.pack
U = struct.unpack

# notes:
# SCMDraft 2 requires, at bare minimum:
scmdraft_req = b"ERA ~DIM ~TILE~UNIT~PUNI~IOWN~THG2~SPRP~WAV ~MASK~STR ~OWNR~SIDE~FORC".split(b"~")

# Starcraft 1.18+ Melee requires, at bare minimum:
# absolutely required:
#   VER TYPE OWNR SIDE COLR ERA DIM MTXM UNIT THG2 SPRP FORC
# not needed at all:
#   IVER IOWN ISOM TILE DD2 
# not required for Melee mode (non-UMS):
#   PUPx UPGx PUNI MASK MRGN WAV PTEx UNIx TECx TRIG MBRF UPRP UPUS SWNM
# special cases:
# VCOD partly: 34 19 00 00 00 00 00 00 00 00 00 00 00 00 00 00 1F 29 7D A6 D7 B0 00 BB CC 31 24 ED 17 4C 13 0B
# STR  kinda (can be empty for melee)
melee_req = b"VER ~TYPE~VCOD~OWNR~SIDE~COLR~ERA ~DIM ~MTXM~UNIT~THG2~STR ~SPRP~FORC".split(b"~")

# TODO: try more things to add to this list.
ums_req = melee_req + b"MRGN~PTEx~PUPx~TECx~TRIG~UNIx~UPGx~UPRP".split(b"~")

both_req = set(melee_req + scmdraft_req)

known_keys = [ # in the order we want to write them.
    # boring filetype stuff
    b"TYPE",
    b"VER ",
    b"IVER",
    b"IVE2",
    b"VCOD",

    b"SPRP", # scenario name and description
    b"COLR", # player colors
    b"FORC", # player forces and settings
    b"SIDE", # player races
    b"IOWN", # editor player types
    b"OWNR", # ingame player types

    b"DIM ", # map dimensions
    b"ERA ", # tileset
    b"MTXM", # ingame tiles
    b"TILE", # editor tiles
    b"ISOM", # editor isometric tiles
    b"MASK", # fog of war per tile
    b"DD2 ", # editor doodads
    b"UNIT", # placed units
    b"THG2", # placed sprites (e.g. doodad)

    b"MRGN", # locations

    b"UNIS", # unit settings (a bunch of stuff, including names)
    b"UNIx", # broodwar override

    b"PUNI", # build unit restrictions

    b"PTEC", # tech restrictions
    b"PTEx", # broodwar override

    b"UPGR", # upgrade restrictions (and defaults, max levels...)
    b"PUPx", # broodwar override

    b"TECS", # tech settings (costs and times)
    b"TECx", # broodwar override

    b"UPGS", # upgrade settings (costs and times)
    b"UPGx", # broodwar override

    b"UPRP", # properties of spawn-with-properties
    b"UPUS", # booleans of which properties are used

    b"MBRF", # mission briefings (triggers without conditions)
    b"TRIG", # triggers

    b"SWNM", # switch names
    b"WAV ", # wav paths (editor only)
    b"STR ", # string data
]

skippable_keys = set(known_keys).symmetric_difference(both_req)

vcod_valid = b"\
\x34\x19\xCA\x77\x99\xDC\x68\x71\x0A\x60\xBF\xC3\xA7\xE7\x75\xA7\
\x1F\x29\x7D\xA6\xD7\xB0\x3A\xBB\xCC\x31\x24\xED\x17\x4C\x13\x0B\
\x65\x20\xA2\xB7\x91\xBD\x18\x6B\x8D\xC3\x5D\xDD\xE2\x7A\xD5\x37\
\xF6\x59\x64\xD4\x63\x9A\x12\x0F\x43\x5C\x2E\x46\xE3\x74\xF8\x2A\
\x08\x6A\x37\x06\x37\xF6\xD6\x3B\x0E\x94\x63\x16\x45\x67\x5C\xEC\
\xD7\x7B\xF7\xB7\x1A\xFC\xD4\x9E\x73\xFA\x3F\x8C\x2E\xC0\xE1\x0F\
\xD1\x74\x09\x07\x95\xE3\x64\xD7\x75\x16\x68\x74\x99\xA7\x4F\xDA\
\xD5\x20\x18\x1F\xE7\xE6\xA0\xBE\xA6\xB6\xE3\x1F\xCA\x0C\xEF\x70\
\x31\xD5\x1A\x31\x4D\xB8\x24\x35\xE3\xF8\xC7\x7D\xE1\x1A\x58\xDE\
\xF4\x05\x27\x43\xBA\xAC\xDB\x07\xDC\x69\xBE\x0A\xA8\x8F\xEC\x49\
\xD7\x58\x16\x3F\xE5\xDB\xC1\x8A\x41\xCF\xC0\x05\x9D\xCA\x1C\x72\
\xA2\xB1\x5F\xA5\xC4\x23\x70\x9B\x84\x04\xE1\x14\x80\x7B\x90\xDA\
\xFA\xDB\x69\x06\xA3\xF3\x0F\x40\xBE\xF3\xCE\xD4\xE3\xC9\xCB\xD7\
\x5A\x40\x01\x34\xF2\x68\x14\xF8\x38\x8E\xC5\x1A\xFE\xD6\x3D\x4B\
\x53\x05\x05\xFA\x34\x10\x45\x8E\xDD\x91\x69\xFE\xAF\xE0\xEE\xF0\
\xF3\x48\x7E\xDD\x9F\xAD\xDC\x75\x62\x7A\xAC\xE5\x31\x1B\x62\x67\
\x20\xCD\x36\x4D\xE0\x98\x21\x74\xFB\x09\x79\x71\x36\x67\xCD\x7F\
\x77\x5F\xD6\x3C\xA2\xA2\xA6\xC6\x1A\xE3\xCE\x6A\x4E\xCD\xA9\x6C\
\x86\xBA\x9D\x3B\xB5\xF4\x76\xFD\xF8\x44\xF0\xBC\x2E\xE9\x6E\x29\
\x23\x25\x2F\x6B\x08\xAB\x27\x44\x7A\x12\xCC\x99\xED\xDC\xF2\x75\
\xC5\x3C\x38\x7E\xF7\x1C\x1B\xC5\xD1\x2D\x94\x65\x06\xC9\x48\xDD\
\xBE\x32\x2D\xAC\xB5\xC9\x32\x81\x66\x4A\xD8\x34\x35\x3F\x15\xDF\
\xB2\xEE\xEB\xB6\x04\xF6\x4D\x96\x35\x42\x94\x9C\x62\x8A\xD3\x61\
\x52\xA8\x7B\x6F\xDC\x61\xFC\xF4\x6C\x14\x2D\xFE\x99\xEA\xA4\x0A\
\xE8\xD9\xFE\x13\xD0\x48\x44\x59\x80\x66\xF3\xE3\x34\xD9\x8D\x19\
\x16\xD7\x63\xFE\x30\x18\x7E\x3A\x9B\x8D\x0F\xB1\x12\xF0\xF5\x8C\
\x0A\x78\x58\xDB\x3E\x63\xB8\x8C\x3A\xAA\xF3\x8E\x37\x8A\x1A\x2E\
\x5C\x31\xF9\xEF\xE3\x6D\xE3\x7E\x9B\xBD\x3E\x13\xC6\x44\xC0\xB9\
\xBC\x3A\xDA\x90\xA4\xAD\xB0\x74\xF8\x57\x27\x89\x47\xE6\x3F\x37\
\xE4\x42\x79\x5A\xDF\x43\x8D\xEE\xB4\x0A\x49\xE8\x3C\xC3\x88\x1A\
\x88\x01\x6B\x76\x8A\xC3\xFD\xA3\x16\x7A\x4E\x56\xA7\x7F\xCB\xBA\
\x02\x5E\x1C\xEC\xB0\xB9\xC9\x76\x1E\x82\xB1\x39\x3E\xC9\x57\xC5\
\x19\x24\x38\x4C\x5D\x2F\x54\xB8\x6F\x5D\x57\x8E\x30\xA1\x0A\x52\
\x6D\x18\x71\x5E\x13\x06\xC3\x59\x1F\xDC\x3E\x62\xDC\xDA\xB5\xEB\
\x1B\x91\x95\xF9\xA7\x91\xD5\xDA\x33\x53\xCE\x6B\xF5\x00\x70\x01\
\x7F\xD8\xEE\xE8\xC0\x0A\xF1\xCE\x63\xEB\xB6\xD3\x78\xEF\xCC\xA5\
\xAA\x5D\xBC\xA4\x96\xAB\xF2\xD2\x61\xFF\xEA\x9A\xA8\x6A\xED\xA2\
\xBD\x3E\xED\x61\x39\xC1\x82\x92\x16\x36\x23\xB1\xB0\xA0\x24\xE5\
\x05\x9B\xA7\xAA\x0D\x12\x9B\x33\x83\x92\x20\xDA\x25\xB0\xEC\xFC\
\x24\xD0\x38\x23\xFC\x95\xF2\x74\x80\x73\xE5\x19\x97\x50\x7D\x44\
\x45\x93\x44\xDB\xA2\xAD\x1D\x69\x44\x14\xEE\xE7\x2C\x7F\x87\xFF\
\x38\x9E\x32\xF1\x4D\xBC\x29\xDA\x42\x27\x26\xFE\xC1\xD2\x2B\xA9\
\xF6\x42\x7A\x0E\xCB\xE8\x7C\xD1\x0F\x5B\xEC\x56\x69\xB7\x61\x31\
\xB4\x6D\xF9\x25\x40\x34\x79\x6D\xFA\x53\xA7\x0B\xFA\xA4\x82\xCE\
\xC3\x45\x49\x61\x0D\x45\x2C\x8F\x28\x49\x60\xF7\xF3\x7D\xC9\x1E\
\x0F\xD0\x89\xC1\x26\x52\xF8\xD3\x4D\x8F\x35\x14\xBA\x9D\x5F\x0B\
\x07\xA9\x4A\x00\xF7\x22\x26\x2F\x3E\x67\xFB\x1F\xA1\x9C\x11\xC6\
\x69\x4F\x5D\x66\x58\x34\x15\x90\x6C\xE5\x54\x46\xAF\x5F\x63\xD6\
\x8A\x0C\x95\xDF\xBD\x0D\xE4\xAF\xBF\x40\x40\x4C\xA3\xF6\x51\x71\
\x29\xED\x26\xF8\x85\x28\x22\xD5\xBF\xBE\xCF\xFA\x28\xC5\x7F\x51\
\xB8\x06\x63\x07\xEC\xBD\x8F\x29\xFA\x55\x7E\x71\x1A\x40\x32\x66\
\xE8\xD4\xDE\x9D\xD4\x5E\xFC\x93\x7A\x3D\xD5\x3B\xCD\x75\x2E\x80\
\x0A\x4F\x74\x87\x1B\xCC\x8F\xEA\x9A\xA9\xDB\x7C\x16\x53\xE5\xEF\
\xAB\x78\xC1\x6E\xA4\x72\x89\x5A\x98\x2C\x70\x50\xFB\xA1\xDF\x1F\
\x6B\xB7\xD9\x44\x07\x80\x82\x56\xFD\xBF\xC0\x83\x0E\x49\xD0\x5B\
\x1E\x68\x6A\x0E\x9A\xC2\x0B\x2F\x8E\x43\xA0\xE1\x99\x0C\xF6\xB2\
\xE0\x7A\x1C\x5E\x2C\xC8\xA0\x45\x3C\x0B\xE9\x88\xAC\xB9\x96\xC6\
\x74\xAE\x83\x2A\xBB\x13\xFA\x65\xEB\x4F\x1F\xA6\xB0\x8A\x8A\xE1\
\x81\xE9\xB8\xB9\xD5\x55\x15\x4E\x45\xF2\xAD\x9B\x3E\xC2\x35\x7E\
\x5F\x92\x2E\x72\xB6\x5B\x68\x23\x6E\xC6\x45\x0E\xE9\x3B\x87\xD4\
\xF4\x41\xC0\xE3\xA8\x05\x44\xBE\xE4\x0F\x8A\x13\x1A\xC4\x37\xF4\
\x5A\x40\x55\xEF\x9D\x79\x1D\x4B\x4A\x79\x3A\x9C\x76\x85\x37\xCC\
\x82\x3D\x0F\xB6\x60\xA6\x93\x7E\xBD\x5C\xC2\xC4\x72\xC7\x7F\x90\
\x4D\x1B\x96\x10\x13\x05\x68\x68\x35\xC0\x7B\xFF\x46\x85\x43\x2A\
\x01\x04\x05\x06\x02\x01\x05\x02\x00\x03\x07\x07\x05\x04\x06\x03\
"

mrgn_default = bytearray(b"\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0" * 255)
# set up the "Anywhere" location correctly. this is always location 64.
mrgn_default[63*20 +  9] = 0x10
mrgn_default[63*20 + 13] = 0x10
mrgn_default[63*20 + 16] = 0x03
mrgn_default = bytes(mrgn_default)

def readit(f):
    f.seek(0, 2)
    fs = f.tell()
    f.seek(0)

    src = {}

    while f.tell() < fs:
        k = f.read(4)

        try:
            name = k.decode('ascii')
        except UnicodeDecodeError:
            print("! skipping invalid section name:", k)
            size = U("<l", f.read(4))[0]
            if size > 0:
                f.read(size)
            continue

        size = U("<l", f.read(4))[0]
        if size < 0:
            print("gotcha!", size)
            return
        print(name, size)
        section = f.read(size)
        if len(section) != size:
            print("bad section size:", len(section))
            print("expected:", size)
            return

        #if k not in b"ERA ~DIM ~MTXM~THG2~UNIT".split(b"~"): continue

        if k not in src:
            src[k] = []
        src[k].append(section)

    return src

def makestringtable(strings):
    count = len(strings)

    if len(strings) > 0 and type(strings[0]) != bytes:
        strings = [str(s).encode("CP1252") for s in strings]

    flat = b"\0" + b"\0".join(strings)

    offsets = []
    o = count * 2 + 3
    for s in strings:
        offsets += [o]
        o += len(s) + 1

    out = P("<H", count)
    for o in offsets:
        out += P("<H", o)
    out += flat

    return out

def rewrite(src):
    if b"STR " not in src:
        return

    ok = b"MRGN~MBRF~TRIG~SPRP~FORC~WAV ~UNIx~SWNM".split(b"~")

    refs = []

    def mark(k, v, i, ki=0, extra=None):
        num = v[i] + v[i+1]*256
        if num != 0:
            print("ref {:04X} found at {}[{}]".format(num, k, i))
            ref = [k, ki, i, num, extra]
            refs.append(ref)

    for k, a in src.items():
        if k not in ok:
            continue

        if k == b"MRGN":
            for ki, v in enumerate(a):
                for i in range(0, len(v), 20):
                    mark(k, v, i+16, ki)

        if k == b"SPRP":
            for ki, v in enumerate(a):
                mark(k, v, 0, ki)
                mark(k, v, 2, ki)

        if k == b"FORC":
            for ki, v in enumerate(a):
                # TODO: handle truncated FORC.
                mark(k, v, 8, ki)
                mark(k, v, 10, ki)
                mark(k, v, 12, ki)
                mark(k, v, 14, ki)

        if k == b"WAV ":
            for ki, v in enumerate(a):
                for i in range(0, 512, 2):
                    if v[i] != 0 and v[i+1] != 0:
                        mark(k, v, i, ki)

        if k == b"UNIx":
            for ki, v in enumerate(a):
                uc = 228
                stringstart = uc + uc*4 + uc*2 + uc + uc*2 + uc*2 + uc*2
                for i in range(stringstart, stringstart + uc*2, 2):
                    if v[i] != 0 and v[i+1] != 0:
                        mark(k, v, i, ki)

        if k == b"SWNM":
            for ki, v in enumerate(a):
                for i in range(0, 1024, 4):
                    if v[i] != 0 and v[i+1] != 0:
                        mark(k, v, i, ki)

        if k == b"TRIG" or k == b"MBRF":
            for ki, v in enumerate(a):
                # trigger size: (20)*16 + (32)*64 + (4+28) = 2400
                for i in range(0, len(v), 2400):
                    # iterate through each of the 64 actions per trigger.
                    for j in range(20*16, 2400, 32):
                        atype = v[i+j+26]
                        if atype == 0:
                            continue
                        #full = U("<l", v[i+j+4:i+j+4+4])[0]
                        #if full != 0: print("DEBUG:", full, atype)
                        #full = U("<l", v[i+j+8:i+j+8+4])[0]
                        #if full != 0: print("DEBUG:", full, atype)
                        mark(k, v, i+j+4, ki)
                        mark(k, v, i+j+8, ki, "wav")

    strings = []
    wavstrings = []

    # dereference.
    STR = src[b"STR "][0] # TODO: rename
    for ref in refs:
        k, ki, i, num, extra = ref
        offset = STR[num*2] + STR[num*2+1]*256
        null = STR.find(b'\0', offset)
        string = STR[offset:null]
        #print("found string at {}[{}] + {:04X} (ref {:04X}):".format(k, ki, offset, num))
        #print(string.decode('CP949', errors='replace'))
        if string not in strings: # FIXME: potentially slow.
            strings.append(string)
            if extra == "wav":
                wavstrings.append(string)

    # create new table.
    src[b"STR "][0] = makestringtable(strings)

    # update references to point at the new table.
    for ref in refs:
        k, ki, i, num, extra = ref
        # TODO: don't repeat ourselves.
        offset = STR[num*2] + STR[num*2+1]*256
        null = STR.find(b'\0', offset)
        string = STR[offset:null]

        assert string in strings # FIXME: potentially slow.

        si = strings.index(string) # FIXME: potentially slow.
        ba = bytearray(src[k][ki]) # FIXME: potentially slow.
        ba[i+0] = (1 + si) % 256
        ba[i+1] = (1 + si) // 256
        src[k][ki] = bytes(ba) # FIXME: potentially slow.

    # fix wav table.
    # NOTE: since we're dealing with raw .chk and not .mpq,
    # you still need to reimport the wavs with an MPQ editor!
    wavs = bytearray(b'\0\0\0\0' * 512)
    for i, string in enumerate(wavstrings):
        assert string in strings # FIXME: potentially slow.
        si = strings.index(string) # FIXME: potentially slow.
        wavs[i*4+0] = (1 + si) % 256
        wavs[i*4+1] = (1 + si) // 256

    src[b"WAV "] = [bytes(wavs)]

def section(k, v):
    assert type(k) == bytes, type(k)
    assert type(v) == bytes, type(v)
    assert len(k) == 4, k
    return k + P("<L", len(v)) + v

def check_dim(v):
    w, h = U("<HH", v)
    if w not in (64, 96, 128, 192, 256):
        return False
    if h not in (64, 96, 128, 192, 256):
        return False
    return True

def writeit(f, src):
    wroteonce = {}
    def W(k, v):
        wroteonce[k] = True
        f.write(section(k, v))

    # set some essentials that there's no need to change.
    # (unless you're writing expansionless maps for some reason)
    src[b"TYPE"] = [b"RAWB"]
    src[b"VER "] = [b"\315\0"]
    src[b"IVER"] = [b"\12\0"]
    src[b"IVE2"] = [b"\13\0"]
    src[b"VCOD"] = [vcod_valid]

    if b"TILE" not in src:
        src[b"TILE"] = src[b"MTXM"]

    w, h = 0, 0

    for k, a in src.items():
        # TODO: consider removing this since we check in the write func anyway.
        assert type(k) == bytes
        assert type(a) == list
        assert len(k) == 4

        if k == b"DIM ": # TODO: move out of loop.
            a = [v for v in a if check_dim(v)]
            w, h = U("<HH", a[0])

    placed_starts = [False] * 12
    k = b"UNIT"
    if k in src:
        for ki, v in enumerate(src[k]):
            for i in range(0, len(v), 36):
                unit = v[i:i+36]
                uid = unit[8] + unit[9]*256
                if uid != 214:
                    continue
                pi = unit[16]
                if pi >= 12:
                    print("# ignoring start location for player", pi)
                placed_starts[pi] = True

    k = b"ERA "
    if k in src:
        for ki, v in enumerate(src[k]):
            era = U("<H", v)[0]
            era &= 7
            src[k][ki] = P("<H", era)

    assert b"DIM " in src # TODO: try to guess map dimensions.

    if b"COLR" not in src:
        src[b"COLR"] = [b"\0\1\2\3\4\5\6\7"]

    # SCMDraft 2 seems to require this to be the correct length.
    if b"MASK" not in src or len(src[b"MASK"][0]) != w * h:
        src[b"MASK"] = [b"\xFF" * (w * h)]

    if b"STR " not in src:
        src[b"STR "] = [b""]

    # SCMDraft 2 calls this "player settings" in its errors.
    if b"OWNR" not in src:
        b = b""
        for i in range(12):
            b += b"\6" if placed_starts[i] else b"\0"
        src[b"OWNR"] = [b]
        assert len(b) == 12, b

    # SCMDraft 2 calls this "player settings" in its errors.
    if b"SIDE" not in src:
        src[b"SIDE"] = [b"\5\5\5\5\5\5\5\5\7\7\7\4"]

    if b"FORC" not in src:
        b = P("<BBBBBBBBHHHHBBBB",
            0,0,0,0,0,0,0,0, # player->force mapping
            0,0,0,0, # strings for forces
            1,1,1,1, # bitmasks: randstart, ally, allyvic, vis, unused...
        )
        src[b"FORC"] = [b]
        assert len(b) == 20, b

    if b"IOWN" not in src:
        # FIXME: possibly invalid.
        b = b""
        for i in range(8):
            b += b"\6" if placed_starts[i] else b"\0"
        b += b"\0\0\0\7"
        src[b"IOWN"] = [b]
        assert len(b) == 12, b

    if b"MRGN" not in src:
        src[b"MRGN"] = [mrgn_default]

    if b"THG2" not in src:
        src[b"THG2"] = [b""]

    if b"SPRP" not in src:
        src[b"SPRP"] = [b"\0\0\0\0"]

    if b"WAV " not in src:
        src[b"WAV "] = [b"\0\0\0\0" * 512]

    #if b"ISOM" not in src:
    #    size = (w // 2 + 1) * (h + 1) * 4
    #    src[b"ISOM"] = [b'\0' * size]

    if b"PUNI" not in src:
        src[b"PUNI"] = [b"\1" * 5700]

    done = {}
    for k in known_keys:
        if k in src:
            for v in src[k]:
                W(k, v)
            done[k] = True

    for k, v in src.items():
        if not done[k]:
            print("# deferred write:", k)
            W(k, v)

    for k in known_keys:
        if k not in skippable_keys and k not in wroteonce:
            print("# never wrote section:", k)

    for k in scmdraft_req:
        if k not in wroteonce:
            print("# map might not open in SCMDraft 2. missing", k)
    for k in melee_req:
        if k not in wroteonce:
            print("# map might not open in StarCraft (Melee). missing", k)
    for k in ums_req:
        if k not in wroteonce:
            print("# map might not open in StarCraft (Use Map Settings). missing", k)

def run(args):
    for fn in args:
        with open(fn, 'rb') as f:
            src = readit(f)

        if src is None:
            print("! failed to rewrite", fn) # FIXME: reuse of terminology.
        ofn = fn.replace('.chk', '') + '.x.chk'
        assert fn != ofn, ofn
        rewrite(src)
        with open(ofn, 'wb') as f:
            writeit(f, src)

if __name__ == '__main__':
    ret = run(sys.argv[1:])
    sys.exit(ret)