#!/usr/bin/env python3

import argparse
import pathlib
import re
import subprocess
from typing import Optional, List


def picasso(
    files: List[str], out: str, header: Optional[str] = None, no_nop: bool = False
) -> None:
    cmd = [
        "picasso",
        "--out=" + out,
        *(["--header" + header] if header else []),
        *(["--no-nop"] if no_nop else []),
        *files,
    ]
    result = subprocess.run(cmd)
    result.check_returncode()


IDENTIFIER_FILTER_REGEX = re.compile(r"[^A-Za-z0-9]+")


def filename_to_identifier(file: str) -> str:
    name = pathlib.Path(file).stem
    return IDENTIFIER_FILTER_REGEX.sub("_", name)


DVLE_EMITTING_DIRECTIVES = set(
    [
        # "uniform",
        "fvec",
        "ivec",
        "bool",
        # "const",
        "constf",
        "consti",
        "constfa",
        "in",
        "out",
        "entry",
        # "nodvle",
        "gsh",
        # "setfi",
        "setf",
        "seti",
        "setb",
    ]
)

NON_DVLE_DIRECTIVES = set(
    [
        "proc",
        "else",
        "end",
        "alias",
    ]
)

DIRECTIVE_REGEX = re.compile(r"^\s*\.(\w+)")


def is_program_emit_dvle(path: str) -> bool:
    has_dvle_directive = False

    with open(path, "r", encoding="utf8") as file:
        while line := file.readline():
            line = line.strip()

            if line.startswith(";"):
                # comment
                continue

            if match := DIRECTIVE_REGEX.match(line):
                directive = match[1]

                if directive == "nodvle":
                    return False

                if directive in DVLE_EMITTING_DIRECTIVES:
                    has_dvle_directive = True
                elif directive not in NON_DVLE_DIRECTIVES:
                    raise RuntimeError(
                        f"Unexpected directive '.{directive}' in shader file '{path}'"
                    )

    return has_dvle_directive


def main() -> None:
    parser = argparse.ArgumentParser(
        prog="picasso_wrapper",
        description="Wrapper for picasso that generates an additional header (ends in '.dvle.h') with enum constants for indexing the DVLE of each program.",
        add_help=False,
    )
    parser.add_argument("-o", "--out", required=True)
    parser.add_argument("-h", "--header")
    parser.add_argument("-n", "--no-nop", action="store_true")
    parser.add_argument("files", nargs="+")

    args = parser.parse_args()

    # if args.help:

    names = [filename_to_identifier(f).upper() for f in args.files]
    unique_names = set()
    for name, path in zip(names, args.files):
        if name in unique_names:
            raise argparse.ArgumentError(
                None, f"File '{path}' produces duplicate constant name {name}."
            )
        unique_names.add(name)

    picasso(args.files, args.out, args.header, args.no_nop)

    bin_name = filename_to_identifier(args.out).upper()

    file_content = ""
    file_content += f"/* Generated by {pathlib.Path(__file__).name} - do not edit */\n"
    file_content += "#pragma once\n\n"
    file_content += f"/** DVLE indices for the programs in '{args.out}' */\n"
    file_content += f"enum {bin_name}_DVLE_enum {{\n"

    for name, path in zip(names, args.files):
        if is_program_emit_dvle(path):
            file_content += f"\t{bin_name}_DVLE_{name}, /**< {path} */\n"
        else:
            file_content += f"\t/* '{path}' does not emit a DVLE entry */\n"

    file_content += "};\n"

    old_content = ""

    out_file = pathlib.Path(args.out).with_suffix(".dvle.h")

    try:
        with open(out_file, "r", encoding="utf8", newline="\n") as file:
            old_content = file.read()
    except:
        pass

    # Skip writing file if nothing's changed (avoid unnecessary rebuilds)
    if old_content == file_content:
        return

    with open(out_file, "w", encoding="utf8", newline="\n") as file:
        file.write(file_content)


if __name__ == "__main__":
    main()