import subprocess
import json
import os
import sys
import colorama
from colorama import Fore, Back, Style
from collections import defaultdict
colorama.init()


def parse_msvc_wall(output):
    i = output.find(b'c1xx.dll)=')
    if i == -1: return None
    j = output.find(b's', i)
    return float(output[i+10:j])


def msvc_get_parsing_time(*fnames, std=None):
    if len(fnames) == 1 and isinstance(fnames, list):
        fnames = fnames[0]

    cmd = ['cl', '/nologo', '/Bt', '/Zs', '/TP', '/w', '.empty.tmp']

    if std is not None:
        cmd.append(f'/std:c++{std}')

    cmd += [f'/FI{fn[1:-1]}' for fn in fnames]
    #print(' '.join(cmd))

    try:
        output = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True).stdout
    except subprocess.CalledProcessError as e:
        if b'C1083' in e.stdout:
            return None
        time = parse_msvc_wall(e.stderr)
        if time is not None:
            return time
        print(f'stderr={e.stderr}')
        print(f'stderr={e.stdout}')
        raise

    return parse_msvc_wall(output)


# https://eel.is/c++draft/headers
headers = set('''
<algorithm>
<flat_set>
<mutex>
<stdexcept>
<any>
<format>
<new>
<stdfloat>
<array>
<forward_list>
<numbers>
<stop_token>
<atomic>
<fstream>
<numeric>
<streambuf>
<barrier>
<functional>
<optional>
<string>
<bit>
<future>
<ostream>
<string_view>
<bitset>
<generator>
<print>
<strstream>
<charconv>
<initializer_list>
<queue>
<syncstream>
<chrono>
<iomanip>
<random>
<system_error>
<codecvt>
<ios>
<ranges>
<thread>
🔗
<compare>
<iosfwd>
<ratio>
<tuple>
🔗
<complex>
<iostream>
<regex>
<type_traits>
🔗
<concepts>
<istream>
<scoped_allocator>
<typeindex>
🔗
<condition_variable>
<iterator>
<semaphore>
<typeinfo>
🔗
<coroutine>
<latch>
<set>
<unordered_map>
🔗
<deque>
<limits>
<shared_mutex>
<unordered_set>
🔗
<exception>
<list>
<source_location>
<utility>
🔗
<execution>
<locale>
<span>
<valarray>
🔗
<expected>
<map>
<spanstream>
<variant>
🔗
<filesystem>
<mdspan>
<sstream>
<vector>
🔗
<flat_map>
<memory>
<stack>
<version>
🔗
<memory_resource>
<stacktrace>
<cassert>
<cfenv>
<climits>
<csetjmp>
<cstddef>
<cstdlib>
<cuchar>
🔗
<cctype>
<cfloat>
<clocale>
<csignal>
<cstdint>
<cstring>
<cwchar>
🔗
<cerrno>
<cinttypes>
<cmath>
<cstdarg>
<cstdio>
<ctime>
<cwctype>
'''.strip().splitlines()) - {'🔗'}
#headers = ['<string_view>', '<chrono>']
max_header_name_len = max(map(len, headers))


def time_to_color(t):
    if t < 0.010: return Fore.BLACK, Style.BRIGHT
    if t < 0.020: return Fore.WHITE, Style.DIM
    if t < 0.040: return '', ''
    if t < 0.060: return Fore.CYAN, Style.BRIGHT
    if t < 0.080: return Fore.WHITE, Style.BRIGHT
    if t < 0.100: return Fore.YELLOW, Style.BRIGHT
    if t < 0.150: return Fore.RED, Style.BRIGHT
    if t < 0.200: return Fore.RED, Style.DIM
    if t < 0.300: return Fore.MAGENTA, Style.BRIGHT
    return Fore.MAGENTA, Style.DIM


reset_colors = Style.RESET_ALL + Fore.RESET

stds = [None, '17', '20', 'latest']


def ttc(time):
    f, s = time_to_color(time)
    return f + s


def info(timings, expected_count=len(stds), reverse=False, delim=' ', fill=' .....'):
    a = [f'{ttc(time)}{time:>6.3f}{reset_colors}' for time in timings]
    b = [fill] * (expected_count - len(timings))
    return delim.join(b + a if reverse else a + b )


def print_slowdown(all_timings, considered_disabled=0.030, min_diff=0.005, min_diff_rel=0.01):
    s = ' | '.join(f'c++{std}' if std else 'default' for std in stds)
    print(f'{s} | slowdown | header')
    print(('-' * 6 + ':|') * (len(stds) + 1) + '-' * max_header_name_len)

    results = []
    for header, timings in all_timings.items():
        timings = timings
        while len(timings) and timings[0] < considered_disabled:
            timings = timings[1:]

        if len(timings) == 0: continue

        slowdown = timings[-1] / timings[0] - 1
        if abs(timings[-1] - timings[0]) > min_diff and abs(slowdown) > min_diff_rel:
            results.append((slowdown, header, timings))

    for slowdown, header, timings in sorted(results, reverse=True):
        print(f'{info(timings, reverse=True, delim=" |")} | {slowdown:>5.0%} | `{header}`')


def get_timings():
    with open('.empty.tmp', 'w+') as f:
        pass
    all_timings = defaultdict(list)
    print(f'Timing standard library headers:')
    for header in headers:
        #print(f'working on {header}...', end='')
        print(f'{info([])} {header}', end='')

        for std in stds:
            timings = []
            for _ in range(10):
                time = msvc_get_parsing_time(header, std=std)
                if time is None:
                    timings.append(float('nan'))
                    break
                timings.append(time)
        
            time = min(timings)
            f, s = time_to_color(time)
            self_colors = f + s
            all_timings[header].append(time)
            print(f'\r{info(all_timings[header])} {header}', end='')
        print(f'\r{info(all_timings[header])} {header}')
    return all_timings


if __name__ == "__main__":
    fn = 'bench_syshdrs.json'
    if os.path.exists(fn) and not '-r' in sys.argv:
        with open(fn) as f:
            timings = json.load(f)
    else:
        timings = get_timings()
        with open(fn, 'w+') as f:
            json.dump(timings, f)
        print('\n' * 3)
    print_slowdown(timings)