#! /usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# 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.

import os
import copy
from urllib.parse import urlparse
from nvidia_pyindex.utils import clean_nvidia_pyindex_from_config

__all__ = [
    "create_pip_config_file",
    "maybe_edit_pip_config_file"
]

NVIDIA_PyPI_SERVERS = [
    "https://pypi.ngc.nvidia.com",
]


def _get_host_from_url(url):
    parsed_uri = urlparse(url)
    return '{uri.netloc}'.format(uri=parsed_uri)


def _format_extra_index_url(url):
    return "%s%s\n" % (" " * 18, url)


def _format_trusted_host(url):
    return "%s%s\n" % (" " * 15, _get_host_from_url(url))


def create_pip_config_file(filepath, config_filedict):

    try:
        print("Creating dir:", os.path.dirname(filepath))
        os.makedirs(os.path.dirname(filepath))
    except FileExistsError:
        pass

    content = generate_configfile_header(config_filedict)
    with open(filepath, 'w') as _file:
        content += """[global]
no-cache-dir = true
index-url = https://pypi.org/simple
extra-index-url =
{extra_index_urls}
trusted-host =
{trusted_hosts}
"""
        extra_index_urls = list()
        trusted_hosts = list()

        for nvidia_pypi_server in NVIDIA_PyPI_SERVERS:
            extra_index_urls.append(_format_extra_index_url(nvidia_pypi_server))
            trusted_hosts.append(_format_trusted_host(nvidia_pypi_server))

        extra_index_urls[-1] = extra_index_urls[-1].rstrip('\n')
        trusted_hosts[-1] = trusted_hosts[-1].rstrip('\n')

        _file.write(
            content.format(
                extra_index_urls="".join(extra_index_urls),
                trusted_hosts="".join(trusted_hosts),
            )
        )


def generate_configfile_header(config_filedict):
    header = """# This file has been autogenerated or modified by NVIDIA PyIndex.
# In case you need to modify your PIP configuration, please be aware that
# some configuration files may have a priority order. Here are the following 
# files that may exists in your machine by order of priority:\n"""

    # Priority order
    # https://github.com/pypa/pip/blob/master/src/pip/_internal/configuration.py#L47
    # OVERRIDE_ORDER = kinds.GLOBAL, kinds.USER, kinds.SITE, kinds.ENV, kinds.ENV_VAR

    header += "#\n# [Priority 1] Site level configuration files\n"
    for idx, file in enumerate(config_filedict["site"][::-1]):
        header += "#\t{idx}. `{file}`\n".format(idx=idx+1, file=file)

    header += "#\n# [Priority 2] User level configuration files\n"
    for idx, file in enumerate(config_filedict["user"][::-1]):
        header += "#\t{idx}. `{file}`\n".format(idx=idx+1, file=file)

    header += "#\n# [Priority 3] Global level configuration files\n"
    for idx, file in enumerate(config_filedict["global"][::-1]):
        header += "#\t{idx}. `{file}`\n".format(idx=idx+1, file=file)

    header += "\n"

    return header


def maybe_edit_pip_config_file(filepath, config_filedict, debug_mode=False):
    
    clean_nvidia_pyindex_from_config(filepath)
    
    with open(filepath, 'r') as _file:
        original_content = _file.readlines()

    # Force `no-cache-dir = true`
    for i in range(len(original_content)):
        if "no-cache-dir" in original_content[i]:
            original_content[i] = "no-cache-dir = true\n"
            try:
                if original_content[i+1][0] == " ":
                    del original_content[i+1]
            except IndexError:
                pass
            break
    else:
        original_content.append("no-cache-dir = true\n")

    # Add the `extra-index-url(s)` and `trusted-host` for nvidia python packages
    for nvidia_pypi_server in NVIDIA_PyPI_SERVERS:

        def edit_pip_environment(env_name, value_to_add):
            temp = []
            file_content = copy.deepcopy(original_content)
            try:
                while file_content:
                    line = file_content.pop(0)
                    if env_name in line:
                        temp.append(line)
                        temp.append(value_to_add)
                        while True:
                            line = file_content.pop(0)
                            temp.append(line)
                            if line[0] != " ":
                                break
                        temp += file_content
                        break
                    else:
                        temp.append(line)
                # the env_name wasn't found in the existing config file
                else:
                    temp += [
                        "{} =\n".format(env_name),
                        value_to_add
                    ]
            except IndexError:
                pass

            return temp

        original_content = edit_pip_environment(
            env_name="extra-index-url",
            value_to_add=_format_extra_index_url(nvidia_pypi_server)
        )

        original_content = edit_pip_environment(
            env_name="trusted-host",
            value_to_add=_format_trusted_host(nvidia_pypi_server)
        )

        if not any("autogenerated or modified by NVIDIA PyIndex" in line for line in original_content):
            original_content = (
                generate_configfile_header(config_filedict) +
                "".join(original_content)
            )

    with open(filepath, 'w') as _file:
        _file.writelines(original_content)
