#!/usr/bin/env python3

import ast
import subprocess
import sys
from collections import Counter
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from itertools import chain
from pathlib import Path


@dataclass
class NestingCounter:
    nesting: int = 0

    @contextmanager
    def enabled(self) -> Iterator[None]:
        self.nesting += 1
        try:
            yield
        finally:
            self.nesting -= 1

    @property
    def active(self) -> bool:
        """Determine whether the level of nesting is currently non-zero"""
        return bool(self.nesting)


def fix_bad_syntax(path: Path) -> None:
    with open(path) as f:
        stub = f.read()

    lines = dict(enumerate(stub.splitlines()))
    tree = ast.parse(stub)
    lines_to_delete: list[int] = []

    class OldSyntaxReplacer(ast.NodeVisitor):
        def __init__(self) -> None:
            self.visiting_orelse = NestingCounter()
            self.visiting_class = NestingCounter()
            self.single_statement_class = False
            self.class_stmt_has_orelse = False

        @staticmethod
        def get_linenos(node: ast.AST) -> tuple[int, int]:
            lineno, end_lineno = node.lineno, node.end_lineno
            assert isinstance(lineno, int)
            assert isinstance(end_lineno, int)
            return lineno, end_lineno

        def delete_node(self, node: ast.AST, *, first_line_only: bool = False) -> None:
            lineno, end_lineno = self.get_linenos(node)
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                lineno -= len(node.decorator_list)
            if self.visiting_class.active and self.single_statement_class and not self.class_stmt_has_orelse:
                lines[lineno - 2] += ' ...'
            if first_line_only:
                lines_to_delete.append(lineno - 1)
            else:
                lines_to_delete.extend(range(lineno - 1, end_lineno))

        def visit_ClassDef(self, node: ast.ClassDef) -> None:
            old_val = self.single_statement_class
            if len(node.body) == 1:
                self.single_statement_class = True
                class_stmt_node = node.body[0]
                if isinstance(class_stmt_node, ast.If):
                    self.class_stmt_has_orelse = bool(class_stmt_node.orelse)
            with self.visiting_class.enabled():
                self.generic_visit(node)
            self.single_statement_class = old_val

        def visit_If(self, node: ast.If) -> None:
            test, body, orelse = node.test, node.body, node.orelse
            if isinstance(test, ast.Compare):
                if not ast.unparse(test).startswith("sys.version_info "):
                    self.generic_visit(node)
                elif ast.unparse(test) == "sys.version_info < (3, 7)":
                    return self.delete_node(node)
                elif ast.unparse(test) == "sys.version_info >= (3, 7)":
                    if self.visiting_orelse.active:
                        lines[node.lineno] = lines[node.lineno].replace("elif sys.version_info >= (3, 7):", "else:")
                    else:
                        self.delete_node(node, first_line_only=True)
                        if body:
                            start, _ = self.get_linenos(test)
                            _, end = self.get_linenos(body[-1])
                            for lineno in range(start, end):
                                lines[lineno] = lines[lineno][4:]
                    if orelse:
                        _, lineno_to_delete = self.get_linenos(body[-1])
                        line = lines[lineno_to_delete]
                        while line.strip() != "else:":
                            lineno_to_delete += 1
                            line = lines[lineno_to_delete]
                        lines_to_delete.append(lineno_to_delete)
                        for child in orelse:
                            self.delete_node(child)
                    for child in body:
                        self.generic_visit(child)
                else:
                    for child in body:
                        self.generic_visit(child)
                    with self.visiting_orelse.enabled():
                        for child in orelse:
                            self.generic_visit(child)
            else:
                self.generic_visit(node)

    OldSyntaxReplacer().visit(tree)

    if not lines_to_delete:
        return

    for lineno in lines_to_delete:
        try:
            del lines[lineno]
        except KeyError:
            continue

    new_stub = '\n'.join(lines.values())
    try:
        new_tree = ast.parse(new_stub)
    except:
        print('OLD\n\n')
        print('\n'.join(stub.splitlines()[140:150]))
        print('\n\nNEW\n\n')
        print('\n'.join(new_stub.splitlines()[140:150]))
        raise
        
    names_count: Counter[str] = Counter()

    class NameFinder(ast.NodeVisitor):
        def visit_Name(self, node: ast.Name) -> None:
            names_count[node.id] += 1

    NameFinder().visit(new_tree)

    if names_count["sys"] == 1:
        for node in new_tree.body:
            if isinstance(node, ast.Import) and node.names[0].name == "sys":
                del lines[node.lineno - 1]
                new_stub = '\n'.join(lines.values())
                break

    with open(path, "w") as f:
        f.write(new_stub + "\n")


def main() -> None:
    print("STARTING RUN: Will attempt to fix new syntax in the stubs directory...\n\n")
    for path in chain(Path("stubs").rglob("*.pyi"), Path("stdlib").rglob("*.pyi")):
        if "@python2" not in path.parts and ("protobuf" not in path.parts or "_pb2" not in str(path)):
            print(f"Attempting to convert {path} to new syntax.")
            fix_bad_syntax(path)

    print("\n\nSTARTING ISORT...\n\n")
    subprocess.run([sys.executable, "-m", "isort", "."])

    print("\n\nSTARTING BLACK...\n\n")
    subprocess.run([sys.executable, "-m", "black", "."])

    print('\n\nRunning "check_new_syntax.py"...\n\n')
    subprocess.run([sys.executable, "tests/check_new_syntax.py"])

    print("\n\nRunning flake8...\n\n")
    subprocess.run([sys.executable, "-m", "flake8", "stdlib stubs"])


if __name__ == "__main__":
    main()