#!/usr/bin/env python3

# Copyright Advanced Micro Devices, Inc., or its affiliates.
# SPDX-License-Identifier: MIT

"""Load kernel graph from the meta data of an assembly file and render it.

Usage:

  kgraph.py path/to/assembly.s

this creates a PDF file at path/to/assembly.pdf.  To automatically
open it (via xdg-open), pass '-x'.

"""

import argparse
import re
import subprocess
from pathlib import Path

import yaml
from dot_diff import diff_dots


def extract_normalised_asm(path: Path, keep_trailing_comments: bool = False):
    """Extract normalized ASM from assembly file."""
    source = path.read_text()

    meta = False

    asm = []
    for line in source.splitlines():
        if keep_trailing_comments:
            line = line.strip()
            if line.startswith("//"):
                continue
        else:
            double_slash = line.find("//")
            if double_slash != -1:
                line = line[:double_slash]
            line = line.rstrip()
        if "kernel_graph_dot" in line:
            continue
        if ".amdgpu_metadata" in line:
            meta = True
        if line and not meta:
            asm.append(line)
        if ".end_amdgpu_metadata" in line:
            meta = False

    return "\n".join(asm)


def extract_asm_dot(path: Path):
    """Extract .kernel_graph meta data from assembly file."""
    source = path.read_text()
    beginMatch = re.search(r"^---$", source, re.MULTILINE)
    endMatch = re.search(r"^\.\.\.$", source, re.MULTILINE)

    assert beginMatch is not None
    assert endMatch is not None

    beginPos = beginMatch.span()[1]
    endPos = endMatch.span()[0]
    assert beginPos < endPos

    meta = yaml.load(source[beginPos:endPos], Loader=yaml.CSafeLoader)
    kernel = meta["amdhsa.kernels"][0]

    def try_get(key):
        if key in kernel:
            return kernel[key]
        return None

    return try_get(".name"), try_get(".kernel_graph_dot"), try_get(".kernel_graph")


def extract_log_dots(path: Path):
    source = path.read_text()
    source_clean = ""
    prev_end = 0
    dots = []
    for graph in re.finditer("digraph {", source):
        this_graph = source[graph.start() :].partition("\n\n")[0]
        before_graph = source[prev_end : graph.start()]
        source_clean += before_graph
        prev_end = graph.start() + len(this_graph)

        dots.append(this_graph)

    source_clean += source[prev_end:]

    return dots, source_clean


def write_dot(dot: str, fname: Path):
    out_fname = fname.with_suffix(".dot")
    out_fname.write_text(dot)
    print(f"Wrote {out_fname}")
    return out_fname


def render_dot(fname: Path, dot: str):
    """Render graph."""
    with fname.open("w") as out:
        subprocess.run(["dot", "-Tpdf"], input=dot.encode(), stdout=out)
    print(f"Wrote {fname}")


def open_dot(fname):
    subprocess.call(["xdg-open", str(fname)])


def open_code(fname):
    subprocess.call(["code", str(fname)])


def process_dot(
    dot: str,
    out_path: Path,
    code_open: bool,
    dot_only: bool,
    xdg_open: bool,
):
    dotfile = write_dot(dot, out_path)
    if code_open:
        open_code(dotfile)
    if not dot_only:
        rendered_fname = out_path.with_suffix(".pdf")
        render_dot(rendered_fname, dot)
        if xdg_open:
            open_dot(rendered_fname)


def process_serialized_graph(graph: dict, out_path: Path):
    with out_path.with_suffix(".yaml").open("w") as f:
        yaml.dump(graph, f, default_flow_style=None, Dumper=yaml.CSafeDumper)
        print(f"Wrote {f.name}")


def process_normalized_asm(path: Path, keep_trailing_comments: bool = True):
    asm = Path(path.stem + "_normalized.s")
    asm.write_text(extract_normalised_asm(path, keep_trailing_comments))
    print(f"Wrote {asm}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Extract and render kernel graph.")
    parser.add_argument("fname", type=str, help="Assembly or log file name")
    parser.add_argument("-d", "--dot-only", default=False, action="store_true")
    parser.add_argument(
        "-x",
        "--xdg-open",
        default=False,
        action="store_true",
        help="Open with xdg-open after generating",
    )
    parser.add_argument(
        "-c",
        "--code-open",
        default=False,
        action="store_true",
        help="Open .dot with VSCode after generating",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        help="Output file base name",
    )
    parser.add_argument(
        "-k",
        "--keep-trailing",
        default=False,
        action="store_true",
        help="Keep trailing comments in normalized assembly file.",
    )
    parser.add_argument(
        "--omit_diff",
        default=False,
        action="store_true",
        help="Omit dot diff coloring",
    )

    args = parser.parse_args()
    path = Path(args.fname)

    foutput = path
    if args.output is not None:
        foutput = Path(args.output)

    if path.suffix == ".s":
        _, dot, graph = extract_asm_dot(path)
        process_dot(
            dot,
            foutput,
            args.code_open,
            args.dot_only,
            args.xdg_open,
        )
        if graph is not None and not args.dot_only:
            process_serialized_graph(graph, foutput)

        process_normalized_asm(path, args.keep_trailing)

    elif path.suffix == ".log":
        dots, source_clean = extract_log_dots(path)

        foutput_clean = Path(str(foutput) + "_clean.log")
        foutput_clean.write_text(source_clean)

        if not args.omit_diff:
            dots = diff_dots(dots)
        for i, dot in enumerate(dots):
            serial_str = f"_{i:04d}"
            foutput_serial = Path(str(foutput) + serial_str)
            process_dot(
                dot,
                foutput_serial,
                args.code_open,
                args.dot_only,
                args.xdg_open,
            )
    else:
        print("Unknown file extension")