#  Copyright (c) 2017-2018 Uber Technologies, Inc.
#
# 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 pickle
import textwrap
import unittest
from typing import Dict

import pytest
from pyarrow.lib import ArrowIOError

from unittest import mock

from petastorm.hdfs.namenode import HdfsNamenodeResolver, HdfsConnector, \
    HdfsConnectError, MaxFailoversExceeded, HAHdfsClient, namenode_failover


class HC:
    """Hadoop constants for testing convenience"""
    WARP_TURTLE = 'WARP-TURTLE'
    FS_WARP_TURTLE = 'hdfs://{}'.format(WARP_TURTLE)
    DEFAULT_NN = 'default:8020'
    WARP_TURTLE_NN1 = 'some.domain.name.net:8020'
    WARP_TURTLE_NN2 = 'some.other.domain.name.net:8020'
    WARP_TURTLE_PATH = '{}/x/y/z'.format(FS_WARP_TURTLE)
    HADOOP_CONFIG_PATH = '/etc/hadoop'


class MockHadoopConfiguration(object):
    def __init__(self):
        self._dict = {}

    def get(self, key):
        val = None
        if key in self._dict:
            val = self._dict[key]
        # print('MockHadoopConfiguration: "{}" == "{}"'.format(key, val))
        return val

    def set(self, key, val):
        self._dict[key] = val


class HdfsNamenodeResolverTest(unittest.TestCase):
    def setUp(self):
        """Initializes a mock hadoop config and a namenode resolver instance, for convenience."""
        self._hadoop_configuration = MockHadoopConfiguration()
        self.suj = HdfsNamenodeResolver(self._hadoop_configuration)

    def test_default_hdfs_service_errors(self):
        """Check error cases with connecting to default namenode"""
        # No default yields RuntimeError
        with self.assertRaises(RuntimeError):
            self.suj.resolve_default_hdfs_service()
        # Bad default FS yields IOError
        self._hadoop_configuration.set('fs.defaultFS', 'invalidFS')
        with self.assertRaises(IOError):
            self.suj.resolve_default_hdfs_service()
        # Random FS host yields IOError
        self._hadoop_configuration.set('fs.defaultFS', 'hdfs://random')
        with self.assertRaises(IOError):
            self.suj.resolve_default_hdfs_service()
        # Valid FS host with no namenode defined yields IOError
        self._hadoop_configuration.set('fs.defaultFS', HC.FS_WARP_TURTLE)
        with self.assertRaises(IOError):
            self.suj.resolve_default_hdfs_service()

    def test_default_hdfs_service_typical(self):
        """Check typical cases resolving default namenode"""
        # One nn
        self._hadoop_configuration.set('fs.defaultFS', HC.FS_WARP_TURTLE)
        self._hadoop_configuration.set('dfs.ha.namenodes.{}'.format(HC.WARP_TURTLE), 'nn1')
        self._hadoop_configuration.set(
            'dfs.namenode.rpc-address.{}.nn1'.format(HC.WARP_TURTLE), HC.WARP_TURTLE_NN1)
        nameservice, namenodes = self.suj.resolve_default_hdfs_service()
        self.assertEqual(HC.WARP_TURTLE, nameservice)
        self.assertEqual(HC.WARP_TURTLE_NN1, namenodes[0])

        # Second of two nns, when the first is undefined
        self._hadoop_configuration.set('dfs.ha.namenodes.{}'.format(HC.WARP_TURTLE), 'nn2,nn1')
        with self.assertRaises(RuntimeError):
            self.suj.resolve_default_hdfs_service()

        # Two valid and defined nns
        self._hadoop_configuration.set(
            'dfs.namenode.rpc-address.{}.nn2'.format(HC.WARP_TURTLE), HC.WARP_TURTLE_NN2)
        nameservice, namenodes = self.suj.resolve_default_hdfs_service()
        self.assertEqual(HC.WARP_TURTLE, nameservice)
        self.assertEqual(HC.WARP_TURTLE_NN2, namenodes[0])
        self.assertEqual(HC.WARP_TURTLE_NN1, namenodes[1])

    def test_resolve_hdfs_name_service(self):
        """Check edge cases with resolving a nameservice"""
        # Most cases already covered by test_default_hdfs_service_ok above...
        # Empty config or no namespace yields None
        self.assertIsNone(HdfsNamenodeResolver({}).resolve_hdfs_name_service(''))
        self.assertIsNone(self.suj.resolve_hdfs_name_service(''))

        # Test a single undefined namenode case, as well as an unconventional multi-NN case;
        # both result in an exception raised
        self._hadoop_configuration.set('fs.defaultFS', HC.FS_WARP_TURTLE)
        self._hadoop_configuration.set('dfs.ha.namenodes.{}'.format(HC.WARP_TURTLE), 'nn1')
        with self.assertRaises(RuntimeError):
            self.suj.resolve_hdfs_name_service(HC.WARP_TURTLE)

        # Test multiple undefined NNs, which will also throw HdfsConnectError
        nns = 'nn1,nn2,nn3,nn4,nn5,nn6,nn7,nn8'
        self._hadoop_configuration.set('dfs.ha.namenodes.{}'.format(HC.WARP_TURTLE), nns)
        with self.assertRaises(RuntimeError):
            self.suj.resolve_hdfs_name_service(HC.WARP_TURTLE)


