# Project: com_google_tcmalloc
import json
import re
import sys

import SCons

Import("env")
Import("has_option")
Import("get_option")

env = env.Clone(
    # Building with hidden visibility interferes with intercepting the
    # libc allocation functions.
    DISALLOW_VISHIDDEN=True,
    NINJA_GENSOURCE_INDEPENDENT=True,
)

if env.Verbose():

    def tcmalloc_scons_print(msg, *args, **kwargs):
        print("[TCMALLOC_TO_SCONS]: " + msg, *args, **kwargs)
else:

    def tcmalloc_scons_print(msg, *args, **kwargs):
        pass


# manually switch this for all the debugging
tcmalloc_extra_debug = False

if tcmalloc_extra_debug:

    def tcmalloc_scons_debug(msg, *args, **kwargs):
        print("[TCMALLOC_TO_SCONS][DEBUG]: " + msg, *args, **kwargs)
else:

    def tcmalloc_scons_debug(msg, *args, **kwargs):
        pass


_bazelToSconsMap = dict(
    (f'@com_google_absl//absl/{k}', [f'$BUILD_DIR/third_party/abseil-cpp/absl_{ve}' for ve in v])
    for k, v in {
        'algorithm:container': [],
        'base:config': [],
        'base:core_headers': [],
        'base:dynamic_annotations': [],
        'container:btree': [],
        'container:fixed_array': [],
        'container:flat_hash_map': ['raw_hash_set'],
        'debugging:leak_check': [],
        'debugging:stacktrace': ['stacktrace'],
        'debugging:symbolize': [],
        'functional:function_ref': [],
        'base:malloc_internal': ['malloc_internal'],
        'memory': [],
        'numeric:bits': [],
        'numeric:int128': [],
        'strings:str_format': ['str_format_internal'],
        'types:optional': [],
        'types:span': [],
    }.items())

sys.path.append(env.Dir('scripts/site-scons').srcnode().abspath)
from bazel_to_scons import BazelEnv, Label


def dumpBazelLibs(baz, target):
    if tcmalloc_extra_debug:
        tcmalloc_scons_debug(f"Dumping tcmalloc deps to: '{target}'", file=sys.stderr)
        with open(target.abspath, 'w') as dump:
            tcmalloc_scons_debug(
                json.dumps({'libraries': baz}, sort_keys=True, indent=4), file=dump)
    else:
        pass


def _remapAbseilDep(label: Label) -> 'list[str]':
    tcmalloc_scons_print(f'Remap abseilDep {label}', file=sys.stderr)
    if str(label) in _bazelToSconsMap:
        out = _bazelToSconsMap[str(label)]
        tcmalloc_scons_print(f'Remap {label} to {out}', file=sys.stderr)
        return out

    pkg = label.package().replace('/', '_')
    tgt = label.target()
    # bazel expands //foo/bar => //foo/bar:bar implicitly. Use short form
    if tgt and not pkg.endswith('/' + tgt):
        tgt = "_" + tgt.replace('/', '_')
    else:
        tgt = ''
    return [f'$BUILD_DIR/third_party/abseil-cpp/{pkg}{tgt}']


def findAbslLibs():
    abslSconscript = env.File('$BUILD_DIR/third_party/abseil-cpp/SConscript').srcnode().abspath
    tcmalloc_scons_debug(f'abslSconscript={abslSconscript}', file=sys.stderr)
    abslLibs = []
    with open(abslSconscript) as inf:
        lines = (s.strip() for s in inf.readlines())
        targetRe = re.compile(r"\s*target=['\"](.*)['\"],")
        for line in lines:
            m = targetRe.match(line)
            if m:
                fq = f'$BUILD_DIR/third_party/abseil-cpp/{m[1]}'
                tcmalloc_scons_debug(f"found {fq} in {line}", file=sys.stderr)
                abslLibs.append(fq)
    return sorted(abslLibs)


