#!/usr/bin/env python3

from pathlib import Path
from typing import List, Set
from glob import glob
from shutil import rmtree
from textwrap import dedent

import os
import re
import sys
import argparse

SAFETYHOOK_ROOT = Path(__file__).resolve().parent
PUBLIC_INCLUDE_PATHS = [
    SAFETYHOOK_ROOT / 'include',
    SAFETYHOOK_ROOT / 'include' / 'safetyhook',
]
INTERNAL_INCLUDE_PATHS = [SAFETYHOOK_ROOT / 'src']
INCLUDE_REGEXP = re.compile(r'^#\s*include\s*"((?:safety).*)"\s*$')
OUTPUT_DIR = SAFETYHOOK_ROOT / 'amalgamated-dist'
FILE_HEADER = ['// DO NOT EDIT. This file is auto-generated by `amalgamate.py`.', '']

parser = argparse.ArgumentParser(description='bundles cpp and hpp files together')
parser.add_argument('--polyfill', action='store_true',
                    help='Replace std::expected with a polyfill so it can be compiled on C++20 or older. https://raw.githubusercontent.com/TartanLlama/expected/master/include/tl/expected.hpp')


# Python versions before 3.10 don't have the root_dir argument for glob, so we 
# crudely emulate it here.
def glob_in_dir(
        pattern: str,
        root_dir: Path,
):
    cwd = os.getcwd()
    root_dir = root_dir.resolve()
    os.chdir(root_dir)
    try:
        for path in glob(pattern, recursive=True):
            yield Path(root_dir) / path
    finally:
        os.chdir(cwd)


def find_include_path(
        include: str,
        search_paths: List[Path],
) -> Path:
    for search_path in search_paths:
        path = search_path / include
        if path.exists():
            return path.absolute()
    else:
        raise FileNotFoundError(f'can\'t find header: {include}')


def merge_headers(
        *,
        header: str,
        search_paths: List[Path],
        covered_headers: Set[Path],
        stack: List[str],
) -> List[str]:
    # Locate and load header contents.
    path = find_include_path(header, search_paths)
    with path.open() as f:
        lines = [x.rstrip() for x in f]

    if path in covered_headers:
        return []

    print(f'Processing header "{header}"')
    covered_headers.add(path)

    # Print the header we emit next & the include stack (if non-root).
    include_stack = []
    if stack:
        include_stack = [
            '//',
            '// Include stack:',
            *(f'//   - {x}' for x in stack)
        ]

    filtered = [
        f'',
        f'//',
        f'// Header: {header}',
        *include_stack,
        f'//',
        f'',
    ]

    # Copy over lines and recursively inline all headers.
    for line in lines:
        match = INCLUDE_REGEXP.match(line)
        if not match:
            filtered.append(line)
            continue

        # Recurse into includes.
        filtered += merge_headers(
            header=match.group(1),
            search_paths=search_paths,
            covered_headers=covered_headers,
            stack=stack + [header],
        )

    return filtered


def merge_sources(*, source_dir: Path, covered_headers: Set[Path]):
    output = [
        '#define NOMINMAX',
        '',
        '#include "safetyhook.hpp"',
        '',
    ]

    for source_file in glob_in_dir('**/*.cpp', source_dir):
        print(f'Processing source file "{source_file}"')

        # Print some comments to show where the code is from.
        output += [
            f'',
            f'//',
            f'// Source file: {source_file.relative_to(source_dir)}',
            f'//',
            f'',
        ]

        # Read source file.
        with (source_dir / source_file).open() as f:
            lines = [x.rstrip() for x in f]

        # Walk source file's lines.
        for line in lines:
            # Emit non-includes as-is.
            match = INCLUDE_REGEXP.match(line)
            if not match:
                output.append(line)
                continue
            path = match.group(1)

            if path in covered_headers:
                continue

            if 'Internal' not in path and 'Generated' not in path:
                print(
                    f'WARN: Including header that looks like it is public '
                    f'and should thus already be covered by `safetyhook.h` '
                    f'during processing of source files: {path}'
                )

            print(f'Processing internal header "{path}"')
            output += merge_headers(
                header=path,
                search_paths=PUBLIC_INCLUDE_PATHS + INTERNAL_INCLUDE_PATHS,
                covered_headers=covered_headers,
                stack=[],
            )

    return output


def do_polyfill(content):
    return content.replace('#include <expected>',
                           dedent('''
                           #if __has_include("tl/expected.hpp")
                               #include "tl/expected.hpp"
                           #elif __has_include("expected.hpp")
                               #include "expected.hpp"
                           #else
                               #error "No <expected> polyfill found"
                           #endif
                           ''')) \
        .replace('std::expected', 'tl::expected') \
        .replace('std::unexpected', 'tl::unexpected')


def main():
    args = parser.parse_args()
    polyfill = args.polyfill is True

    if OUTPUT_DIR.exists():
        print('Output directory exists. Deleting.')
        rmtree(OUTPUT_DIR)

    OUTPUT_DIR.mkdir()

    covered_headers = set()
    with open(OUTPUT_DIR / 'safetyhook.hpp', 'w') as f:
        content = '\n'.join(FILE_HEADER + merge_headers(
            header='safetyhook.hpp',
            search_paths=PUBLIC_INCLUDE_PATHS,
            covered_headers=covered_headers,
            stack=[],
        ))
        if polyfill:
            content = do_polyfill(content)
        f.write(content)

    print(covered_headers)

    with open(OUTPUT_DIR / 'safetyhook.cpp', 'w') as f:
        content = '\n'.join(FILE_HEADER + merge_sources(
            source_dir=SAFETYHOOK_ROOT / 'src',
            covered_headers=covered_headers,
        ))
        if polyfill:
            content = do_polyfill(content)
        f.write(content)


if __name__ == '__main__':
    main()