"""msvc9compiler monkey patch test

This test ensures that importing setuptools is sufficient to replace
the standard find_vcvarsall function with our patched version that
finds the Visual C++ for Python package.
"""

import os
import shutil
import tempfile
import unittest
import distutils.errors

import pytest

from . import contexts

# importing only setuptools should apply the patch
__import__('setuptools')

pytest.importorskip("distutils.msvc9compiler")

class MockReg:
    """Mock for distutils.msvc9compiler.Reg. We patch it
    with an instance of this class that mocks out the
    functions that access the registry.
    """

    def __init__(self, hkey_local_machine={}, hkey_current_user={}):
        self.hklm = hkey_local_machine
        self.hkcu = hkey_current_user

    def __enter__(self):
        self.original_read_keys = distutils.msvc9compiler.Reg.read_keys
        self.original_read_values = distutils.msvc9compiler.Reg.read_values

        _winreg = getattr(distutils.msvc9compiler, '_winreg', None)
        winreg = getattr(distutils.msvc9compiler, 'winreg', _winreg)

        hives = {
            winreg.HKEY_CURRENT_USER: self.hkcu,
            winreg.HKEY_LOCAL_MACHINE: self.hklm,
        }

        def read_keys(cls, base, key):
            """Return list of registry keys."""
            hive = hives.get(base, {})
            return [k.rpartition('\\')[2]
                    for k in hive if k.startswith(key.lower())]

        def read_values(cls, base, key):
            """Return dict of registry keys and values."""
            hive = hives.get(base, {})
            return dict((k.rpartition('\\')[2], hive[k])
                        for k in hive if k.startswith(key.lower()))

        distutils.msvc9compiler.Reg.read_keys = classmethod(read_keys)
        distutils.msvc9compiler.Reg.read_values = classmethod(read_values)

        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        distutils.msvc9compiler.Reg.read_keys = self.original_read_keys
        distutils.msvc9compiler.Reg.read_values = self.original_read_values


class TestMSVC9Compiler(unittest.TestCase):

    def test_find_vcvarsall_patch(self):
        mod_name = distutils.msvc9compiler.find_vcvarsall.__module__
        assert mod_name == "setuptools.msvc9_support", "find_vcvarsall unpatched"

        find_vcvarsall = distutils.msvc9compiler.find_vcvarsall
        query_vcvarsall = distutils.msvc9compiler.query_vcvarsall

        # No registry entries or environment variable means we should
        # not find anything
        with contexts.environment(VS90COMNTOOLS=None):
            with MockReg():
                self.assertIsNone(find_vcvarsall(9.0))

                expected = distutils.errors.DistutilsPlatformError
                with pytest.raises(expected) as exc:
                    query_vcvarsall(9.0)
                assert 'aka.ms/vcpython27' in str(exc)

        key_32 = r'software\microsoft\devdiv\vcforpython\9.0\installdir'
        key_64 = r'software\wow6432node\microsoft\devdiv\vcforpython\9.0\installdir'

        # Make two mock files so we can tell whether HCKU entries are
        # preferred to HKLM entries.
        mock_installdir_1 = tempfile.mkdtemp()
        mock_vcvarsall_bat_1 = os.path.join(mock_installdir_1, 'vcvarsall.bat')
        open(mock_vcvarsall_bat_1, 'w').close()
        mock_installdir_2 = tempfile.mkdtemp()
        mock_vcvarsall_bat_2 = os.path.join(mock_installdir_2, 'vcvarsall.bat')
        open(mock_vcvarsall_bat_2, 'w').close()
        try:
            # Ensure we get the current user's setting first
            with MockReg(
                hkey_current_user={key_32: mock_installdir_1},
                hkey_local_machine={
                    key_32: mock_installdir_2,
                    key_64: mock_installdir_2,
                }
            ):
                self.assertEqual(mock_vcvarsall_bat_1, find_vcvarsall(9.0))

            # Ensure we get the local machine value if it's there
            with MockReg(hkey_local_machine={key_32: mock_installdir_2}):
                self.assertEqual(mock_vcvarsall_bat_2, find_vcvarsall(9.0))

            # Ensure we prefer the 64-bit local machine key
            # (*not* the Wow6432Node key)
            with MockReg(
                hkey_local_machine={
                    # This *should* only exist on 32-bit machines
                    key_32: mock_installdir_1,
                    # This *should* only exist on 64-bit machines
                    key_64: mock_installdir_2,
                }
            ):
                self.assertEqual(mock_vcvarsall_bat_1, find_vcvarsall(9.0))
        finally:
            shutil.rmtree(mock_installdir_1)
            shutil.rmtree(mock_installdir_2)