@pytest.fixture()
def mock_hadoop_home_directory(tmpdir):
    """Create hadoop site files once"""
    tmpdir_path = tmpdir.strpath
    os.makedirs('{}{}'.format(tmpdir_path, HC.HADOOP_CONFIG_PATH))
    with open('{}{}/core-site.xml'.format(tmpdir_path, HC.HADOOP_CONFIG_PATH), 'wt') as f:
        f.write(textwrap.dedent("""\
            <?xml version="1.0"?>
            <?xml-stylesheet type="text/xsl" href="configuration.xsl"?>
            <configuration>
              <property>
                <name>fs.defaultFS</name>
                <value>hdfs://{0}</value>
              </property>
            </configuration>
            """.format(HC.WARP_TURTLE)))
    with open('{}{}/hdfs-site.xml'.format(tmpdir_path, HC.HADOOP_CONFIG_PATH), 'wt') as f:
        f.write(textwrap.dedent("""\
            <?xml version="1.0"?>
            <?xml-stylesheet type="text/xsl" href="configuration.xsl"?>
            <configuration>
              <property>
                <name>dfs.ha.namenodes.{0}</name>
                <value>nn2,nn1</value>
              </property>
              <property>
                <name>dfs.namenode.rpc-address.{0}.nn1</name>
                <value>{1}</value>
              </property>
              <property>
                <name>dfs.namenode.rpc-address.{0}.nn2</name>
                <value>{2}</value>
              </property>
              <property>
                <name>dfs.ha.namenodes.foobar</name>
                <value>nn</value>
              </property>
            </configuration>
            """.format(HC.WARP_TURTLE, HC.WARP_TURTLE_NN1, HC.WARP_TURTLE_NN2)))
    return tmpdir_path


def _test_default_hdfs_service(mock_hadoop_home_directory, env_var):
    # Trigger env var evaluation
    suj = HdfsNamenodeResolver()
    assert env_var == suj._hadoop_env
    assert mock_hadoop_home_directory == suj._hadoop_path
    # List of namenodes returned nominally
    nameservice, namenodes = suj.resolve_default_hdfs_service()
    assert HC.WARP_TURTLE == nameservice
    assert HC.WARP_TURTLE_NN2 == namenodes[0]
    assert HC.WARP_TURTLE_NN1 == namenodes[1]
    # Exception raised for badly defined nameservice (XML issue)
    with pytest.raises(RuntimeError):
        suj.resolve_hdfs_name_service('foobar')
    # None for nonexistent nameservice (intentional design)
    assert suj.resolve_hdfs_name_service('nonexistent') is None


def test_env_hadoop_home_prefix_install(mock_hadoop_home_directory):
    # The second+third env vars won't cause an error
    with mock.patch.dict(os.environ, {'HADOOP_PREFIX': '{}/no/where/here'.format(mock_hadoop_home_directory),
                                      'HADOOP_INSTALL': '{}/no/where/here'.format(mock_hadoop_home_directory),
                                      'HADOOP_HOME': mock_hadoop_home_directory}, clear=True):
        _test_default_hdfs_service(mock_hadoop_home_directory, 'HADOOP_HOME')


def test_env_hadoop_prefix_only(mock_hadoop_home_directory):
    with mock.patch.dict(os.environ, {'HADOOP_PREFIX': mock_hadoop_home_directory}, clear=True):
        _test_default_hdfs_service(mock_hadoop_home_directory, 'HADOOP_PREFIX')


def test_env_hadoop_install_only(mock_hadoop_home_directory):
    with mock.patch.dict(os.environ, {'HADOOP_INSTALL': mock_hadoop_home_directory}, clear=True):
        _test_default_hdfs_service(mock_hadoop_home_directory, 'HADOOP_INSTALL')


