#!/usr/bin/env python

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Script to convert type hints to follow PEP-0585

For more information, see https://peps.python.org/pep-0585

To run from the repository's root directory:
    python convert-types.py
"""

from __future__ import annotations

from collections.abc import Callable, Iterator
import difflib
from glob import glob
import logging
import re
import sys
from typing import NamedTuple, TypeVar

# TODO:
# - False positives with lambdas and dict comprehensions due to the `:` misinterpreted as a type hint
# - Sort imports case insensitive (e.g. PIL, Flask)
# - Type hint arguments can be lists like `Callable[[a, b], c]`

# Cases not covered:
# - Multi-line imports like `from M import (\nA,\nB,\n)`
# - Importing `typing` directly like `import typing` and `x: typing.Any`
# - Parsing types with `|` syntax like Union or Optional
# - typing.re.Match --> re.Match
# - typing.re.Pattern --> re.Pattern


BUILTIN_TYPES = {"Tuple", "List", "Dict", "Set", "FrozenSet", "Type"}
COLLECTIONS_TYPES = {"Deque", "DefaultDict", "OrderedDict", "Counter", "ChainMap"}
COLLECTIONS_ABC_TYPES = {
    "Awaitable",
    "Coroutine",
    "AsyncIterable",
    "AsyncIterator",
    "AsyncGenerator",
    "Iterable",
    "Iterator",
    "Generator",
    "Reversible",
    "Container",
    "Collection",
    "Callable",
    "AbstractSet",
    "MutableSet",
    "Mapping",
    "MutableMapping",
    "Sequence",
    "MutableSequence",
    "ByteString",
    "MappingView",
    "KeysView",
    "ItemsView",
    "ValuesView",
}
CONTEXTLIB_TYPES = {"ContextManager", "AsyncContextManager"}
RE_TYPES = {"Match", "Pattern"}

RENAME_TYPES = {
    "Tuple": "tuple",  # builtin
    "List": "list",  # builtin
    "Dict": "dict",  # builtin
    "Set": "set",  # builtin
    "FrozenSet": "frozenset",  # builtin
    "Type": "type",  # builtin
    "Deque": "deque",  # collections
    "DefaultDict": "defaultdict",  # collections
    "AbstractSet": "Set",  # collections.abc
    "ContextManager": "AbstractContextManager",  # contextlib
    "AsyncContextManager": "AbstractAsyncContextManager",  # contextlib
}

# Parser a = String -> (a, String)
a = TypeVar("a")
Parser = Callable[[str], tuple[a, str]]


class TypeHint(NamedTuple):
    name: str
    args: list[TypeHint]

    def __repr__(self) -> str:
        match (self.name, self.args):
            case ("Optional", [x]):
                return f"{x} | None"
            case ("Union", args):
                return " | ".join(map(str, args))
            case (name, []):
                return name
            case (name, args):
                return f"{name}[{', '.join(map(str, args))}]"

    def patch(self, types: set[str]) -> TypeHint:
        if self.name in types:
            name = RENAME_TYPES.get(self.name, self.name)
        else:
            name = self.name
        return TypeHint(name, [arg.patch(types) for arg in self.args])


def patch_file(file_path: str, dry_run: bool = False, quiet: bool = False) -> None:
    with open(file_path) as f:
        before = f.read()
    try:
        lines = [line.rstrip() for line in before.splitlines()]
        if types := find_typing_imports(lines):
            lines = insert_import_annotations(lines)
            lines = [patched for line in lines for patched in patch_imports(line)]
            lines = sort_imports(lines)
            after = patch_type_hints("\n".join(lines), types) + "\n"
            if before == after:
                return

            if not dry_run:
                with open(file_path, "w") as f:
                    f.write(after)
                print(file_path)
            elif not quiet:
                print(f"| {file_path}")
                print(f"+--{'-' * len(file_path)}")
                diffs = difflib.context_diff(
                    before.splitlines(keepends=True),
                    after.splitlines(keepends=True),
                    fromfile="Before changes",
                    tofile="After changes",
                    n=1,
                )
                sys.stdout.writelines(diffs)
                print(f"+{'=' * 100}")
                print("| Press [ENTER] to continue to the next file")
                input()
    except Exception:
        logging.exception(f"Could not process file: {file_path}")


def insert_import_annotations(lines: list[str]) -> list[str]:
    new_import = "from __future__ import annotations"
    if new_import in lines:
        return lines

    match find_import(lines):
        case None:
            return lines
        case i:
            if lines[i].startswith("from __future__ import "):
                return lines[:i] + [new_import] + lines[i:]
            return lines[:i] + [new_import, ""] + lines[i:]


def find_typing_imports(lines: list[str]) -> set[str]:
    return {
        name.strip()
        for line in lines
        if line.startswith("from typing import ")
        for name in line.split("import")[1].split(",")
    }


def find_import(lines: list[str]) -> int | None:
    for i, line in enumerate(lines):
        if line.startswith(("import ", "from ")):
            return i
    return None


def get_imports_group(lines: list[str]) -> tuple[list[str], list[str]]:
    for i, line in enumerate(lines):
        if not line.strip() or line.startswith("#"):
            return (lines[:i], lines[i:])
    return ([], lines)


def import_name(line: str) -> str:
    match line.split():
        case ["import", name, *_]:
            return name
        case ["from", name, "import", *_]:
            return name
    raise ValueError(f"not an import: {line}")


def sort_imports(lines: list[str]) -> list[str]:
    match find_import(lines):
        case None:
            return lines
        case i:
            (imports, left) = get_imports_group(lines[i:])
            if imports:
                return lines[:i] + sorted(imports, key=import_name) + sort_imports(left)
            return left


def patch_imports(line: str) -> Iterator[str]:
    if not line.startswith("from typing import "):
        yield line
        return

    types = find_typing_imports([line])
    collections_types = types.intersection(COLLECTIONS_TYPES)
    collections_abc_types = types.intersection(COLLECTIONS_ABC_TYPES)
    contextlib_types = types.intersection(CONTEXTLIB_TYPES)
    re_types = types.intersection(RE_TYPES)
    typing_types = (
        types
        - BUILTIN_TYPES
        - COLLECTIONS_TYPES
        - COLLECTIONS_ABC_TYPES
        - CONTEXTLIB_TYPES
        - RE_TYPES
        - {"Optional", "Union"}
    )

    rename = lambda name: RENAME_TYPES.get(name, name)
    if collections_types:
        names = sorted(map(rename, collections_types))
        yield f"from collections import {', '.join(names)}"
    if collections_abc_types:
        names = sorted(map(rename, collections_abc_types))
        yield f"from collections.abc import {', '.join(names)}"
    if contextlib_types:
        names = sorted(map(rename, contextlib_types))
        yield f"from contextlib import {', '.join(names)}"
    if re_types:
        names = sorted(map(rename, re_types))
        yield f"from re import {', '.join(names)}"
    if typing_types:
        names = sorted(map(rename, typing_types))
        yield f"from typing import {', '.join(names)}"


def patch_type_hints(txt: str, types: set[str]) -> str:
    if m := re.search(rf"(?:->|:) *(\w+)", txt):
        (typ, left) = parse_type_hint(txt[m.start(1) :])
        return f"{txt[:m.start(1)]}{typ.patch(types)}{patch_type_hints(left, types)}"
    return txt


# Parser combinators
def parse_text(src: str, txt: str) -> tuple[str, str]:
    if src.startswith(txt):
        return (src[: len(txt)], src[len(txt) :])
    raise SyntaxError("text")


def parse_identifier(src: str) -> tuple[str, str]:
    if m := re.search(r"[\w\._]+", src):
        return (m.group(), src[m.end() :])
    raise SyntaxError("name")


def parse_zero_or_more(src: str, parser: Parser[a]) -> tuple[list[a], str]:
    try:
        (x, src) = parser(src)
        (xs, src) = parse_zero_or_more(src, parser)
        return ([x] + xs, src)
    except SyntaxError:
        return ([], src)


def parse_comma_separated(src: str, parser: Parser[a]) -> tuple[list[a], str]:
    def parse_next(src: str) -> tuple[a, str]:
        (_, src) = parse_text(src, ",")
        (_, src) = parse_zero_or_more(src, lambda src: parse_text(src, " "))
        return parser(src)

    try:
        (x, src) = parser(src)
        (xs, src) = parse_zero_or_more(src, parse_next)
        return ([x] + xs, src)
    except SyntaxError:
        return ([], src)


def parse_type_hint(src: str) -> tuple[TypeHint, str]:
    (name, src) = parse_identifier(src)
    try:
        (_, src) = parse_text(src, "[")
        (args, src) = parse_comma_separated(src, parse_type_hint)
        (_, src) = parse_text(src, "]")
        return (TypeHint(name, args), src)
    except SyntaxError:
        return (TypeHint(name, []), src)


def run(patterns: list[str], dry_run: bool = False, quiet: bool = False):
    for pattern in patterns:
        for filename in glob(pattern, recursive=True):
            patch_file(filename, dry_run, quiet)


if __name__ == "__main__":
    import argparse

    assert sys.version_info.major == 3, "Requires Python 3"
    assert sys.version_info.minor >= 10, "Requires Python >= 3.10 for pattern matching"

    parser = argparse.ArgumentParser()
    parser.add_argument("patterns", nargs="*", default=["**/*.py"])
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--quiet", action="store_true")
    args = parser.parse_args()

    run(**args.__dict__)