#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------------------------------------
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
# CANN Open Software License Agreement Version 2.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# -----------------------------------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import subprocess
import sys
from distutils import sysconfig

try:
    from shutil import which
except ImportError:
    from distutils.spawn import find_executable as which

_COMPAT_TENSORFLOW_VERSION = "2.6"
_PYTHON_BIN_PATH_ENV = "ADAPTER_TARGET_PYTHON_PATH"
_ASCEND_HOME_PATH_ENV = "ASCEND_HOME_PATH"
_RETRY_COUNT_MAX = 5


def run_command(cmd):
    """Execute command"""
    output = subprocess.check_output(cmd)
    return output.decode('UTF-8').strip()


def get_input(q):
    """Get user input from the keyboard"""
    try:
        try:
            ans = raw_input(q)
        except NameError:
            ans = input(q)
    except EOFError:
        ans = ''
    return ans


def real_config_path(file):
    """Get complete file path"""
    os.makedirs("tools", exist_ok=True)
    return os.path.join("tools", file)


def setup_python(env_path):
    """Get python install path."""
    default_python_bin_path = sys.executable
    ask_python_bin_path = ('Please specify the location of python with valid '
                           'tensorflow 2.6 site-packages installed. [Default '
                           'is %s]\n(You can make this quiet by set env '
                           '[ADAPTER_TARGET_PYTHON_PATH]): ') % default_python_bin_path
    custom_python_bin_path = env_path if env_path else default_python_bin_path

    is_success = False
    retry_count = 0
    while not is_success:
        retry_count += 1
        if retry_count > _RETRY_COUNT_MAX:
            break

        if not custom_python_bin_path:
            python_bin_path = get_input(ask_python_bin_path)
        else:
            python_bin_path = custom_python_bin_path
            custom_python_bin_path = None
        if not python_bin_path:
            python_bin_path = default_python_bin_path
        # Check if the path is valid
        if os.path.isfile(python_bin_path) and os.access(python_bin_path, os.X_OK):
            pass
        elif not os.path.exists(python_bin_path):
            print('Invalid python path: %s cannot be found.' % python_bin_path)
            continue
        else:
            print('%s is not executable.  Is it the python binary?' % python_bin_path)
            continue

        try:
            compile_args = run_command([
                python_bin_path, '-c',
                'import distutils.sysconfig; import tensorflow as tf; print(tf.__version__ + "|" + '
                'tf.sysconfig.get_lib('') + "|" + "|".join(tf.sysconfig.get_compile_flags()) + "|" + '
                'distutils.sysconfig.get_python_inc())'
            ]).split("|")
        except subprocess.CalledProcessError:
            print('Invalid python path: %s tensorflow not installed.' %
                  python_bin_path)
            continue
        if not compile_args[0].startswith(_COMPAT_TENSORFLOW_VERSION):
            print('Invalid python path: %s compat tensorflow version is %s'
                    ' got %s.' % (python_bin_path, _COMPAT_TENSORFLOW_VERSION,
                                compile_args[0]))
            continue
        is_success = True

    if is_success:
        # Write tools/PYTHON_BIN_PATH
        with open(real_config_path('PYTHON_BIN_PATH'), 'w') as f:
            f.write(python_bin_path)

        with open(real_config_path('PYTHON_LD_LIBRARY'), 'w') as f:
            sys_vars = sysconfig.get_config_vars()
            if os.path.isfile(os.path.join(sys_vars['LIBPL'], sys_vars['LDLIBRARY'])):
                f.write(os.path.join(sys_vars['LIBPL'], sys_vars['LDLIBRARY']))
            else:
                f.write(os.path.join(sys_vars['LIBDIR'], sys_vars['LDLIBRARY']))

        with open(real_config_path('COMPILE_FLAGS'), 'w') as f:
            for flag in compile_args[2:-1]:
                f.write("".join([flag, '\n']))
            f.write("".join(["-I", compile_args[-1], '\n']))
        with open(real_config_path('TF_INSTALLED_PATH'), 'w') as f:
            f.write(compile_args[1])

    return is_success


def setup_ascend(env_path):
    """Get ascend install path."""
    default_ascend_path = os.path.realpath("/usr/local/Ascend/cann")
    ask_ascend_path = ('Please specify the location of ascend. [Default is '
                       '%s]\n(You can make this quiet by set env [ASCEND_HOME_PATH]): ') % default_ascend_path
    custom_ascend_path = env_path

    is_success = False
    retry_count = 0
    while not is_success:
        retry_count += 1
        if retry_count > _RETRY_COUNT_MAX:
            break

        if not custom_ascend_path:
            ascend_path = get_input(ask_ascend_path)
        else:
            ascend_path = custom_ascend_path
            custom_ascend_path = None
        if not ascend_path:
            ascend_path = default_ascend_path
        # Check if the path is valid
        if os.path.isdir(ascend_path) and os.access(ascend_path, os.X_OK):
            pass
        elif not os.path.exists(ascend_path):
            print('Invalid ascend path: %s cannot be found.' % ascend_path)
            continue
        is_success = True

    if is_success:
        with open(real_config_path('ASCEND_INSTALLED_PATH'), 'w') as f:
            f.write(ascend_path)

    return is_success


def main():
    """Entry point for configuration"""
    env_snapshot = dict(os.environ)
    if not setup_python(env_snapshot.get(_PYTHON_BIN_PATH_ENV)):
        sys.exit(1)
    if not setup_ascend(env_snapshot.get(_ASCEND_HOME_PATH_ENV)):
        sys.exit(1)


if __name__ == '__main__':
    main()