#!/usr/bin/env python3
# ################################################################
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under both the BSD-style license (found in the
# LICENSE file in the root directory of this source tree) and the GPLv2 (found
# in the COPYING file in the root directory of this source tree).
# You may select, at your option, one of the above-listed licenses.
# ################################################################

"""Flatten Cargo's staticlib into Meson's libzstd static archive."""

from __future__ import print_function

import argparse
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile


def run(command, **kwargs):
    try:
        return subprocess.run(command, check=True, **kwargs)
    except subprocess.CalledProcessError as error:
        print('command failed: {}'.format(' '.join(command)), file=sys.stderr)
        return error


def extract_members(archiver, archive, destination, prefix):
    members = subprocess.check_output(
        [archiver, 't', str(archive)], universal_newlines=True
    ).splitlines()
    if not members:
        raise RuntimeError('{} contains no object members'.format(archive))

    objects = []
    for index, member in enumerate(members):
        object_file = destination / '{}-{:04d}.o'.format(prefix, index)
        with object_file.open('wb') as output:
            run_result = run(
                [archiver, 'p', str(archive), member], stdout=output
            )
        if isinstance(run_result, subprocess.CalledProcessError):
            raise RuntimeError('could not extract {} from {}'.format(member, archive))
        objects.append(object_file)
    return objects


def merge_with_ar(args, output):
    with tempfile.TemporaryDirectory(prefix='zstd-rust-archive-') as temporary:
        temporary_path = Path(temporary)
        c_objects = extract_members(
            args.archiver, Path(args.c_archive), temporary_path, 'c'
        )
        rust_objects = extract_members(
            args.archiver, Path(args.rust_archive), temporary_path, 'rust'
        )
        temporary_output = temporary_path / output.name
        run_result = run(
            [args.archiver, 'rcs', str(temporary_output)]
            + [str(object_file) for object_file in c_objects + rust_objects]
        )
        if isinstance(run_result, subprocess.CalledProcessError):
            raise RuntimeError('could not create {}'.format(output))
        shutil.move(str(temporary_output), str(output))


def merge_with_msvc(args, output):
    temporary_output = output.with_name(output.name + '.tmp')
    run_result = run(
        [
            args.archiver,
            '/OUT:{}'.format(temporary_output),
            args.c_archive,
            args.rust_archive,
        ]
    )
    if isinstance(run_result, subprocess.CalledProcessError):
        raise RuntimeError('could not create {}'.format(output))
    os.replace(str(temporary_output), str(output))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--archiver', required=True)
    parser.add_argument('--c-archive', required=True)
    parser.add_argument('--rust-archive', required=True)
    parser.add_argument('--output', required=True)
    parser.add_argument('--msvc', action='store_true')
    args = parser.parse_args()

    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    try:
        if args.msvc:
            merge_with_msvc(args, output)
        else:
            merge_with_ar(args, output)
    except (OSError, RuntimeError, subprocess.CalledProcessError) as error:
        print('could not merge Rust archive: {}'.format(error), file=sys.stderr)
        return 1
    return 0


if __name__ == '__main__':
    sys.exit(main())