#!/usr/bin/env python3
import os
import re
import subprocess
import sys

# This script parses Dawn CMakeLists.txt files and creates a .bzl file with the files extracted.
#
# It generates file lists dynamically, including querying custom Python generators
# with the --print-cmake-outputs flag to ensure build-time autogenerated files
# are cleanly and automatically tracked without any manual hardcoding.

class CMakeValue:
    """Represents a value bound to a CMake variable."""
    def __init__(self, value):
        self.value = value

class CMakeParser:
    def __init__(self, prefix, rel_path, dawn_root):
        self.prefix = prefix
        self.rel_path = rel_path
        self.dawn_root = dawn_root
        self.variables = {
            "DAWN_INCLUDE_DIR": [CMakeValue("include")],
            "DAWN_SRC_DIR": [CMakeValue("src")],
        }
        self.generated_files = set()
        self.generator_vars = {}
        # output_vars[prefix][group] = set of files
        self.output_vars = {}

    def resolve_var(self, var_name):
        """Resolves a CMake variable, expanding nested variable references."""
        if var_name not in self.variables:
            return [CMakeValue(f"${{{var_name}}}")]

        vals = self.variables[var_name]
        resolved = []
        for cv in vals:
            v = cv.value
            if isinstance(v, str) and v.startswith("${") and v.endswith("}"):
                inner = v[2:-1]
                if inner != var_name:
                    resolved.extend(self.resolve_var(inner))
                else:
                    resolved.append(cv)
            else:
                resolved.append(cv)
        return resolved

    def expand_string(self, s):
        def repl(match):
            var_name = match.group(1)
            resolved = self.resolve_var(var_name)
            # Use the first resolved value for simplicity in path expansion
            return str(resolved[0].value) if resolved else ""

        return re.sub(r'\${([^}]+)}', repl, s)

    def clean_path(self, path):
        """Normalises variable expressions and relative paths into clean workspace-relative paths.

        Specifically, it:
          1. Fully expands any embedded CMake variables.
          2. Strips workspace-level prefixes like '${Dawn_SOURCE_DIR}/' or '${CMake_SOURCE_DIR}/'.
          3. Prepends the parsed target directory's relative path (e.g. 'src/dawn/common/')
             only if the file is relative to the subdirectory and doesn't start with root
             directories like 'include/', 'src/', or 'generator/'.

        Examples:
            "GPUInfo.cpp" (rel_path="src/dawn/common")        -> "src/dawn/common/GPUInfo.cpp"
            "${Dawn_SOURCE_DIR}/src/dawn/native/Device.cpp"   -> "src/dawn/native/Device.cpp"
            "include/dawn/webgpu.h"                           -> "include/dawn/webgpu.h"
            "generator/dawn_json_generator.py"                -> "generator/dawn_json_generator.py"
        """
        path = self.expand_string(path).replace("\\", "/")
        # Handle generator and other paths starting with variables or specific directories
        if path.startswith("${Dawn_SOURCE_DIR}/"):
            path = path[len("${Dawn_SOURCE_DIR}/"):]
        if path.startswith("${CMake_SOURCE_DIR}/"):
            path = path[len("${CMake_SOURCE_DIR}/"):]

        # Map WebGPU generated headers to match the exact custom outputs of Bazel genrules.
        # Note: webgpu_cpp_chained_struct.h belongs to the universal/standard WebGPU namespace
        # (include/webgpu/), while other generated C++ headers are Dawn-specific extensions
        # (include/dawn/). We must route them correctly here to prevent Bazel analysis mismatches.
        if path.startswith("webgpu-headers/"):
            basename = os.path.basename(path)
            if basename == "webgpu_cpp_chained_struct.h":
                path = "include/webgpu/webgpu_cpp_chained_struct.h"
            else:
                path = "include/dawn/" + basename

        if not path.startswith("include/") and not path.startswith("src/") and not path.startswith("generator/"):
            path = os.path.join(self.rel_path, path)
        return os.path.normpath(path).replace("\\", "/")

    def _find_closing_paren(self, content, start_pos):
        """Finds the index of the matching closing parenthesis, handling nested ones."""
        count = 0
        for i in range(start_pos, len(content)):
            if content[i] == '(':
                count += 1
            elif content[i] == ')':
                count -= 1
                if count == 0:
                    return i
        return -1

    def parse(self, content):
        # Remove comments
        content = re.sub(r'#.*', '', content)

        pos = 0
        while pos < len(content):
            # Find next command
            match = re.search(r'(\w+)\s*\(', content[pos:])
            if not match:
                break

            cmd_name = match.group(1).lower()
            start_paren = pos + match.end() - 1

            end_paren = self._find_closing_paren(content, start_paren)
            if end_paren == -1:
                break

            args_str = content[start_paren+1:end_paren]
            # Parse arguments, respecting quotes
            args_raw = re.findall(r'"([^"]*)"|([^\s()]+)', args_str)
            args = [a[0] or a[1] for a in args_raw if a[0] or a[1]]

            self.handle_command(cmd_name, args)
            pos = end_paren + 1

    def _run_generator(self, script, extra_params):
        """Runs the generator Python script with --print-cmake-outputs to get generated files."""
        expanded_params = [
            self.expand_string(p).replace("${Dawn_SOURCE_DIR}/", "")
            for p in extra_params
        ]

        cmd = [
            sys.executable,
            os.path.abspath(os.path.join(self.dawn_root, script)),
            "--template-dir", "generator/templates",
            "--root-dir", ".",
            "--output-dir", ".",  # Required but ignored by --print-cmake-outputs
        ] + expanded_params + ["--print-cmake-outputs"]

        try:
            res = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True,
                cwd=self.dawn_root
            )
            # Split outputs on any combination of whitespace or semicolons
            outputs = re.split(r'[\s;]+', res.stdout.strip())
            return [os.path.normpath(o) for o in outputs if o]
        except subprocess.CalledProcessError as e:
            msg = (f"Error running generator {script}:\n"
                   f"Command: {' '.join(cmd)}\n"
                   f"Stderr: {e.stderr}\n"
                   f"Stdout: {e.stdout}")
            raise RuntimeError(msg) from e
        except Exception as e:
            raise RuntimeError(f"Unexpected error running generator {script}: {e}") from e

    def handle_command(self, name, args):
        """Dispatches CMake commands to their respective specialized handler methods."""
        # For commands, we expect args to be non-empty
        if not args:
            return

        # Variable bindings
        elif name == "set":
            self._handle_set_command(args)
        elif name == "list":
            self._handle_list_command(args)

        # Custom Dawn constructs
        elif name in ["dawngenerator", "dawnjsongenerator"]:
            self._handle_generator_command(name, args)
        elif name == "dawn_add_library":
            self._handle_add_library_command(args)

    def _handle_set_command(self, args):
        """Parses variable bindings created via CMake's 'set' command."""
        var_name = args[0]
        values = args[1:]
        self.variables[var_name] = [CMakeValue(val) for val in values]

    def _handle_list_command(self, args):
        """Parses variable modifications created via CMake's 'list' command."""
        if len(args) >= 3 and args[0].upper() == "APPEND":
            var_name = args[1]
            values = args[2:]
            if var_name not in self.variables:
                self.variables[var_name] = []
            self.variables[var_name].extend([CMakeValue(val) for val in values])

    def _handle_generator_command(self, name, args):
        """Parses and executes custom code generator commands in Dawn's build.

        This handles two custom CMake commands used in Dawn:
          1. DawnGenerator: Explicitly specifies the Python generation script via "SCRIPT".
          2. DawnJSONGenerator: A custom CMake wrapper that implicitly runs "dawn_json_generator.py"
             and translates the "TARGET" parameter into the "--targets" CLI flag.

        Both commands query the underlying Python generator scripts at rule fetch-time using the
        special "--print-cmake-outputs" flag to dynamically discover all files that will be
        produced at compile time, storing them into the respective OUTPUT_HEADERS or
        OUTPUT_SOURCES variable destinations.

        Examples of upstream CMake usage:
          # Example 1: DawnGenerator
          DawnGenerator(
              SCRIPT "${Dawn_SOURCE_DIR}/generator/dawn_version_generator.py"
              OUTPUT_HEADERS DAWN_VERSION_AUTOGEN_HEADERS
              EXTRA_PARAMETERS "--dawn-dir" "${Dawn_SOURCE_DIR}"
          )

          # Example 2: DawnJSONGenerator
          DawnJSONGenerator(
              TARGET "headers"
              PRINT_NAME "Dawn headers"
              OUTPUT_HEADERS DAWN_HEADERS_GEN_HEADERS
          )
        """
        # Skip running opengl_loader_generator.py entirely and bail out early as it has external
        # XML dependencies. Its generated files are handled statically in the Bazel BUILD target.
        if any("opengl_loader_generator.py" in arg for arg in args):
            return

        script = None
        output_headers_var = None
        output_sources_var = None
        extra_params = []

        # Emulate the custom DawnJSONGenerator CMake function defined in generator/CMakeLists.txt
        if name == "dawnjsongenerator":
            script = "generator/dawn_json_generator.py"
            # Add default dawnjsongenerator extra parameters
            extra_params.extend([
                "--dawn-json", "src/dawn/dawn.json",
                "--wire-json", "src/dawn/dawn_wire.json",
            ])
            # Once https://crbug.com/dawn/878ee25c690101566294e1ddcfd34a152adc9dc2 lands
            # we no longer need to support Dawn versions older than the commit above,
            # we can make this mandatory and remove the os.path.exists check.
            native_json = "src/dawn/dawn_native.json"
            if os.path.exists(os.path.join(self.dawn_root, native_json)):
                extra_params.extend(["--native-json", native_json])

        i = 0
        while i < len(args):
            arg_upper = args[i].upper()
            if arg_upper == "SCRIPT":
                script = args[i+1]
                i += 2
            elif arg_upper == "OUTPUT_HEADERS":
                output_headers_var = args[i+1]
                i += 2
            elif arg_upper == "OUTPUT_SOURCES":
                output_sources_var = args[i+1]
                i += 2
            elif arg_upper == "TARGET" and name == "dawnjsongenerator":
                extra_params.extend(["--targets", args[i+1]])
                i += 2
            elif arg_upper == "EXTRA_PARAMETERS":
                # Collect all subsequent arguments until another keyword or end
                i += 1
                keywords = {"SCRIPT", "OUTPUT_HEADERS", "OUTPUT_SOURCES", "EXTRA_PARAMETERS",
                            "TARGET", "PRINT_NAME"}
                while i < len(args) and args[i].upper() not in keywords:
                    extra_params.append(args[i])
                    i += 1
            else:
                i += 1

        if script:
            cleaned_script = self.clean_path(script)
            outputs = self._run_generator(cleaned_script, extra_params)

            headers = []
            sources = []
            for o in outputs:
                ext = os.path.splitext(o)[1].lower()
                if ext in [".h", ".hpp", ".inl", ".inc"]:
                    headers.append(o)
                else:
                    sources.append(o)

            # Clean output file paths
            cleaned_headers = [self.clean_path(h) for h in headers]
            cleaned_sources = [self.clean_path(s) for s in sources]
            all_cleaned_outputs = cleaned_headers + cleaned_sources

            # Register in self.generated_files to prevent folding into standard lists
            for co in all_cleaned_outputs:
                self.generated_files.add(co)

            # Strip "DAWN_" prefix and replace "_SOURCES" or "_HEADERS" with
            # "_SRCS" and "_PRIV_HDRS"
            var_to_use = output_sources_var or output_headers_var
            if var_to_use:
                base_var = var_to_use[5:] if var_to_use.startswith("DAWN_") else var_to_use
                if base_var.endswith("_SOURCES"):
                    base_var = base_var[:-8]
                elif base_var.endswith("_HEADERS"):
                    base_var = base_var[:-8]

                srcs_var = f"{base_var}_SRCS"
                priv_hdrs_var = f"{base_var}_PRIV_HDRS"

                if cleaned_sources:
                    self.generator_vars[srcs_var] = set(cleaned_sources)
                if cleaned_headers:
                    self.generator_vars[priv_hdrs_var] = set(cleaned_headers)

            # Register these files in self.variables
            if output_headers_var:
                self.variables[output_headers_var] = [CMakeValue(h) for h in headers]
            if output_sources_var:
                self.variables[output_sources_var] = [CMakeValue(s) for s in sources]

    def _handle_add_library_command(self, args):
        """Parses target library declarations created via custom dawn_add_library function."""
        target_name = args[0]
        if target_name.startswith("dawn_"):
            prefix_name = target_name[5:]
        else:
            prefix_name = target_name

        old_prefix = self.prefix
        self.prefix = prefix_name.upper()

        # We iterate through the arguments from index 1 (skipping target_name).
        all_keywords = {
            "FORCE_STATIC", "FORCE_SHARED", "FORCE_OBJECT", "HEADER_ONLY", "ENABLE_EMSCRIPTEN",
            "UTILITY_TARGET", "HEADERS", "PRIVATE_HEADERS", "SOURCES", "DEPENDS", "PRIVATE_DEPENDS"
        }
        current_keyword = None
        for arg in args[1:]:
            arg_upper = arg.upper()
            if arg_upper in all_keywords:
                current_keyword = arg_upper
            elif current_keyword in ["HEADERS", "PRIVATE_HEADERS", "SOURCES"]:
                self.add_to_outputs(current_keyword.lower(), [arg])

        self.prefix = old_prefix

    def add_to_outputs(self, var_name, values):
        """Binds extracted files/headers to their respective target prefix."""
        # Determine if we are collecting public headers (strictly HEADERS keyword)
        is_public_headers = (var_name.lower() == "headers")

        for val in values:
            if val.startswith("${") and val.endswith("}"):
                # Resolve and loop through variable expansion
                var_inner = val[2:-1]
                if var_inner in self.variables:
                    for cv in self.resolve_var(var_inner):
                        self._append_file_to_group(cv.value, is_public_headers)
            else:
                # Add single direct value
                self._append_file_to_group(val, is_public_headers)

    def _append_file_to_group(self, path, is_public_headers):
        """Cleans and appends a single file path to its target group."""
        cleaned = self.clean_path(path)
        if "${" in cleaned or cleaned in self.generated_files:
            return

        # Simple, generic platform and backend classifications based on standard keywords
        backend_suffix = ""
        lower_path = cleaned.lower()

        # 1. Android Specifics (including Android Vulkan/OpenGL files)
        if any(kw in lower_path for kw in ["android", "ahb", "ahardwarebuffer"]):
            backend_suffix = "_ANDROID"
        # 2. Unix / POSIX specific Vulkan/System files (FD, DmaBuf, OpaqueFD, Zircon)
        elif any(kw in lower_path for kw in ["dmabuf", "opaque_fd", "zircon"]) or lower_path.endswith("fd.cpp") or lower_path.endswith("fd.h"):
            backend_suffix = "_UNIX"
        # 3. Core Backends (Vulkan / OpenGL / Metal / D3D / Null / WebGPU)
        elif "/vulkan/" in cleaned:
            backend_suffix = "_VULKAN"
        elif "/metal/" in cleaned or "_metal" in lower_path:
            backend_suffix = "_METAL"
        elif "/opengl/" in cleaned:
            backend_suffix = "_OPENGL"
        elif any(kw in lower_path for kw in ["d3d", "dxgi"]):
            backend_suffix = "_D3D"
        elif "/null/" in cleaned:
            backend_suffix = "_NULL"
        elif "/webgpu/" in cleaned:
            backend_suffix = "_WEBGPU"
        # 3. Generic Apple / macOS / iOS
        elif (any(kw in lower_path for kw in ["_mac", "/mac/", "mac_", "ios", "apple", "cocoa",
                                              "iosurface", "objc", "osx"])
              or cleaned.endswith((".mm", ".m"))):
            backend_suffix = "_APPLE"
        # 4. Generic Windows
        elif any(kw in lower_path for kw in ["win", "windows"]):
            backend_suffix = "_WIN32"
        # 5. Generic Unix / X11 / Posix
        elif any(kw in lower_path for kw in ["posix", "linux", "unix", "x11"]):
            backend_suffix = "_USE_X11" if "x11" in lower_path else "_UNIX"
        # 6. DRM / GBM Specifics (Direct Rendering Manager / Generic Buffer Management)
        elif any(kw in lower_path for kw in ["drm", "gbm"]):
            backend_suffix = "_DRM"

        prefix = f"{self.prefix}{backend_suffix}"
        if prefix not in self.output_vars:
            self.output_vars[prefix] = {"HDRS": set(), "PRIV_HDRS": set(), "SRCS": set()}

        # Group by file extension and public header status
        ext = os.path.splitext(cleaned)[1].lower()
        if ext in [".h", ".hpp", ".inl", ".inc"]:
            group = "HDRS" if is_public_headers else "PRIV_HDRS"
        else:
            group = "SRCS"

        self.output_vars[prefix][group].add(cleaned)

