import argparse
from collections import OrderedDict
import glob
import os
import shutil
from construct import *
from tqdm import tqdm

setglobalstringencoding(None)

ScrapFile = Struct(
    "path" / PascalString(Int32ul),
    "size" / Int32ul,
    "offset" / Int32ul,
    "data" / OnDemandPointer(this.offset, Bytes(this.size)),
)
DummyFile = Struct("path" / PascalString(Int32ul), "size" / Int32ul, "offset" / Int32ul)

PackedHeader = Struct(
    Const(b"BFPK"), Const(b"\0\0\0\0"), "files" / PrefixedArray(Int32ul, ScrapFile)
)
DummyHeader = Struct(
    Const(b"BFPK"), Const(b"\0\0\0\0"), "files" / PrefixedArray(Int32ul, DummyFile)
)
parser = argparse.ArgumentParser(description="Unpack and Repack .packed files")
parser.add_argument(
    "-u", "--unpack", action="store_true", help="unpack file to 'extracted' directory"
)
parser.add_argument(
    "-r", "--repack", action="store_true", help="repack file from 'extracted' directory"
)

parser.add_argument(
    "--reset", action="store_true", default=False, help="restore backup"
)

parser.add_argument(
    "scrap_dir",
    metavar="Scrapland Directory",
    type=str,
    default=".",
    help="Scrapland installation directory",
)
options = parser.parse_args()
scrap_dir = os.path.abspath(options.scrap_dir)

if options.reset:
    print("Restoring Backups and removing extracted folder...")
    for packed_file in glob.glob(os.path.join(scrap_dir, "*.packed.bak")):
        outfile = os.path.basename(packed_file)
        orig_filename = outfile[:-4]
        if os.path.isfile(outfile):
            print("deleting", orig_filename)
            os.remove(orig_filename)
            print("moving", outfile, "->", orig_filename)
            shutil.move(outfile, orig_filename)
        target_folder = os.path.join("extracted", os.path.basename(orig_filename))
        print("deleting", target_folder)
        shutil.rmtree(target_folder)
    if os.path.isdir("extracted"):
        input("Press enter to remove rest of extracted folder")
        shutil.rmtree("extracted")
    exit("Done!")

if not (options.unpack or options.repack):
    parser.print_help()
    exit()
pstatus = ""
if options.unpack:
    if os.path.isdir("extracted"):
        print("Removing extracted folder")
        shutil.rmtree("extracted")
    for packed_file in glob.glob(os.path.join(scrap_dir, "*.packed")):
        os.chdir(scrap_dir)
        BN = os.path.basename(packed_file)
        target_folder = os.path.join("extracted", os.path.basename(packed_file))
        os.makedirs(target_folder, exist_ok=True)
        os.chdir(target_folder)
        print("Unpacking {}".format(os.path.basename(packed_file)))
        with open(packed_file, "rb") as pkfile:
            data = PackedHeader.parse_stream(pkfile)
            print("Offset:", hex(pkfile.tell()))
            for file in tqdm(data.files, ascii=True):
                folder, filename = os.path.split(file.path)
                if folder:
                    os.makedirs(folder, exist_ok=True)
                with open(file.path, "wb") as outfile:
                    outfile.write(file.data())
        print("\r" + " " * len(pstatus) + "\r", end="", flush=True)
        os.chdir(scrap_dir)

if options.unpack and options.repack:
    input(
        "Press enter to rebuild *.packed files from folders in 'extracted' dir..."
    )  # noqa
    pass


def file_gen(files, offset=0):
    for real_path, size, path in files:
        file = dict(path=path, offset=offset, size=size)
        yield file
        offset += file["size"]


def make_header(files, offset=0):
    files_list = list(file_gen(files, offset))
    return DummyHeader.build(dict(files=files_list))


if options.repack:
    for folder in glob.glob(os.path.join(scrap_dir, "extracted", "*.packed")):
        data = []
        filename = os.path.join(scrap_dir, os.path.basename(folder))
        for root, folders, files in os.walk(folder):
            for file in sorted(files):
                file = os.path.join(root, file)
                rel_path = bytes(
                    file.replace(folder, "").replace("\\", "/").lstrip("/"),
                    "windows-1252",
                )
                size = os.stat(file).st_size
                data.append((file, size, rel_path))
        print("Found {} files for {}".format(len(data), filename))
        offset = len(make_header(data))
        print("Writing", filename)
        header = make_header(data, offset)
        with open(filename, "wb") as outfile:
            outfile.write(header)
            for file, size, rel_path in tqdm(data, ascii=True):
                outfile.write(open(file, "rb").read())
print("Done!")