def test_env_bad_hadoop_home_with_hadoop_install(mock_hadoop_home_directory):
    with mock.patch.dict(os.environ, {'HADOOP_HOME': '{}/no/where/here'.format(mock_hadoop_home_directory),
                                      'HADOOP_INSTALL': mock_hadoop_home_directory}, clear=True):
        with pytest.raises(IOError):
            # Trigger env var evaluation
            HdfsNamenodeResolver()


def test_unmatched_env_var(mock_hadoop_home_directory):
    with mock.patch.dict(os.environ, {'HADOOP_HOME_X': mock_hadoop_home_directory}, clear=True):
        suj = HdfsNamenodeResolver()
        # No successful connection
        with pytest.raises(RuntimeError):
            suj.resolve_default_hdfs_service()


def test_bad_hadoop_path(mock_hadoop_home_directory):
    with mock.patch.dict(os.environ, {'HADOOP_HOME': '{}/no/where/here'.format(mock_hadoop_home_directory)},
                         clear=True):
        with pytest.raises(IOError):
            HdfsNamenodeResolver()


def test_missing_or_empty_core_site(mock_hadoop_home_directory):
    with mock.patch.dict(os.environ, {'HADOOP_HOME': mock_hadoop_home_directory}):
        # Make core-site "disappear" and make sure we raise an error
        cur_path = '{}{}/core-site.xml'.format(mock_hadoop_home_directory, HC.HADOOP_CONFIG_PATH)
        new_path = '{}{}/core-site.xml.bak'.format(mock_hadoop_home_directory, HC.HADOOP_CONFIG_PATH)
        os.rename(cur_path, new_path)
        with pytest.raises(IOError):
            HdfsNamenodeResolver()
        # Make an empty file
        with open(cur_path, 'wt') as f:
            f.write('')
        # Re-trigger env var evaluation
        suj = HdfsNamenodeResolver()
        with pytest.raises(RuntimeError):
            suj.resolve_default_hdfs_service()
        # restore file for other tests to work
        os.rename(new_path, cur_path)


class HdfsMockError(Exception):
    pass


class MockHdfs(object):
    """
    Any operation in the mock class raises an exception for the first N failovers, and then returns
    True after those N calls.
    """

    def __init__(self, n_failovers=0, user=None):
        self._n_failovers = n_failovers
        self._user = user

    def __getattribute__(self, attr):
        """
        The Mock HDFS simply calls check_failover, regardless of the filesystem operator invoked.
        """

        def op(*args, **kwargs):
            """ Mock operator """
            return self._check_failovers()

        # Of course, exclude any protected/private method calls
        if not attr.startswith('_'):
            return op
        return object.__getattribute__(self, attr)

    def _check_failovers(self):
        if self._n_failovers == -1:
            # Special case to exercise the unhandled exception path
            raise HdfsMockError('Some random HDFS exception!')

        if self._n_failovers > 0:
            self._n_failovers -= 1
            raise ArrowIOError('org.apache.hadoop.ipc.RemoteException'
                               '(org.apache.hadoop.ipc.StandbyException): '
                               'Operation category READ is not supported in state standby. '
                               'Visit https://s.apache.org/sbnn-error\n'
                               '{} namenode failover(s) remaining!'.format(self._n_failovers))
        return True

    def __reduce__(self):
        raise AssertionError('A connection object can not be pickled. If we try to pickle it, it means '
                             'it leaks somehow with a closure that holds it and we need to make sure it '
                             'does not happen.')


class MockHdfsConnector(HdfsConnector):
    # static member for static hdfs_connect_namenode to access
    _n_failovers = 0
    _fail_n_next_connect = 0
    _connect_attempted: Dict[str, int] = {}

    @classmethod
    def reset(cls):
        cls._n_failovers = 0
        cls._fail_n_next_connect = 0
        cls._connect_attempted = {}

    @classmethod
    def set_n_failovers(cls, failovers):
        cls._n_failovers = failovers

    @classmethod
    def set_fail_n_next_connect(cls, fails):
        cls._fail_n_next_connect = fails

    @classmethod
    def connect_attempted(cls, host):
        if host in cls._connect_attempted:
            return cls._connect_attempted[host]
        else:
            return 0

    @classmethod
    def hdfs_connect_namenode(cls, url, driver='libhdfs3', user=None):
        netloc = '{}:{}'.format(url.hostname or 'default', url.port or 8020)
        if netloc not in cls._connect_attempted:
            cls._connect_attempted[netloc] = 0
        cls._connect_attempted[netloc] += 1
        # We just want to check connection attempt, but also raise an error if
        # 'default' or fail counter
        if cls._fail_n_next_connect != 0 or netloc == HC.DEFAULT_NN:
            if cls._fail_n_next_connect != 0:
                cls._fail_n_next_connect -= 1
            raise ArrowIOError('ERROR! Mock pyarrow hdfs connect to {} using driver {}, '
                               'fail counter: {}'
                               .format(netloc, driver, cls._fail_n_next_connect))
        # Return a mock HDFS object with optional failovers, so that this connector mock can
        # be shared for the HAHdfsClient failover tests below.
        hdfs = MockHdfs(cls._n_failovers, user=user)
        if cls._n_failovers > 0:
            cls._n_failovers -= 1
        return hdfs


