#!/usr/bin/env python3
import os
import sys
import shutil
import subprocess
import urllib.request
import tomllib
import difflib
from pathlib import Path, PurePath, PureWindowsPath
from zipfile import ZipFile
from argparse import ArgumentParser
from multiprocessing import Pool

EXTRACT_DIR = 'extract'
EXPECTED_DIR = 'expected'
OUTPUT_DIR = 'output'
CMD_BASE = ['-iwad', 'miniwad.wad', '-nodraw', '-noblit', '-nosound', '-nogui']


def download(url):
    response = urllib.request.urlopen(url)
    with open('tmp', 'wb') as stream:
        stream.write(response.read())
    print('downloaded ' + url)
    return 'tmp'


def extract(zipname, filename):
    with ZipFile(zipname, 'r') as zf:
        with zf.open(filename) as stream:
            with open(Path(EXTRACT_DIR, filename), 'wb') as out:
                out.write(stream.read())


def download_and_extract(record):
    if not Path(EXTRACT_DIR, record['wad']).exists():
        zipname = download(record['wad_url'])
        extract(zipname, record['wad'])
        if 'deh' in record:
            extract(zipname, record['deh'])
        os.remove(zipname)

    if not Path(EXTRACT_DIR, record['demo']).exists():
        zipname = download(record['demo_url'])
        extract(zipname, record['demo'])
        os.remove(zipname)


def build_command_line(record):
    cmd = []
    cmd += CMD_BASE
    cmd += ['-file', record['wad']]
    if 'deh' in record:
        cmd += ['-deh', record['deh']]
    if 'gameversion' in record:
        cmd += ['-gameversion', record['gameversion']]
    cmd += ['-timedemo', record['demo']]
    if 'statdump' in record:
        cmd += ['-statdump', Path(OUTPUT_DIR, record['statdump'])]
    if 'levelstat' in record:
        cmd += ['-levelstat']
    return cmd


def call_port(args):
    source_port, record = args
    cmd = [source_port] + build_command_line(record)

    if 'levelstat' in record:
        base_dir = record['levelstat']
        Path(base_dir).mkdir(exist_ok=True)

        subprocess.run(cmd, cwd=base_dir, capture_output=True)

        path = Path(base_dir, 'levelstat.txt')
        if not path.exists():
            with open(path, 'w') as f:
                pass
        shutil.copyfile(path, Path(OUTPUT_DIR, record['levelstat']))
        shutil.rmtree(base_dir)
    else:
        subprocess.run(cmd, capture_output=True)


def compare_output(record):
    if 'levelstat' in record:
        name = record['levelstat']
    else:
        name = record['statdump']

    with open(Path(EXPECTED_DIR, name), 'r') as f:
        expected = f.read().splitlines()
    with open(Path(OUTPUT_DIR, name), 'r') as f:
        actual = f.read().splitlines()

    # Ignore whitespace (diff -w)
    def normalize(line):
        return ''.join(line.split())

    diff = list(difflib.unified_diff(
        [normalize(line) for line in expected],
        [normalize(line) for line in actual],
        lineterm=''
    ))

    if diff:
        print("name: " + name)
        sys.stdout.writelines('\n'.join(diff) + '\n')
        return True
    return False


def run_program(args):
    source_port = PurePath(args.source_port)
    if type(source_port) is PureWindowsPath:
        source_port = source_port.with_suffix('.exe')
    source_port = Path(source_port).resolve()
    if not source_port.exists():
        sys.exit("Doom port is not found.")

    with open('config.toml', 'rb') as stream:
        config = tomllib.load(stream)

    extract('miniwad.zip', 'miniwad.wad')

    for record in config['demo']:
        download_and_extract(record)

    os.environ['SDL_VIDEODRIVER'] = 'dummy'
    os.environ['DOOMWADDIR'] = str(Path(Path().resolve(), EXTRACT_DIR))

    with Pool(processes=args.jobs) as pool:
        pool.map(call_port, [(source_port, record) for record in config['demo']])

    differences = False

    for record in config['demo']:
        if compare_output(record):
            differences = True

    if differences:
        sys.exit(1)
    else:
        print("Success.")


if __name__ == "__main__":
    parser = ArgumentParser(description="Execute demos for Doom port in a batch.")
    parser.add_argument('--jobs', dest='jobs', default=1, type=int, help="Set the number of jobs.")
    parser.add_argument('--port', dest='source_port', default="doom", type=str, help="Path to Doom port.")
    args = parser.parse_args()
    run_program(args)
