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

from __future__ import unicode_literals

import logging
import os
import sys
import threading

from multiprocessing import Process, Pipe

import maxminddb

try:
    import maxminddb.extension
except ImportError:
    maxminddb.extension = None

from maxminddb import open_database, InvalidDatabaseError
from maxminddb.compat import FileNotFoundError
from maxminddb.const import (MODE_AUTO, MODE_MMAP_EXT, MODE_MMAP, MODE_FILE,
                             MODE_MEMORY)

if sys.version_info[:2] == (2, 6):
    import unittest2 as unittest
else:
    import unittest

if sys.version_info[0] == 2:
    unittest.TestCase.assertRaisesRegex = unittest.TestCase.assertRaisesRegexp
    unittest.TestCase.assertRegex = unittest.TestCase.assertRegexpMatches


class BaseTestReader(object):
    def test_reader(self):
        for record_size in [24, 28, 32]:
            for ip_version in [4, 6]:
                file_name = (
                    'tests/data/test-data/MaxMind-DB-test-ipv' +
                    str(ip_version) + '-' + str(record_size) + '.mmdb')
                reader = open_database(file_name, self.mode)

                self._check_metadata(reader, ip_version, record_size)

                if ip_version == 4:
                    self._check_ip_v4(reader, file_name)
                else:
                    self._check_ip_v6(reader, file_name)
                reader.close()

    def test_decoder(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        record = reader.get('::1.1.1.0')

        self.assertEqual(record['array'], [1, 2, 3])
        self.assertEqual(record['boolean'], True)
        self.assertEqual(record['bytes'], bytearray(b'\x00\x00\x00*'))
        self.assertEqual(record['double'], 42.123456)
        self.assertAlmostEqual(record['float'], 1.1)
        self.assertEqual(record['int32'], -268435456)
        self.assertEqual({
            'mapX': {
                'arrayX': [7, 8, 9],
                'utf8_stringX': 'hello'
            },
        }, record['map'])

        self.assertEqual(record['uint16'], 100)
        self.assertEqual(record['uint32'], 268435456)
        self.assertEqual(record['uint64'], 1152921504606846976)
        self.assertEqual(record['utf8_string'], 'unicode! ☯ - ♫')

        self.assertEqual(1329227995784915872903807060280344576,
                         record['uint128'])
        reader.close()

    def test_no_ipv4_search_tree(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-no-ipv4-search-tree.mmdb',
            self.mode)

        self.assertEqual(reader.get('1.1.1.1'), '::0/64')
        self.assertEqual(reader.get('192.1.1.1'), '::0/64')
        reader.close()

    def test_ipv6_address_in_ipv4_database(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-ipv4-24.mmdb', self.mode)
        with self.assertRaisesRegex(ValueError, 'Error looking up 2001::. '
                                    'You attempted to look up an IPv6 address '
                                    'in an IPv4-only database'):
            reader.get('2001::')
        reader.close()

    def test_broken_database(self):
        reader = open_database('tests/data/test-data/'
                               'GeoIP2-City-Test-Broken-Double-Format.mmdb',
                               self.mode)
        with self.assertRaisesRegex(InvalidDatabaseError,
                                    "The MaxMind DB file's data "
                                    "section contains bad data \(unknown data "
                                    "type or corrupt data\)"):
            reader.get('2001:220::')
        reader.close()

    def test_ip_validation(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        self.assertRaisesRegex(ValueError,
                               "'not_ip' does not appear to be an IPv4 or "
                               "IPv6 address", reader.get, ('not_ip'))
        reader.close()

    def test_missing_database(self):
        self.assertRaisesRegex(FileNotFoundError, "No such file or directory",
                               open_database, 'file-does-not-exist.mmdb',
                               self.mode)

    def test_nondatabase(self):
        self.assertRaisesRegex(InvalidDatabaseError,
                               'Error opening database file \(README.rst\). '
                               'Is this a valid MaxMind DB file\?',
                               open_database, 'README.rst', self.mode)

    def test_too_many_constructor_args(self):
        cls = self.readerClass[0]
        self.assertRaises(TypeError, cls, 'README.md', self.mode, 1)

    def test_bad_constructor_mode(self):
        cls = self.readerClass[0]
        self.assertRaisesRegex(
            ValueError,
            'Unsupported open mode \(100\)',
            cls,
            'README.md',
            mode=100)

    def test_no_constructor_args(self):
        cls = self.readerClass[0]
        self.assertRaisesRegex(
            TypeError,
            ' 1 required positional argument|\(pos 1\) not found|takes at least 2 arguments',
            cls)

    def test_too_many_get_args(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        self.assertRaises(TypeError, reader.get, ('1.1.1.1', 'blah'))
        reader.close()

    def test_no_get_args(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        self.assertRaises(TypeError, reader.get)
        reader.close()

    def test_metadata_args(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        self.assertRaises(TypeError, reader.metadata, ('blah'))
        reader.close()

    def test_metadata_unknown_attribute(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        metadata = reader.metadata()
        with self.assertRaisesRegex(AttributeError, "'Metadata' object has no "
                                    "attribute 'blah'"):
            metadata.blah
        reader.close()

    def test_close(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        reader.close()

    def test_double_close(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        reader.close()
        self.assertIsNone(reader.close(),
                          'Double close does not throw an exception')

    def test_closed_get(self):
        if self.mode == MODE_MEMORY:
            return
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        reader.close()
        self.assertRaisesRegex(ValueError,
                               'Attempt to read from a closed MaxMind DB.'
                               '|closed', reader.get, ('1.1.1.1'))

    # XXX - Figure out whether we want to have the same behavior on both the
    #       extension and the pure Python reader. If we do, the pure Python
    #       reader will need to throw an exception or the extension will need
    #       to keep the metadata in memory.
    def test_closed_metadata(self):
        reader = open_database(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
        reader.close()

        # The primary purpose of this is to ensure the extension doesn't
        # segfault
        try:
            metadata = reader.metadata()
        except IOError as ex:
            self.assertEqual('Attempt to read from a closed MaxMind DB.',
                             str(ex), 'extension throws exception')
        else:
            self.assertIsNotNone(metadata,
                                 'pure Python implementation returns value')

    def test_multiprocessing(self):
        self._check_concurrency(Process)

    def test_threading(self):
        self._check_concurrency(threading.Thread)

    if sys.version_info[0] == 2:

        def test_byte_ip_on_python2(self):
            reader = open_database(
                'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)
            record = reader.get(b'::1.1.1.0')

    def _check_concurrency(self, worker_class):
        reader = open_database('tests/data/test-data/GeoIP2-Domain-Test.mmdb',
                               self.mode)

        def lookup(pipe):
            try:
                for i in range(32):
                    reader.get('65.115.240.{i}'.format(i=i))
                pipe.send(1)
            except:
                pipe.send(0)
            finally:
                if worker_class is Process:
                    reader.close()
                pipe.close()

        pipes = [Pipe() for _ in range(32)]
        procs = [worker_class(target=lookup, args=(c, )) for (p, c) in pipes]
        for proc in procs:
            proc.start()
        for proc in procs:
            proc.join()

        reader.close()

        count = sum([p.recv() for (p, c) in pipes])

        self.assertEqual(count, 32, 'expected number of successful lookups')

    def _check_metadata(self, reader, ip_version, record_size):
        metadata = reader.metadata()

        self.assertEqual(2, metadata.binary_format_major_version,
                         'major version')
        self.assertEqual(metadata.binary_format_minor_version, 0)
        self.assertGreater(metadata.build_epoch, 1373571901)
        self.assertEqual(metadata.database_type, 'Test')

        self.assertEqual({
            'en': 'Test Database',
            'zh': 'Test Database Chinese'
        }, metadata.description)
        self.assertEqual(metadata.ip_version, ip_version)
        self.assertEqual(metadata.languages, ['en', 'zh'])
        self.assertGreater(metadata.node_count, 36)

        self.assertEqual(metadata.record_size, record_size)

    def _check_ip_v4(self, reader, file_name):
        for i in range(6):
            address = '1.1.1.' + str(pow(2, i))
            self.assertEqual({
                'ip': address
            },
                             reader.get(address),
                             'found expected data record for ' + address +
                             ' in ' + file_name)

        pairs = {
            '1.1.1.3': '1.1.1.2',
            '1.1.1.5': '1.1.1.4',
            '1.1.1.7': '1.1.1.4',
            '1.1.1.9': '1.1.1.8',
            '1.1.1.15': '1.1.1.8',
            '1.1.1.17': '1.1.1.16',
            '1.1.1.31': '1.1.1.16'
        }
        for key_address, value_address in pairs.items():
            data = {'ip': value_address}

            self.assertEqual(data,
                             reader.get(key_address),
                             'found expected data record for ' + key_address +
                             ' in ' + file_name)

        for ip in ['1.1.1.33', '255.254.253.123']:
            self.assertIsNone(reader.get(ip))

    def _check_ip_v6(self, reader, file_name):
        subnets = [
            '::1:ffff:ffff', '::2:0:0', '::2:0:40', '::2:0:50', '::2:0:58'
        ]

        for address in subnets:
            self.assertEqual({
                'ip': address
            },
                             reader.get(address),
                             'found expected data record for ' + address +
                             ' in ' + file_name)

        pairs = {
            '::2:0:1': '::2:0:0',
            '::2:0:33': '::2:0:0',
            '::2:0:39': '::2:0:0',
            '::2:0:41': '::2:0:40',
            '::2:0:49': '::2:0:40',
            '::2:0:52': '::2:0:50',
            '::2:0:57': '::2:0:50',
            '::2:0:59': '::2:0:58'
        }

        for key_address, value_address in pairs.items():
            self.assertEqual({
                'ip': value_address
            },
                             reader.get(key_address),
                             'found expected data record for ' + key_address +
                             ' in ' + file_name)

        for ip in ['1.1.1.33', '255.254.253.123', '89fa::']:
            self.assertIsNone(reader.get(ip))


def has_maxminddb_extension():
    return maxminddb.extension and hasattr(maxminddb.extension, 'Reader')


@unittest.skipIf(not has_maxminddb_extension() and
                 not os.environ.get('MM_FORCE_EXT_TESTS'),
                 'No C extension module found. Skipping tests')
class TestExtensionReader(BaseTestReader, unittest.TestCase):
    mode = MODE_MMAP_EXT

    if has_maxminddb_extension():
        readerClass = [maxminddb.extension.Reader]


class TestAutoReader(BaseTestReader, unittest.TestCase):
    mode = MODE_AUTO

    if has_maxminddb_extension():
        readerClass = [maxminddb.extension.Reader]
    else:
        readerClass = [maxminddb.reader.Reader]


class TestMMAPReader(BaseTestReader, unittest.TestCase):
    mode = MODE_MMAP
    readerClass = [maxminddb.reader.Reader]


class TestFileReader(BaseTestReader, unittest.TestCase):
    mode = MODE_FILE
    readerClass = [maxminddb.reader.Reader]


class TestMemoryReader(BaseTestReader, unittest.TestCase):
    mode = MODE_MEMORY
    readerClass = [maxminddb.reader.Reader]


class TestOldReader(unittest.TestCase):
    def test_old_reader(self):
        reader = maxminddb.Reader(
            'tests/data/test-data/MaxMind-DB-test-decoder.mmdb')
        record = reader.get('::1.1.1.0')

        self.assertEqual(record['array'], [1, 2, 3])
        reader.close()