import argparse
import pathlib
import io
import sys
from typing import List


def get_bazel_files(root_dir: pathlib.Path) -> List[pathlib.Path]:
    """Finds files matching "BUILD.bazel" or "*.bzl" under the given directory.

    Args:
        root_dir: The root directory to search in.

    Returns:
        A list of pathlib.Path objects representing the matching files.
    """
    return [
        path
        for path in root_dir.rglob("*")
        if path.match("BUILD.bazel") or path.match("*.bzl")
    ]

def print_file_info(
        file_path: pathlib.Path,
        root_dir: pathlib.Path,
        writer: io.TextIOWrapper,
):
    """Prints the relative path, line count, and prefixed content of a file.

    Args:
        file_path: The path to the file.
        root_dir: The root directory to calculate the relative path from.
        writer: Where to write the output
    """
    relative_path = file_path.relative_to(root_dir)
    with open(file_path, "r", encoding="utf-8") as file:
        content = file.readlines()
        line_count = len(content)
        writer.write("\n")
        writer.write("--- /dev/null\n")
        writer.write(f"+++ {relative_path}\n")
        writer.write(f"@@ -0,0 +1,{line_count} @@\n")
        for line in content:
            writer.write(f"+{line.rstrip()}\n")

def print_all(root_dir: pathlib.Path, writer: io.TextIOWrapper):
    """Iterate through all BUILD.bazel and *.bzl files and write them to writer

    Args:
        root_dir: The root directory to calculate the relative path from.
        writer: Where to write the output
    """
    for file_path in get_bazel_files(root_dir=root_dir):
        print_file_info(file_path=file_path, root_dir=root_dir, writer=writer)

def main() -> None:
    """Finds all BUILD.bazel files and *.bzl files and generates a diff"""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--root-dir",
        type=pathlib.Path,
        required=True
    )
    parser.add_argument(
        "-o",
        dest="output",
        type=pathlib.Path,
    )
    args = parser.parse_args()
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        with open(args.output, mode="w", encoding="utf-8") as o:
            print_all(root_dir=args.root_dir, writer=o)
    else:
        print_all(
            root_dir=args.root_dir,
            writer=io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8")
        )

if __name__ == "__main__":
    main()