class HdfsConnectorTest(unittest.TestCase):
    """Check correctness of connecting to a list of namenodes. """

    @classmethod
    def setUpClass(cls):
        """Initializes a mock HDFS namenode connector to track connection attempts."""
        cls.NAMENODES = [HC.WARP_TURTLE_NN1, HC.WARP_TURTLE_NN2]
        cls.suj = MockHdfsConnector()

    def setUp(self):
        self.suj.reset()

    def test_connect_to_either_namenode_ok(self):
        """ Test connecting OK to first of name node URLs. """
        self.assertIsNotNone(self.suj.connect_to_either_namenode(self.NAMENODES))
        self.assertEqual(0, self.suj.connect_attempted(HC.DEFAULT_NN))
        self.assertEqual(1, self.suj.connect_attempted(HC.WARP_TURTLE_NN1))
        self.assertEqual(0, self.suj.connect_attempted(HC.WARP_TURTLE_NN2))

    def test_connect_to_either_with_user(self):
        mock_name = "mock-manager"
        mocked_hdfs = self.suj.connect_to_either_namenode(self.NAMENODES, user=mock_name)
        self.assertEqual(mocked_hdfs._user, mock_name)

    def test_connect_to_either_namenode_ok_one_failed(self):
        """ With one failver, test that both namenode URLS are attempted, with 2nd connected. """
        self.suj.set_fail_n_next_connect(1)
        self.assertIsNotNone(self.suj.connect_to_either_namenode(self.NAMENODES))
        self.assertEqual(0, self.suj.connect_attempted(HC.DEFAULT_NN))
        self.assertEqual(1, self.suj.connect_attempted(HC.WARP_TURTLE_NN1))
        self.assertEqual(1, self.suj.connect_attempted(HC.WARP_TURTLE_NN2))

    def test_connect_to_either_namenode_exception_two_failed(self):
        """ With 2 failvers, test no connection, and no exception is raised. """
        self.suj.set_fail_n_next_connect(2)
        with self.assertRaises(HdfsConnectError):
            self.suj.connect_to_either_namenode(self.NAMENODES)
        self.assertEqual(0, self.suj.connect_attempted(HC.DEFAULT_NN))
        self.assertEqual(1, self.suj.connect_attempted(HC.WARP_TURTLE_NN1))
        self.assertEqual(1, self.suj.connect_attempted(HC.WARP_TURTLE_NN2))

    def test_connect_to_either_namenode_exception_four_failed(self):
        """ With 4 failvers, test that exception is raised. """
        self.suj.set_fail_n_next_connect(4)
        with self.assertRaises(HdfsConnectError):
            self.suj.connect_to_either_namenode(self.NAMENODES)
        with self.assertRaises(HdfsConnectError):
            self.suj.connect_to_either_namenode(self.NAMENODES)
        self.assertEqual(0, self.suj.connect_attempted(HC.DEFAULT_NN))
        self.assertEqual(2, self.suj.connect_attempted(HC.WARP_TURTLE_NN1))
        self.assertEqual(2, self.suj.connect_attempted(HC.WARP_TURTLE_NN2))


