#!/usr/bin/env python

# Copyright 2020 Stanford University, Los Alamos National Laboratory
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from __future__ import absolute_import, division, print_function, unicode_literals

import argparse
import os
import subprocess

def find_flexflow_header(ffhome_dir):
    def try_prefix(prefix_dir):
        flexflow_ch_path = os.path.join(prefix_dir, 'python', 'flexflow_c.h')
        flexflow_cxxh_path = os.path.join(prefix_dir, 'include', 'model.h')
        if os.path.exists(flexflow_ch_path) and os.path.exists(flexflow_cxxh_path):
            flexflow_cxxh_dir = os.path.join(prefix_dir, 'include')
            return flexflow_cxxh_dir, flexflow_ch_path

    result = try_prefix(ffhome_dir)
    if result:
        return result

    raise Exception('Unable to locate flexflow_c.h and flexflow.h header file')

def build(output_dir, libname, ffhome_dir):
    flexflow_cxxh_dir, flexflow_ch_path = find_flexflow_header(ffhome_dir)

    header = subprocess.check_output(['gcc', '-I', flexflow_cxxh_dir, '-E', '-P', flexflow_ch_path]).decode('utf-8')

    with open(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'flexflow_cffi_header.py.in')) as f:
        content = f.read()
    content = content.format(header=repr(header), libname=repr(libname))

    if output_dir is None:
        output_dir = os.path.dirname(os.path.realpath(__file__))

    with open(os.path.join(output_dir, 'flexflow_cffi_header.py'), 'wb') as f:
        f.write(content.encode('utf-8'))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--ffhome-dir', required=True)
    parser.add_argument('--libname', required=True)
    parser.add_argument('--output-dir', required=False)
    args = parser.parse_args()

    build(args.output_dir, args.libname, args.ffhome_dir)

[((54, 13, 54, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((26, 27, 26, 77), 'os.path.join', 'os.path.join', ({(26, 40, 26, 50): 'prefix_dir', (26, 52, 26, 60): '"""python"""', (26, 62, 26, 76): '"""flexflow_c.h"""'}, {}), "(prefix_dir, 'python', 'flexflow_c.h')", False, 'import os\n'), ((27, 29, 27, 75), 'os.path.join', 'os.path.join', ({(27, 42, 27, 52): 'prefix_dir', (27, 54, 27, 63): '"""include"""', (27, 65, 27, 74): '"""model.h"""'}, {}), "(prefix_dir, 'include', 'model.h')", False, 'import os\n'), ((28, 11, 28, 43), 'os.path.exists', 'os.path.exists', ({(28, 26, 28, 42): 'flexflow_ch_path'}, {}), '(flexflow_ch_path)', False, 'import os\n'), ((28, 48, 28, 82), 'os.path.exists', 'os.path.exists', ({(28, 63, 28, 81): 'flexflow_cxxh_path'}, {}), '(flexflow_cxxh_path)', False, 'import os\n'), ((29, 32, 29, 67), 'os.path.join', 'os.path.join', ({(29, 45, 29, 55): 'prefix_dir', (29, 57, 29, 66): '"""include"""'}, {}), "(prefix_dir, 'include')", False, 'import os\n'), ((41, 13, 41, 100), 'subprocess.check_output', 'subprocess.check_output', ({(41, 37, 41, 99): "['gcc', '-I', flexflow_cxxh_dir, '-E', '-P', flexflow_ch_path]"}, {}), "(['gcc', '-I', flexflow_cxxh_dir, '-E', '-P',\n    flexflow_ch_path])", False, 'import subprocess\n'), ((48, 37, 48, 63), 'os.path.realpath', 'os.path.realpath', ({(48, 54, 48, 62): '__file__'}, {}), '(__file__)', False, 'import os\n'), ((50, 14, 50, 65), 'os.path.join', 'os.path.join', ({(50, 27, 50, 37): 'output_dir', (50, 39, 50, 64): '"""flexflow_cffi_header.py"""'}, {}), "(output_dir, 'flexflow_cffi_header.py')", False, 'import os\n'), ((43, 43, 43, 69), 'os.path.realpath', 'os.path.realpath', ({(43, 60, 43, 68): '__file__'}, {}), '(__file__)', False, 'import os\n')]