#!/usr/bin/env python3

import re
from datetime import datetime

import cython

GEN_START = "##### START: GENERATED LIST OF GENERATED TYPES #####\n"
GEN_END = "##### END: GENERATED LIST OF GENERATED TYPES #####\n"

map_py_type_to_name = {
    bool: 'bint',
    float: 'py_float',
    int: 'py_int',
    complex: 'py_complex',
    None: 'Any',
}.get


def non_generated_lines(lines, start=GEN_START, end=GEN_END):
    lines = iter(lines)
    for line in lines:
        if line == start:
            for line in lines:
                if line == end:
                    break
            continue
        yield line


def find_known_names(file_path):
    match_name = re.compile(r"(?:(?:class|type)\s+)?(\w+)\s*[=:\[](?:\w|\s)").match

    with open(file_path) as f:
        return {
            match.group(1)
            for match in map(match_name, non_generated_lines(f))
            if match is not None
        }


def replace_type_list(file_path, type_lines):
    with open(file_path) as f:
        lines = f.readlines()

    try:
        start_index = lines.index(GEN_START)
        end_index = lines.index(GEN_END, start_index)
    except ValueError:
        raise RuntimeError(f"Failed to find generated section in {file_path}")

    new_lines = lines[:start_index+1]
    new_lines.append(f'# Generated by "Tools/cython-generate-shadow-pyi.py" on {datetime.now()}\n')
    new_lines.append("\n")

    new_lines.extend(type_lines)

    new_lines.append("\n")
    new_lines.extend(lines[end_index:])

    if lines[start_index+2:] == new_lines[start_index+2:]:
        # No changes except for the generation date => don't change file.
        return

    with open(file_path, 'w') as f:
        f.writelines(new_lines)


def map_type(pytype):
    try:
        py_type_name = map_py_type_to_name(pytype)
        if py_type_name is not None:
            return py_type_name
    except TypeError:
        # Type is not hashable.
        pass

    if isinstance(pytype, (cython.const, cython.volatile)):
        return f"{type(pytype).__name__}[{map_type(pytype._basetype)}]"

    if isinstance(pytype, cython.typedef):
        return map_type(pytype._basetype)

    if isinstance(pytype, type) and issubclass(pytype, cython.PointerType):
        base_type = map_type(pytype._basetype)
        if issubclass(pytype, cython.ArrayType):
            return f"array[{base_type}, {pytype._n}]"
        else:
            return f"pointer[{base_type}]"

    raise ValueError(f"Unmappable type '{pytype}({type(pytype).__mro__})")


def map_types(namespace, ignore=()):
    for type_name, pytype in namespace.items():
        if type_name in ignore:
            continue
        if type_name.startswith('_'):
            continue
        is_type_alias = isinstance(pytype, cython.typedef)
        if not (is_type_alias or isinstance(pytype, type)):
            continue

        try:
            py_type_name = map_type(pytype)
        except (ValueError, AttributeError, TypeError):
            print(f"Not mapping type '{type_name}' from {pytype}")
            continue

        if type_name == py_type_name:
            continue

        yield type_name, py_type_name, is_type_alias


def generate_type_lines(types):
    # Sort 'pointers_const_type_name' by (type name, const, pointers)
    def sort_key(item):
        return '_'.join(reversed(item[0].lower().split('_')))

    for type_name, py_type_code, is_type_alias in sorted(types, key=sort_key):
        yield f"{type_name}{' : TypeAlias' if is_type_alias else ''} = {py_type_code}\n"


def main(file_path):
    namespace = vars(cython)
    declared_names = find_known_names(file_path)

    types = map_types(namespace, ignore=declared_names)

    type_lines = generate_type_lines(types)
    replace_type_list(file_path, type_lines)


if __name__ == "__main__":
    import sys
    shadow_py = sys.argv[1] if len(sys.argv) > 1 else "Cython/Shadow.pyi"
    main(shadow_py)
