#!/usr/bin/env python3
import os
import shutil

def safe_symlink(src, dst):
    if os.path.exists(dst) or os.path.islink(dst):
        if os.path.isdir(dst) and not os.path.islink(dst):
            shutil.rmtree(dst)
        else:
            os.unlink(dst)
    os.symlink(src, dst)

def main():
    # 1. Clean up any old mock_repos and recreate it from our static packaging template
    dest_dir = "./mock_repos"
    if os.path.exists(dest_dir):
        shutil.rmtree(dest_dir)
    shutil.copytree("debian/mock_repos", dest_dir)
    print("Copied static mock templates to mock_repos/")

    # 2. Write workspace_root.bzl into rules_jvm_external mock
    rje_dir = os.path.join(dest_dir, "rules_jvm_external")
    with open(os.path.join(rje_dir, "workspace_root.bzl"), "w") as f:
        f.write(f'WORKSPACE_ROOT = "{os.path.abspath(os.getcwd())}"\n')
    print("Generated rules_jvm_external/workspace_root.bzl")

    # 3. Setup protobuf mock dynamic files
    protobuf_dir = os.path.join(dest_dir, "protobuf")
    
    # Copy well-known protos from system path
    proto_src_dir = os.path.join(protobuf_dir, "google", "protobuf")
    os.makedirs(proto_src_dir, exist_ok=True)
    system_proto_dir = "/usr/include/google/protobuf"
    well_known_protos = ["any.proto", "wrappers.proto", "descriptor.proto", "timestamp.proto", "duration.proto", "empty.proto"]
    for proto in well_known_protos:
        src_path = os.path.join(system_proto_dir, proto)
        dest_path = os.path.join(proto_src_dir, proto)
        if os.path.exists(src_path):
            shutil.copy2(src_path, dest_path)
        else:
            with open(dest_path, "w") as f:
                f.write(f"// Mock {proto}\n")
    print("Copied well-known proto schemas")

    # Symlink system Jars into the packages
    safe_symlink("/usr/share/java/protobuf.jar", os.path.join(protobuf_dir, "protobuf-java.jar"))
    safe_symlink("/usr/share/java/protobuf-util.jar", os.path.join(protobuf_dir, "protobuf-java-util.jar"))
    
    java_core_dir = os.path.join(protobuf_dir, "java", "core")
    java_util_dir = os.path.join(protobuf_dir, "java", "util")
    safe_symlink("/usr/share/java/protobuf.jar", os.path.join(java_core_dir, "protobuf-java.jar"))
    safe_symlink("/usr/share/java/protobuf-util.jar", os.path.join(java_util_dir, "protobuf-java-util.jar"))
    print("Linked system Protobuf Java libraries")

    # Setup blake3 mock dynamic files: bridge system include path (usr/include/blake3.h -> c/blake3.h)
    blake3_dir = os.path.join(dest_dir, "blake3")
    blake3_c_dir = os.path.join(blake3_dir, "c")
    os.makedirs(blake3_c_dir, exist_ok=True)
    system_header = "/usr/include/blake3.h"
    mock_header = os.path.join(blake3_c_dir, "blake3.h")
    if os.path.exists(system_header) and not os.path.exists(mock_header):
        os.symlink(system_header, mock_header)
    print("Bridged system blake3 header")

    # Setup rules_jvm_external mock: generate a valid, minimal 22-byte empty zip archive
    jvm_dir = os.path.join(dest_dir, "rules_jvm_external")
    os.makedirs(jvm_dir, exist_ok=True)
    empty_jar_path = os.path.join(jvm_dir, "empty.jar")
    import zipfile
    with zipfile.ZipFile(empty_jar_path, "w") as z:
        pass
    print("Generated valid empty Jar template")

    # Setup zstd-jni mock dynamic files
    zstd_dir = os.path.join(dest_dir, "zstd-jni")
    safe_symlink("/usr/share/java/zstd-jni.jar", os.path.join(zstd_dir, "zstd-jni.jar"))
    print("Linked system Zstd-JNI libraries")

    # 4. Symlink system rules packages (bazel_skylib, rules_cc, etc.)
    system_mappings = {
        "bazel_skylib": "/usr/share/bazel/tools/skylib",
        "bazel_features": "/usr/share/bazel/tools/features",
        "rules_python": "/usr/share/bazel/rules/python",
        "rules_java": "/usr/share/bazel/rules/java",
        "platforms": "/usr/share/bazel/tools/platforms",
        "rules_license": "/usr/share/bazel/rules/license",
        "rules_pkg": "/usr/share/bazel/rules/pkg",
        "rules_cc": "/usr/share/bazel/rules/cc",
        "rules_proto": "/usr/share/bazel/rules/proto"
    }

    for mod, system_dir in system_mappings.items():
        mod_dir = os.path.join(dest_dir, mod)
        os.makedirs(mod_dir, exist_ok=True)
        # Write clean Bzlmod metadata if they don't exist
        if not os.path.exists(os.path.join(mod_dir, "MODULE.bazel")):
            with open(os.path.join(mod_dir, "MODULE.bazel"), "w") as f:
                f.write(f'module(name = "{mod}")\n')
        if not os.path.exists(os.path.join(mod_dir, "BUILD")):
            with open(os.path.join(mod_dir, "BUILD"), "w") as f:
                f.write("# Dummy BUILD\n")

        # Symlink all files from system rules folder
        if os.path.exists(system_dir):
            for item in os.listdir(system_dir):
                src = os.path.join(system_dir, item)
                dst = os.path.join(mod_dir, item)
                safe_symlink(src, dst)
            print(f"Symlinked system {mod} rules")

        # Ensure LICENSE and AUTHORS are present (even as dummies) to satisfy source tarball packaging!
        for metadata_file in ["LICENSE", "AUTHORS"]:
            metadata_path = os.path.join(mod_dir, metadata_file)
            if not os.path.exists(metadata_path):
                with open(metadata_path, "w") as f:
                    f.write(f"Mock {metadata_file} for Bzlmod packaging\n")

    # Setup zlib mock to link against system zlib
    zlib_dir = os.path.join(dest_dir, "zlib")
    os.makedirs(zlib_dir, exist_ok=True)
    with open(os.path.join(zlib_dir, "MODULE.bazel"), "w") as f:
        f.write('module(name = "zlib")\n')
    with open(os.path.join(zlib_dir, "BUILD"), "w") as f:
        f.write('''
cc_library(
    name = "zlib",
    linkopts = ["-lz"],
    visibility = ["//visibility:public"],
)
''')
    print("Generated system-linking mock for zlib")

    # 5. Generate dummy mocks for other overridden unused transitive dependencies
    dummy_modules = [
        "grpc",
        "rules_testing",
        "googletest",
        "with_cfg.bzl",
        "apple_support",
        "abseil-cpp",
        "c-ares",
        "rules_go",
        "rules_kotlin",
        "upb",
        "buildozer"
    ]

    for mod in dummy_modules:
        mod_dir = os.path.join(dest_dir, mod)
        os.makedirs(mod_dir, exist_ok=True)
        with open(os.path.join(mod_dir, "MODULE.bazel"), "w") as f:
            f.write(f'module(name = "{mod}")\n')
        with open(os.path.join(mod_dir, "BUILD"), "w") as f:
            f.write("# Dummy BUILD\n")
    print("Generated dummy mocks for unused transitive dependencies")

    # Setup rules_graalvm mock with dummy native_image macro
    graalvm_dir = os.path.join(dest_dir, "rules_graalvm")
    os.makedirs(os.path.join(graalvm_dir, "graalvm"), exist_ok=True)
    with open(os.path.join(graalvm_dir, "MODULE.bazel"), "w") as f:
        f.write('module(name = "rules_graalvm")\n')
    with open(os.path.join(graalvm_dir, "BUILD"), "w") as f:
        f.write("# Dummy BUILD\n")
    with open(os.path.join(graalvm_dir, "graalvm", "BUILD"), "w") as f:
        f.write("# Dummy BUILD\n")
    with open(os.path.join(graalvm_dir, "graalvm", "defs.bzl"), "w") as f:
        f.write('''
def native_image(**kwargs):
    pass
''')
    print("Generated rules_graalvm mock with dummy native_image macro")

    # Setup remote_java_tools and remote_java_tools_linux mocks
    for repo in ["remote_java_tools", "remote_java_tools_linux"]:
        repo_dir = os.path.join(dest_dir, repo)
        os.makedirs(repo_dir, exist_ok=True)
        # These are repositories, they don't need MODULE.bazel if they are overridden as repositories,
        # but they need a WORKSPACE file.
        with open(os.path.join(repo_dir, "WORKSPACE"), "w") as f:
            f.write(f'workspace(name = "{repo}")\n')
        with open(os.path.join(repo_dir, "BUILD"), "w") as f:
            f.write('''
package(default_visibility = ["//visibility:public"])
filegroup(name = "ijar", srcs = ["ijar_bin"])
filegroup(name = "singlejar", srcs = ["singlejar_bin"])
filegroup(name = "xmldoclet", srcs = ["xmldoclet_deploy.jar"])
filegroup(name = "javac_jar", srcs = ["javac-9+181-rxtx.jar"])
filegroup(name = "JavaBuilder", srcs = ["JavaBuilder_deploy.jar"])
filegroup(name = "VanillaJavaBuilder", srcs = ["VanillaJavaBuilder_deploy.jar"])
filegroup(name = "GenClass", srcs = ["GenClass_deploy.jar"])
filegroup(name = "turbine_direct", srcs = ["TurbineDirect_deploy.jar"])
filegroup(name = "turbine_direct_graal", srcs = ["turbine_direct_graal_bin"])
filegroup(name = "Runner", srcs = ["Runner_deploy.jar"])
filegroup(name = "JacocoCoverage", srcs = ["jacoco_coverage_runner_deploy.jar"])
filegroup(name = "one_version_tool", srcs = ["one_version_tool_bin"])
filegroup(name = "java_tools", srcs = [])
filegroup(name = "TurbineDirect", srcs = ["TurbineDirect_deploy.jar"])
filegroup(name = "jacoco_coverage_runner_filegroup", srcs = ["jacoco_coverage_runner_deploy.jar"])
filegroup(name = "ijar_cc_binary", srcs = ["ijar_bin"])
filegroup(name = "singlejar_cc_bin", srcs = ["singlejar_bin"])
filegroup(name = "ijar_prebuilt_binary", srcs = ["ijar_bin"])
filegroup(name = "prebuilt_singlejar", srcs = ["singlejar_bin"])
''')
    print("Generated remote_java_tools and remote_java_tools_linux mocks")

    # Setup remotejdk mocks for all architectures and versions
    versions = ["8", "11", "15", "16", "17", "21", "22", "23", "24"]
    platforms = [
        "linux", "linux_aarch64", "linux_arm64", "linux_ppc64le", "linux_ppc64", "linux_ppc",
        "linux_s390x", "linux_riscv64", "linux_loongarch64", "linux_loong64",
        "linux_mips64el", "linux_mips64", "linux_mips", "linux_arm", "linux_armhf",
        "linux_sparc64", "linux_sparc", "linux_x86_64", "linux_amd64",
        "macos", "macos_aarch64", "macos_x86_64", "win", "win_arm64", "win_x86_64", "windows"
    ]
    for ver in versions:
        for plat in platforms:
            for prefix in ["remotejdk", "remote_jdk"]:
                jdk = f"{prefix}{ver}_{plat}"
                repo_dir = os.path.join(dest_dir, jdk)
                os.makedirs(repo_dir, exist_ok=True)
                with open(os.path.join(repo_dir, "WORKSPACE"), "w") as f:
                    f.write(f'workspace(name = "{jdk}")\n')
                with open(os.path.join(repo_dir, "BUILD"), "w") as f:
                    f.write('''
package(default_visibility = ["//visibility:public"])
java_runtime(
    name = "jdk",
    srcs = [],
    java_home = "/usr",
)
''')
    print("Generated remotejdk mocks for all architectures")

    # Defensive self-healing check: ensure mock protoc script is executable!
    protoc_script = os.path.join(protobuf_dir, "protoc.sh")
    if os.path.exists(protoc_script):
        os.chmod(protoc_script, 0o755)

    print("Successfully initialized all offline Bzlmod mocks!")

if __name__ == "__main__":
    main()
