import ast
from pathlib import Path
from typing import Iterable, Set, Dict, Tuple
import re
from collections import defaultdict
import pandas as pd


def iter_python_scripts(root: str) -> Iterable[Path]:
    for p in Path(root).rglob("*"):
        if p.name.endswith(".py"):
            yield p


def read_file(path: Path) -> str:
    return path.read_text()


class Visitor(ast.NodeVisitor):
    def __init__(self) -> None:
        super().__init__()
        self.nodes: Set[str] = set()

    def visit_Assign(self, node: ast.Assign):
        if isinstance(node.value, ast.Str) and re.match(r"^MLFLOW_[A-Z0-9_]+$", node.value.s):
            self.nodes.add((node.value.s, node.lineno))
        self.generic_visit(node)


def main() -> None:
    envs: Dict[str, Set(Tuple(str, int))] = dict()
    for d in ["mlflow"]:
        for path in iter_python_scripts(d):
            if str(path) == "mlflow/environment_variables.py":
                continue
            visitor = Visitor()
            src = read_file(path)
            root = ast.parse(src)
            visitor.visit(root)
            if visitor.nodes:
                envs[str(path)] = visitor.nodes

    data = []
    for path, vals in envs.items():
        for (name, lineno) in vals:
            data.append(
                (
                    f"`{name}`",
                    "https://github.com/mlflow/mlflow/blob/{}/{}#L{}".format(
                        "1b89add428da0a4453e7523a20d13f04f6291f37", path, lineno
                    ),
                    "",
                    "",
                )
            )

    print(
        pd.DataFrame(data, columns=["Name", "Location", "Assignee", "PR"]).to_markdown(index=False)
    )


if __name__ == "__main__":
    main()