#!/usr/bin/env python3

import argparse
import os
import re
import subprocess
import sys

# Parse arguments

parser = argparse.ArgumentParser(
    description="Generate a std compatibility module"
)
parser.add_argument("--src", help=(
    "Specify the location of the rust source code. The default is "
    "`$(rustc --print sysroot)/lib/rustlib/src/rust/src`"
))
args = parser.parse_args()

if args.src is None:
    output = subprocess.run(["rustc", "--print", "sysroot"],
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    args.src = os.path.join(output.stdout.decode("utf-8").strip(),
                            "lib", "rustlib", "src", "rust", "src")


# Read files

modules_regex = re.compile(
    r"^(?:\S.*)?pub\s+(?:mod\s+|use\s+(?:[a-zA-Z_][a-zA-Z0-9_]*::)*)"
    r"([a-zA-Z_][a-zA-Z0-9_]*);",
    re.MULTILINE
)


def modules(crate):
    root = os.path.join(args.src, crate)
    lib = os.path.join(root, "lib.rs")
    with open(lib) as f:
        contents = f.read()

    modules = dict()
    for match in modules_regex.finditer(contents):
        module = match.group(1)
        unstable = False

        path = os.path.join(root, module + ".rs")
        if not os.path.isfile(path):
            path = os.path.join(root, module, "mod.rs")
        try:
            with open(path, "r") as f:
                unstable = "#![unstable" in f.read()
                if unstable:
                    print(
                        f"Module '{module}' from '{crate}' appears unstable",
                        file=sys.stderr
                    )
        except OSError:
            pass

        modules[module] = unstable
    return modules


def generate(module, unstable, *namespaces):
    out = f"pub mod {module} {{\n"

    if module == "prelude":
        return None

    for namespace in namespaces:
        out += "    "

        cfgs = []
        if namespace != "core":
            cfgs.append(f"feature = \"{namespace}\"")
        if unstable:
            cfgs.append("feature = \"unstable\"")

        if len(cfgs) == 1:
            out += f"#[cfg({cfgs[0]})] "
        elif len(cfgs) > 1:
            out += "#[cfg(all(" + ", ".join(cfgs) + "))] "

        out += f"pub use __{namespace}::{module}::*;\n"

    if module == "collections":
        prefix = (
            "    #[cfg(all("
            "feature = \"alloc\", "
            "feature = \"compat_hash\""
            "))] pub use hashbrown::"
        )
        out += (
            prefix + "HashMap;\n" +
            prefix + "HashSet;\n"
        )
    elif module == "sync":
        prefix = (
            "    #[cfg(all("
            "feature = \"alloc\", "
            "feature = \"compat_sync\""
            "))] pub use spin::"
        )
        out += (
            prefix + "Mutex;\n" +
            prefix + "MutexGuard;\n" +
            prefix + "Once;\n" +
            prefix + "RwLock;\n" +
            prefix + "RwLockReadGuard;\n" +
            prefix + "RwLockWriteGuard;\n"
        )

    out += "}"
    return out


core = modules("libcore")
alloc = modules("liballoc")

generated = {}

core_keys = set(core.keys())
alloc_keys = set(alloc.keys())

for module in core_keys & alloc_keys:
    # TODO: separate these
    unstable = core[module] or alloc[module]
    generated[module] = generate(module, unstable, "core", "alloc")
for module in core_keys - alloc_keys:
    unstable = core[module]
    generated[module] = generate(module, unstable, "core")
for module in alloc_keys - core_keys:
    unstable = alloc[module]
    generated[module] = generate(module, unstable, "alloc")

generated["prelude"] = """pub mod prelude {
    pub mod v1 {
        // Prelude
        pub use __core::prelude::v1::*;
        #[cfg(all(feature = "alloc", feature = "unstable"))]
        pub use __alloc::prelude::v1::*;
        #[cfg(all(feature = "alloc", not(feature = "unstable")))]
        pub use __alloc::{
            borrow::ToOwned,
            boxed::Box,
            // UNSTABLE: slice::SliceConcatExt,
            string::String,
            string::ToString,
            vec::Vec,
        };

        // Other imports
        #[cfg(feature = "alloc")]
        pub use __alloc::{format, vec};
        #[cfg(feature = "compat_macros")]
        pub use crate::{print, println, eprint, eprintln, dbg};
    }
}"""

print("""//! Generated by generate.py located at the repository root
//! ./generate.py > src/generated.rs""")
for module in sorted(generated.items(), key=lambda i: i[0]):
    print(module[1])