def main():
    if len(sys.argv) < 3:
        print("Usage: generate_dawn_files.py <dawn_root> <output_bzl>")
        sys.exit(1)

    dawn_root = sys.argv[1]
    output_bzl = sys.argv[2]

    # Target directories and their base prefixes
    targets = [
        ("DAWN_ROOT", "src/dawn"),
        ("COMMON", "src/dawn/common"),
        ("NATIVE", "src/dawn/native"),
        ("WIRE", "src/dawn/wire"),
        ("PLATFORM", "src/dawn/platform"),
        ("UTILS", "src/dawn/utils"),
        ("PARTITION_ALLOC", "src/dawn/partition_alloc"),
    ]

    all_output_vars = {}

    for prefix, rel_path in targets:
        cmake_path = os.path.join(dawn_root, rel_path, "CMakeLists.txt")
        if not os.path.exists(cmake_path):
            raise FileNotFoundError(f"CMakeLists.txt not found at {cmake_path}")

        with open(cmake_path) as f:
            content = f.read()

        parser = CMakeParser(prefix, rel_path, dawn_root)
        parser.parse(content)

        # Verify that we extracted some files for this target
        has_files = False
        for groups in parser.output_vars.values():
            for files in groups.values():
                if files:
                    has_files = True
                    break
            if has_files:
                break

        if not has_files:
            raise ValueError(
                f"No file lists were generated for target '{prefix}' at '{cmake_path}'. "
                f"This indicates a parser failure or unexpected CMake structure."
            )

        # Merge parser outputs into global variable map
        for p, groups in parser.output_vars.items():
            for group, files in groups.items():
                if files:
                    var_name = f"{p}_{group}"
                    if var_name not in all_output_vars:
                        all_output_vars[var_name] = set()
                    all_output_vars[var_name].update(files)

        # Merge parser generator outputs into global variable map
        for g_var, files in parser.generator_vars.items():
            if files:
                if g_var not in all_output_vars:
                    all_output_vars[g_var] = set()
                all_output_vars[g_var].update(files)

    with open(output_bzl, "w") as f:
        f.write('"""\nThis file contains the list of files used to build Dawn.\n\nGenerated by generate_dawn_files.py\n"""\n\n')
        for var_name in sorted(all_output_vars.keys()):
            files = sorted(list(all_output_vars[var_name]))
            f.write(f"{var_name} = [\n")
            for file_path in files:
                f.write(f'    "{file_path}",\n')
            f.write("]\n\n")

if __name__ == "__main__":
    main()