#!/usr/bin/env python3
# Copyright 2025 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

"""MLIR passes markdown document postprocessor for website generation.

The markdown files generated by `iree-tblgen` and `mlir-tblgen` are generic,
but we want to process them through the same static site generator (mkdocs)
as the rest of our website. This script allows us to make certain style and
structure changes to each file in a directory.

Usage (typically invoked by generate_extra_files.sh, edits files in-place):
  postprocess_passes_docs.py ${WEBSITE_DOCS_BUILD_DIRECTORY}
"""

import argparse
import fileinput
import os
import pathlib
import re


def main(args):
    directory = args.directory
    files = list(pathlib.Path(directory).iterdir())

    # Add explicit custom_edit_url links as these markdown files were generated
    # from other source files.
    passes_sources_map = {
        # Dialects:
        "Check.md": "compiler/src/iree/compiler/Modules/Check/Transforms",
        "Flow.md": "compiler/src/iree/compiler/Dialect/Flow/Transforms",
        "Encoding.md": "compiler/src/iree/compiler/Dialect/Encoding/Transforms",
        "HAL.md": "compiler/src/iree/compiler/Dialect/HAL/Transforms",
        "HALInline.md": "compiler/src/iree/compiler/Modules/HAL/Inline/Transforms",
        "HALLoader.md": "compiler/src/iree/compiler/Modules/HAL/Loader/Transforms",
        "IOParameters.md": "compiler/src/iree/compiler/Modules/IO/Parameters/Transforms",
        "IREECodegen.md": "compiler/src/iree/compiler/Codegen/Dialect/Codegen/Transforms",
        "IREEGPU.md": "compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms",
        "IREEVectorExt.md": "compiler/src/iree/compiler/Codegen/Dialect/VectorExt/Transforms",
        "LinalgExt.md": "compiler/src/iree/compiler/Dialect/LinalgExt/Transforms",
        "Stream.md": "compiler/src/iree/compiler/Dialect/Stream/Transforms",
        "TensorExt.md": "compiler/src/iree/compiler/Dialect/TensorExt/Transforms",
        "Util.md": "compiler/src/iree/compiler/Dialect/Util/Transforms",
        "VM.md": "compiler/src/iree/compiler/Dialect/VM/Transforms",
        "VMVX.md": "compiler/src/iree/compiler/Dialect/VMVX/Transforms",
        # Pipelines:
        "ConstEval.md": "compiler/src/iree/compiler/ConstEval",
        "DispatchCreation.md": "compiler/src/iree/compiler/DispatchCreation",
        "GlobalOptimization.md": "compiler/src/iree/compiler/GlobalOptimization",
        "InputConversion.md": "compiler/src/iree/compiler/InputConversion/Common",
        "Preprocessing.md": "compiler/src/iree/compiler/Preprocessing/Common",
        # Codegen:
        "CodegenCommonCPU.md": "Codegen/Common/CPU",
        "CodegenCommonGPU.md": "Codegen/Common/GPU",
        "CodegenCommon.md": "Codegen/Common",
        "CodegenDialectGPU.md": "Codegen/Dialect/GPU/Transforms",
        "CodegenDialectVectorExt.md": "Codegen/Dialect/VectorExt/Transforms",
        "CodegenLLVMCPU.md": "Codegen/LLVMCPU",
        "CodegenLLVMGPU.md": "Codegen/LLVMGPU",
        "CodegenSPIRV.md": "Codegen/SPIRV",
        "CodegenVMVX.md": "Codegen/VMVX",
    }
    base_url = "https://github.com/iree-org/iree/tree/main/"
    for file in files:
        filename = pathlib.Path(file).name
        relative_path = passes_sources_map.get(filename, None)
        if not relative_path:
            continue

        full_url = base_url + relative_path
        with open(file, "r+", encoding="utf8") as f:
            content = f.read()
            f.seek(0, 0)
            frontmatter = f"""---
custom_edit_url: {full_url}
---
"""
            f.write(frontmatter + os.linesep + content)


def parse_arguments():
    parser = argparse.ArgumentParser(description="Passes doc postprocessor.")

    parser.add_argument(
        "directory",
        help="Passes docs directory to edit in-place.",
    )
    return parser.parse_args()


if __name__ == "__main__":
    main(parse_arguments())