def _mapDepToScons(lab: str, base: str = '') -> str:
    if re.match(r'^@com_google_absl//', lab):
        return _remapAbseilDep(Label(lab))
    lab = re.sub(r'^:', f'//{Label(base).package()}:', lab)
    lab = re.sub(r'^//', '', lab)
    lab = re.sub(r'(.*):(.*)', r'\1_\2', lab)
    lab = lab.replace("/", "_")
    return [lab]


def slurpBlaze(target, source, exports, env):
    bazel = BazelEnv(env, env.Dir("dist").srcnode().abspath, debug=tcmalloc_scons_debug)
    bazel.run()
    bazel.pruneTestOnlyLibraries()
    bazel.eliminateHeadersFromSources()
    bazel.eliminateSourcelessDeps()
    bzl = bazel.libraries()
    dumpBazelLibs(bzl, target)
    resolved = bazel.resolveDeps(exports)

    unknowns = [(x, resolved[x]) for x in resolved if 'unknown' in resolved[x]]
    abslImports = {}
    for unk in sorted(unknowns):
        lab = Label(unk[0])
        if lab.remote() == 'com_google_absl':
            abslImports[str(lab)] = _remapAbseilDep(lab)
    tcmalloc_scons_debug(f"{json.dumps({'abslImports': abslImports}, indent=4)}", file=sys.stderr)

    tcmalloc_scons_print('Final render into env.Library calls', file=sys.stderr)
    for libName in sorted(resolved.keys()):
        if Label(libName).remote() or libName in _bazelToSconsMap or libName not in bzl:
            continue
        libDef = bzl[libName]
        # It's the abseil name
        lab = _mapDepToScons(libName)[0]
        tcmalloc_scons_debug(f'libName: {libName:60s} => {lab}', file=sys.stderr)
        tcmalloc_scons_debug(f'    {json.dumps(list(libDef), indent=4)}', file=sys.stderr)
        kwargs = {'target': lab}
        for src in libDef.get('srcs', []):
            src = f'dist/{Label(libName).package()}/{src}'
            tcmalloc_scons_debug(f'srcs for lib={libName} -> src={src}', file=sys.stderr)
            kwargs.setdefault('source', []).append(src)
        for dep in libDef.get('deps', set()):
            scons_deps = _mapDepToScons(dep, base=libName)
            tcmalloc_scons_debug(f'lib={libName}: dep={dep} => {scons_deps}', file=sys.stderr)
            kwargs.setdefault('LIBDEPS', []).extend(scons_deps)
        if 'LIBDEPS' in kwargs:
            kwargs['LIBDEPS'] = sorted(list(set(kwargs['LIBDEPS'])))

        # libunwind and tcmalloc are both added as global dependencies. Skip
        # inheriting global dependencies to avoid a circular dependency.
        kwargs['LIBDEPS_NO_INHERIT'] = [
            '$BUILD_DIR/third_party/unwind/unwind',
            "$BUILD_DIR/third_party/tcmalloc/tcmalloc",
            "$BUILD_DIR/third_party/gperftools/tcmalloc_minimal",
        ]

        for cf in libDef.get('copts', []):
            kwargs.setdefault('CCFLAGS', [e for e in env.get('CCFLAGS', [])]).append(cf)
        tcmalloc_scons_print(f'env.Library(**{json.dumps(kwargs, indent=4)})', file=sys.stderr)
        env.BazelLibrary(**kwargs)

    return 0


env = env.Clone()
env.InjectThirdParty(libraries=['abseil-cpp'])

if get_option('link-model') == 'dynamic':
    env.Append(CPPDEFINES=[('MONGO_TCMALLOC_DYNAMIC_BUILD', 1)])

slurpBlaze(
    target=env.File('tcmalloc_deps.json').srcnode(), source=[],
    exports=['//tcmalloc', '//tcmalloc:tcmalloc_extension'], env=env)