#!/usr/bin/env python3

import os
import re
import shutil
import struct
import subprocess
import sys

COMMAND_TO_METHOD = {
    "bash-completions": "bash_completions",
}
METHOD_TO_COMMAND = {method: command for command, method in COMMAND_TO_METHOD.items()}

WAIT_FOR_PROCESS = "wait_for_process"

class DevHelper:
    SHADER_EXTENSIONS = (".vert", ".frag", ".comp", ".geom", ".tesc", ".tese")

    def __init__(self):
        self.last_exit_status = 0

    # Core logic
    def call(self, *args, **kwargs):
        if not args:
            return

        command, rest = args[0], args[1:]
        method_name = COMMAND_TO_METHOD.get(command, command)
        method = getattr(self, method_name, None)
        if callable(method) and not method_name.startswith("_") and method_name != "call":
            return method(*rest, **kwargs)

    # Primary command wrappers
    def build(self, *args, **kwargs):
        command = f"rm -rf build && cmake -S . -B build -DCMAKE_BUILD_TYPE=Debug"
        self._exec_in_root(command, replace=True)

    def renderdoc(self, *args, **kwargs):
        qrenderdoc = self._find_qrenderdoc()
        self._exec_in_root(qrenderdoc, replace=True)

    def shaders(self, *args, **kwargs):
        shaders_dir = os.path.join(self._project_root(), "engine", "src", "shaders")
        if not os.path.isdir(shaders_dir):
            raise ValueError(f"Shaders directory not found: {shaders_dir}")

        glslc = self._find_glslc()

        shader_paths = sorted(
            os.path.join(root, filename)
            for root, _dirs, files in os.walk(shaders_dir)
            for filename in files
            if filename.endswith(self.SHADER_EXTENSIONS)
        )

        if not shader_paths:
            print(f"No shaders found in {shaders_dir}")
            return

        for shader_path in shader_paths:
            self._generate_shader_header(glslc, shader_path)

    # Custom helpers

    def ownit(self, *args, **kwargs):
        file_or_directory = args[0] if args else None
        if file_or_directory is None:
            raise ValueError("Must provide a file or directory path.")

        if sys.platform.startswith("linux"):
            print(f"Take ownership of the file or directory? {file_or_directory}")
            command = f"sudo chown -R {self._user_id()}:{self._group_id()} {file_or_directory}"
            self._exec(command, replace=True)
        else:
            raise NotImplementedError(f"Not implemented for platform {sys.platform}")

    def bash_completions(self, *args, **kwargs):
        completions = sorted(
            METHOD_TO_COMMAND.get(name, name)
            for name in vars(DevHelper)
            if not name.startswith("_") and name != "call" and callable(getattr(DevHelper, name))
        )
        print(" ".join(completions))

    # Private helpers
    def _find_glslc(self):
        glslc = shutil.which("glslc")
        if glslc is None:
            raise RuntimeError(
                "glslc not found on PATH (install the Vulkan SDK)"
            )
        return glslc

    def _find_qrenderdoc(self):
        qrenderdoc = shutil.which("qrenderdoc")
        if qrenderdoc is None:
            raise RuntimeError(
                "qrenderdoc not found on PATH (install RenderDoc)"
            )
        return qrenderdoc

    def _generate_shader_header(self, glslc, shader_path):
        relative_path = os.path.relpath(shader_path, self._project_root())
        print(f"Compiling shader: {relative_path}")

        result = subprocess.run(
            [glslc, shader_path, "-o", "-"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        if result.returncode != 0:
            raise RuntimeError(
                f"failed to compile {shader_path}: {result.stderr.decode()}"
            )

        spirv = result.stdout
        if len(spirv) % 4 != 0:
            raise RuntimeError(
                f"invalid SPIR-V output for {shader_path} (size not a multiple of 4)"
            )

        word_count = len(spirv) // 4
        words = struct.unpack(f"<{word_count}I", spirv)

        filename = os.path.basename(shader_path)
        variable_name = f"{filename.replace('.', '_')}_spv"
        header_path = shader_path + ".spv.h"

        body_lines = [
            "    " + ", ".join(f"0x{word:08x}" for word in words[i:i + 4]) + ","
            for i in range(0, len(words), 4)
        ]

        header = "\n".join([
            "#pragma once",
            "",
            "#include <array>",
            "#include <cstdint>",
            "",
            "namespace sophia::shaders {",
            "",
            f"inline constexpr std::array<uint32_t, {word_count}> {variable_name} = {{",
            *body_lines,
            "};",
            "",
            "}  // namespace sophia::shaders",
            "",
        ])

        with open(header_path, "w") as f:
            f.write(header)

        print(f"Wrote {os.path.relpath(header_path, self._project_root())} "
              f"({word_count * 4} bytes)")

    def _wait_for_process_with_logging(self, command):
        process = self._exec(command, capture=True)
        for line in process.stdout:
            print(line, end="")
        process.wait()
        self.last_exit_status = process.returncode

    def _exec_in_root(self, command, **kwargs):
        return self._exec(f"cd {self._project_root()} && {command}", **kwargs)

    def _exec(self, command, *, replace=False, capture=False):
        print(f"Running: {command}")
        if replace:
            os.execv("/bin/sh", ["/bin/sh", "-c", command])
        if capture:
            return subprocess.Popen(
                ["/bin/sh", "-c", f"{command} 2>&1"],
                stdout=subprocess.PIPE,
                text=True,
            )
        return subprocess.run(["/bin/sh", "-c", command])

    def _project_root(self):
        return os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))

    def _take_over_needed(self, file_or_directory):
        result = subprocess.run(
            ["/bin/sh", "-c", f"find {file_or_directory} -not -user {self._user_id()} -print -quit | grep -q ."],
        )
        return result.returncode == 0

    def _user_id(self):
        if not sys.platform.startswith("linux"):
            raise NotImplementedError(f"Not implemented for platform {sys.platform}")
        return subprocess.run(["id", "-u"], stdout=subprocess.PIPE, text=True).stdout.strip()

    def _group_id(self):
        if not sys.platform.startswith("linux"):
            raise NotImplementedError(f"Not implemented for platform {sys.platform}")
        return subprocess.run(["id", "-g"], stdout=subprocess.PIPE, text=True).stdout.strip()


if __name__ == "__main__":
    DevHelper().call(*sys.argv[1:])