class HAHdfsClientTest(unittest.TestCase):
    """
    The HDFS testing functions are enumerated explicitly below for simplicity and clarity, but it
    should impose but a minute maintenance overhead, since MockHdfs class requires no enumeration.
    """

    @classmethod
    def setUpClass(cls):
        """Initializes namenodes list and mock HDFS namenode connector."""
        cls.NAMENODES = [HC.WARP_TURTLE_NN1, HC.WARP_TURTLE_NN2]

    def setUp(self):
        """Reset mock HDFS failover count."""
        MockHdfsConnector.reset()

    def test_unhandled_exception(self):
        """Exercise the unhandled exception execution path."""
        MockHdfsConnector.set_n_failovers(-1)
        with self.assertRaises(HdfsMockError) as e:
            getattr(HAHdfsClient(MockHdfsConnector, [HC.WARP_TURTLE_NN1]), 'ls')('random')
        self.assertTrue('random HDFS exception' in str(e.exception))

    def test_invalid_namenode_list(self):
        """Make sure robust to invalid namenode list."""
        MockHdfsConnector.set_n_failovers(-1)
        with self.assertRaises(HdfsConnectError) as e:
            getattr(HAHdfsClient(MockHdfsConnector, []), 'ls')('random')
        self.assertTrue('Unable to connect' in str(e.exception))
        with self.assertRaises(HdfsConnectError) as e:
            getattr(HAHdfsClient(MockHdfsConnector, [None]), 'ls')('random')
        self.assertTrue('Unable to connect' in str(e.exception))

    def test_client_pickles_correctly(self):
        """
        Does HAHdfsClient pickle properly?

        Check that all attributes are equal, with the exception of the HDFS object, which is fine
        as long as the types are the same.
        """
        mock_name = "mock-manager"
        client = HAHdfsClient(MockHdfsConnector, self.NAMENODES, user=mock_name)
        client_unpickled = pickle.loads(pickle.dumps(client))
        self.assertEqual(client._connector_cls, client_unpickled._connector_cls)
        self.assertEqual(client._list_of_namenodes, client_unpickled._list_of_namenodes)
        self.assertEqual(client._index_of_nn, client_unpickled._index_of_nn)
        self.assertEqual(client._user, client_unpickled._user)
        self.assertEqual(type(client._hdfs), type(client_unpickled._hdfs))

    def _try_failover_combos(self, func, *args, **kwargs):
        """Common tests for each of the known HDFS operators, with varying failover counts."""
        MockHdfsConnector.set_n_failovers(1)
        suj = HAHdfsClient(MockHdfsConnector, self.NAMENODES)
        self.assertTrue(getattr(suj, func)(*args, **kwargs))

        MockHdfsConnector.set_n_failovers(namenode_failover.MAX_FAILOVER_ATTEMPTS)
        suj = HAHdfsClient(MockHdfsConnector, self.NAMENODES)
        self.assertTrue(getattr(suj, func)(*args, **kwargs))

        MockHdfsConnector.set_n_failovers(namenode_failover.MAX_FAILOVER_ATTEMPTS + 1)
        suj = HAHdfsClient(MockHdfsConnector, self.NAMENODES)
        with self.assertRaises(MaxFailoversExceeded) as e:
            getattr(suj, func)(*args, **kwargs)
        self.assertEqual(len(e.exception.failed_exceptions),
                         namenode_failover.MAX_FAILOVER_ATTEMPTS + 1)
        self.assertEqual(e.exception.max_failover_attempts, namenode_failover.MAX_FAILOVER_ATTEMPTS)
        self.assertEqual(e.exception.__name__, func)
        self.assertTrue('Failover attempts exceeded' in str(e.exception))

    def test_cat(self):
        self._try_failover_combos('cat', 'random')

    def test_chmod(self):
        self._try_failover_combos('chmod', 'random', 0)

    def test_chown(self):
        self._try_failover_combos('chown', 'random', 'user')

    def test_delete(self):
        self._try_failover_combos('delete', 'random', recursive=True)

    def test_df(self):
        self._try_failover_combos('df')

    def test_disk_usage(self):
        self._try_failover_combos('disk_usage', 'random')

    def test_download(self):
        self._try_failover_combos('download', 'random', None)

    def test_exists(self):
        self._try_failover_combos('exists', 'random')

    def test_get_capacity(self):
        self._try_failover_combos('get_capacity')

    def test_get_space_used(self):
        self._try_failover_combos('get_space_used')

    def test_info(self):
        self._try_failover_combos('info', 'random')

    def test_ls(self):
        self._try_failover_combos('ls', 'random', detail=True)

    def test_mkdir(self):
        self._try_failover_combos('mkdir', 'random', create_parents=False)

    def test_open(self):
        self._try_failover_combos('open', 'random', 'rb')

    def test_rename(self):
        self._try_failover_combos('rename', 'random', 'new_random')

    def test_rm(self):
        self._try_failover_combos('rm', 'random', recursive=True)

    def test_upload(self):
        self._try_failover_combos('upload', 'random', None)


if __name__ == '__main__':
    unittest.main()
