#!/usr/bin/env python
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.

import sys

from pathlib import Path
from shutil import copy
from setuptools import setup
from subprocess import call

import os
import platform
import xml.etree.ElementTree

from distutils.cmd import Command
from setuptools.command.build_py import build_py

## CONFIGURATION #############################################################
def pkgPlatform():
    if platform.system() == "Windows": 
        return "win10-x64"
    elif platform.system() == "Darwin":
        return "osx-x64"
    elif platform.system() == "Linux":
        return "linux-x64"
    else:
        return "unknown"

VERSION = os.getenv('ASSEMBLY_VERSION', '0.0.0')
PLATFORM = os.getenv('PKG_PLATFORM', pkgPlatform())
SRC_DIR = Path(__file__).parent.resolve()
QSHARP_PACKAGE_ROOT = SRC_DIR / "qsharp"

## NUGET DEPENDENCIES ########################################################
# We need to copy nuget dependencies into site-packages, so we first call nuget restore
# and then a nuget_copy
class NugetRestoreCommand(Command):
    description = 'Calls nuget restore to get the binaries from the QDK.'
    user_options = [
        ('packages-dir=', None, 'Specifies the location in which to look for NuGet packages.'),
        ('packages-config=', None, 'Specifies the packages.config file used to find NuGet dependencies.')
    ]

    def initialize_options(self):
        self.packages_dir = SRC_DIR / "packages"
        self.packages_config = SRC_DIR / "packages.config"

    def finalize_options(self):
        pass

    def run(self):
        call(["nuget", "restore", str(self.packages_config), "-PackagesDirectory", str(self.packages_dir)])

class CopyNugetDllsCommand(Command):
    description = 'Copies binaries dependencies into this Python package for installation.'
    user_options = [
        ('packages-dir=', None, 'Specifies the location in which to look for NuGet packages.'),
        ('packages-config=', None, 'Specifies the packages.config file used to find NuGet dependencies.')
    ]

    def initialize_options(self):
        self.packages_dir = SRC_DIR / "packages"
        self.packages_config = SRC_DIR / "packages.config"

        if self.packages_config.exists():
            self.packages_data = xml.etree.ElementTree.parse(str(self.packages_config))
        else:
            raise IOError("NuGet packages.config file {} not found.".format(self.packages_config))

    def finalize_options(self):
        pass

    def run(self):
        self.packages_dir = Path(self.packages_dir)
        if not self.packages_dir.exists():
            self.run_command('nuget_restore')
            
        lib_dlls = []

        for package_dir_node in self.packages_data.iter('package'):
            package_dir = self.packages_dir / "{0[id]}.{0[version]}".format(package_dir_node.attrib)
            framework = package_dir_node.attrib['targetFramework']

            print("Found NuGet dependency:\n  {}\n  framework: {}".format(package_dir, framework))
            pkg_lib_dir = package_dir / "lib" / framework
            pkg_runtime_dir = package_dir / "runtimes" / PLATFORM / "native"

            if pkg_lib_dir.exists():
                lib_dlls += list(pkg_lib_dir.glob("*.dll"))

            if pkg_runtime_dir.exists():
                lib_dlls += list(pkg_runtime_dir.glob("*.dll"))

        # Copy the libraries we found adjacent to qsharp/__init__.py so that they
        # appear as package data.
        print ("Copying DLLs into package root:")
        for lib_dll in lib_dlls:
            print("  - {}".format(lib_dll))
            copy(str(lib_dll), str(QSHARP_PACKAGE_ROOT))

class BuildPyCommand(build_py):
    def write_version(self):
        with open('qsharp/version.py', 'w') as f:
            f.write("""
# AUTOGENERATED, DO NOT MODIFY
version = "{}"
""".format(VERSION))

    def run(self):
        qsimDll = Path(QSHARP_PACKAGE_ROOT / "Microsoft.Quantum.Simulation.Simulators.dll")
        if not qsimDll.exists():
            self.run_command('nuget_copy')

        self.write_version()
        build_py.run(self)

## SETUP MAIN ################################################################

setup(
    cmdclass={
        'nuget_restore': NugetRestoreCommand,
        'nuget_copy': CopyNugetDllsCommand,
        'build_py': BuildPyCommand
    },
    name='qsharp',
    # NB: "-preview" is not supported as a Python version identifier.
    version=VERSION,
    description='Python interoperability layer for the Quantum Development Kit',
    author='Microsoft',
    url='http://microsoft.com/quantum/',
    packages=[
        'qsharp',
        'qsharp.tests'
    ],
    package_data={
        'qsharp': ["*.dll"]
    },
    zip_safe=False
)