'''
Copyright (C) 2015 Dato, Inc.
All rights reserved.

This software may be modified and distributed under the terms
of the BSD license. See the DATO-PYTHON-LICENSE file for details.
'''
import commands
import json
import logging
import os
import re
import tempfile
import unittest
import pandas

import graphlab
import graphlab.connect.main as glconnect
import graphlab.sys_util as _sys_util
from graphlab.test.util import create_server, start_test_tcp_server
from pandas.util.testing import assert_frame_equal


def _test_save_load_object_helper(testcase, obj, url):
    """
    Helper function to test save and load a server side object to a given url.
    """
    def cleanup(url):
        """
        Remove the saved file from temp directory.
        """
        protocol = None
        path = None
        splits = url.split("://")
        if len(splits) > 1:
            protocol = splits[0]
            path = splits[1]
        else:
            path = url
        if not protocol or protocol is "local" or protocol is "remote":
            tempdir = tempfile.gettempdir()
            pattern = path + ".*"
            for f in os.listdir(tempdir):
                if re.search(pattern, f):
                    os.remove(os.path.join(tempdir, f))

    if isinstance(obj, graphlab.SGraph):
        obj.save(url + ".graph")
        newobj = graphlab.load_graph(url + ".graph")
        testcase.assertItemsEqual(obj.get_fields(), newobj.get_fields())
        testcase.assertDictEqual(obj.summary(), newobj.summary())
    elif isinstance(obj, graphlab.Model):
        obj.save(url + ".model")
        newobj = graphlab.load_model(url + ".model")
        testcase.assertItemsEqual(obj.list_fields(), newobj.list_fields())
        testcase.assertEqual(type(obj), type(newobj))
    elif isinstance(obj, graphlab.SFrame):
        obj.save(url + ".frame_idx")
        newobj = graphlab.load_sframe(url + ".frame_idx")
        testcase.assertEqual(obj.shape, newobj.shape)
        testcase.assertEqual(obj.column_names(), newobj.column_names())
        testcase.assertEqual(obj.column_types(), newobj.column_types())
        assert_frame_equal(obj.head(obj.num_rows()).to_dataframe(),
                           newobj.head(newobj.num_rows()).to_dataframe())
    else:
        raise TypeError
    cleanup(url)


def create_test_objects():
    vertices = pandas.DataFrame({'vid': ['1', '2', '3'],
                                 'color': ['g', 'r', 'b'],
                                 'vec': [[.1, .1, .1], [.1, .1, .1], [.1, .1, .1]]})
    edges = pandas.DataFrame({'src_id': ['1', '2', '3'],
                              'dst_id': ['2', '3', '4'],
                              'weight': [0., 0.1, 1.]})

    graph = graphlab.SGraph().add_vertices(vertices, 'vid').add_edges(edges, 'src_id', 'dst_id')
    sframe = graphlab.SFrame(edges)
    model = graphlab.pagerank.create(graph)
    return (graph, sframe, model)


class LocalFSConnectorTests(unittest.TestCase):

    @classmethod
    def setUpClass(self):
        self.tempfile = tempfile.NamedTemporaryFile().name
        (self.graph, self.sframe, self.model) = create_test_objects()

    def _test_read_write_helper(self, url, content):
        url = graphlab.util._make_internal_url(url)
        glconnect.get_unity().__write__(url, content)
        content_read = glconnect.get_unity().__read__(url)
        self.assertEquals(content_read, content)
        if os.path.exists(url):
            os.remove(url)

    def test_object_save_load(self):
        for prefix in ['', 'local://', 'remote://']:
            _test_save_load_object_helper(self, self.graph, prefix + self.tempfile)
            _test_save_load_object_helper(self, self.model, prefix + self.tempfile)
            _test_save_load_object_helper(self, self.sframe, prefix + self.tempfile)

    def test_basic(self):
        self._test_read_write_helper(self.tempfile, 'hello world')
        self._test_read_write_helper("local://" + self.tempfile + ".csv", 'hello,world,woof')
        self._test_read_write_helper("remote://" + self.tempfile + ".csv", 'hello,world,woof')

    def test_gzip(self):
        self._test_read_write_helper(self.tempfile + ".gz", 'hello world')
        self._test_read_write_helper(self.tempfile + ".csv.gz", 'hello world')
        self._test_read_write_helper("local://" + self.tempfile + ".csv.gz", 'hello world')
        self._test_read_write_helper("remote://" + self.tempfile + ".csv.gz", 'hello world')

    def test_exception(self):
        self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("/root/tmp"))
        self.assertRaises(IOError, lambda: glconnect.get_unity().__write__("/root/tmp", '.....'))
        self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("/root/tmp"))
        self.assertRaises(IOError, lambda: glconnect.get_unity().__write__("/root/tmp", '.....'))
        self.assertRaises(IOError, lambda: self.graph.save("/root/tmp.graph"))
        self.assertRaises(IOError, lambda: self.sframe.save("/root/tmp.frame_idx"))
        self.assertRaises(IOError, lambda: self.model.save("/root/tmp.model"))
        self.assertRaises(IOError, lambda: graphlab.load_graph("/root/tmp.graph"))
        self.assertRaises(IOError, lambda: graphlab.load_sframe("/root/tmp.frame_idx"))
        self.assertRaises(IOError, lambda: graphlab.load_model("/root/tmp.model"))


class RemoteFSConnectorTests(unittest.TestCase):

    @classmethod
    def setUpClass(self):
        glconnect.stop()
        auth_token = 'graphlab_awesome'
        self.server = start_test_tcp_server(auth_token=auth_token)
        glconnect.launch(self.server.get_server_addr(), auth_token=auth_token)
        self.tempfile = tempfile.NamedTemporaryFile().name
        (self.graph, self.sframe, self.model) = create_test_objects()

    @classmethod
    def tearDownClass(self):
        glconnect.stop()
        self.server.stop()

    def _test_read_write_helper(self, url, content):
        url = graphlab.util._make_internal_url(url)
        glconnect.get_unity().__write__(url, content)
        content_read = glconnect.get_unity().__read__(url)
        self.assertEquals(content_read, content)

    def test_basic(self):
        self._test_read_write_helper("remote://" + self.tempfile, 'hello,world,woof')

    def test_gzip(self):
        self._test_read_write_helper("remote://" + self.tempfile + ".csv.gz", 'hello,world,woof')

    def test_object_save_load(self):
        prefix = "remote://"
        _test_save_load_object_helper(self, self.graph, prefix + self.tempfile)
        _test_save_load_object_helper(self, self.model, prefix + self.tempfile)
        _test_save_load_object_helper(self, self.sframe, prefix + self.tempfile)

    def test_exception(self):
        self.assertRaises(ValueError, lambda: self._test_read_write_helper(self.tempfile, 'hello world'))
        self.assertRaises(ValueError, lambda: self._test_read_write_helper("local://" + self.tempfile + ".csv.gz", 'hello,world,woof'))
        self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("remote:///root/tmp"))
        self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("remote:///root/tmp"))
        self.assertRaises(IOError, lambda: glconnect.get_unity().__write__("remote:///root/tmp", '.....'))
        self.assertRaises(IOError, lambda: self.graph.save("remote:///root/tmp.graph"))
        self.assertRaises(IOError, lambda: self.sframe.save("remote:///root/tmp.frame_idx"))
        self.assertRaises(IOError, lambda: self.model.save("remote:///root/tmp.model"))
        self.assertRaises(IOError, lambda: graphlab.load_graph("remote:///root/tmp.graph"))
        self.assertRaises(IOError, lambda: graphlab.load_sframe("remote:///root/tmp.frame_idx"))
        self.assertRaises(IOError, lambda: graphlab.load_model("remote:///root/tmp.model"))


class HttpConnectorTests(unittest.TestCase):

    @classmethod
    def setUpClass(self):
        self.url = "http://s3-us-west-2.amazonaws.com/testdatasets/a_to_z.txt.gz"

    def _test_read_helper(self, url, content_expected):
        url = graphlab.util._make_internal_url(url)
        content_read = glconnect.get_unity().__read__(url)
        self.assertEquals(content_read, content_expected)

    def test_read(self):
        expected = "\n".join([str(unichr(i + ord('a'))) for i in range(26)])
        expected = expected + "\n"
        self._test_read_helper(self.url, expected)

    def test_exception(self):
        self.assertRaises(IOError, lambda: glconnect.get_unity().__write__(self.url, '.....'))

@unittest.skip("Disabling HDFS Connector Tests")
class HDFSConnectorTests(unittest.TestCase):
    # This test requires hadoop to be installed and avaiable in $PATH.
    # If not, the tests will be skipped.
    @classmethod
    def setUpClass(self):
        self.has_hdfs = len(_sys_util.get_hadoop_class_path()) > 0
        self.tempfile = tempfile.NamedTemporaryFile().name
        (self.graph, self.sframe, self.model) = create_test_objects()

    def _test_read_write_helper(self, url, content_expected):
        url = graphlab.util._make_internal_url(url)
        glconnect.get_unity().__write__(url, content_expected)
        content_read = glconnect.get_unity().__read__(url)
        self.assertEquals(content_read, content_expected)
        # clean up the file we wrote
        status, output = commands.getstatusoutput('hadoop fs -test -e ' + url)
        if status is 0:
            commands.getstatusoutput('hadoop fs -rm ' + url)

    def test_basic(self):
        if self.has_hdfs:
            self._test_read_write_helper("hdfs://" + self.tempfile, 'hello,world,woof')
        else:
            logging.getLogger(__name__).info("No hdfs avaiable. Test pass.")

    def test_gzip(self):
        if self.has_hdfs:
            self._test_read_write_helper("hdfs://" + self.tempfile + ".gz", 'hello,world,woof')
            self._test_read_write_helper("hdfs://" + self.tempfile + ".csv.gz", 'hello,world,woof')
        else:
            logging.getLogger(__name__).info("No hdfs avaiable. Test pass.")

    def test_object_save_load(self):
        if self.has_hdfs:
            prefix = "hdfs://"
            _test_save_load_object_helper(self, self.graph, prefix + self.tempfile)
            _test_save_load_object_helper(self, self.model, prefix + self.tempfile)
            _test_save_load_object_helper(self, self.sframe, prefix + self.tempfile)
        else:
            logging.getLogger(__name__).info("No hdfs avaiable. Test pass.")

    def test_exception(self):
        bad_url = "hdfs:///root/"
        if self.has_hdfs:
            self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("hdfs:///"))
            self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("hdfs:///tmp"))
            self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("hdfs://" + self.tempfile))
            self.assertRaises(IOError, lambda: glconnect.get_unity().__write__(bad_url + "/tmp", "somerandomcontent"))
            self.assertRaises(IOError, lambda: self.graph.save(bad_url + "x.graph"))
            self.assertRaises(IOError, lambda: self.sframe.save(bad_url + "x.frame_idx"))
            self.assertRaises(IOError, lambda: self.model.save(bad_url + "x.model"))
            self.assertRaises(IOError, lambda: graphlab.load_graph(bad_url + "mygraph"))
            self.assertRaises(IOError, lambda: graphlab.load_sframe(bad_url + "x.frame_idx"))
            self.assertRaises(IOError, lambda: graphlab.load_model(bad_url + "x.model"))
        else:
            logging.getLogger(__name__).info("No hdfs avaiable. Test pass.")


@unittest.skip("Disabling S3 Connector Tests")
class S3ConnectorTests(unittest.TestCase):
    # This test requires aws cli to be installed. If not, the tests will be skipped.
    @classmethod
    def setUpClass(self):
        status, output = commands.getstatusoutput('aws s3api list-buckets')
        self.has_s3 = (status is 0)
        self.standard_bucket = None
        self.regional_bucket = None
        # Use aws cli s3api to find a bucket with "gl-testdata" in the name, and use it as out test bucket.
        # Temp files will be read from /written to the test bucket's /tmp folder and be cleared on exist.
        if self.has_s3:
            try:
                json_output = json.loads(output)
                bucket_list = [b['Name'] for b in json_output['Buckets']]
                assert 'gl-testdata' in bucket_list
                assert 'gl-testdata-oregon' in bucket_list
                self.standard_bucket = 'gl-testdata'
                self.regional_bucket = 'gl-testdata-oregon'
                self.tempfile = tempfile.NamedTemporaryFile().name
                (self.graph, self.sframe, self.model) = create_test_objects()
            except:
                logging.getLogger(__name__).warning("Fail parsing ioutput of s3api into json. Please check your awscli version.")
                self.has_s3 = False

    def _test_read_write_helper(self, url, content_expected):
        s3url = graphlab.util._make_internal_url(url)
        glconnect.get_unity().__write__(s3url, content_expected)
        content_read = glconnect.get_unity().__read__(s3url)
        self.assertEquals(content_read, content_expected)
        (status, output) = commands.getstatusoutput('aws s3 rm --region us-west-2 ' + url)
        if status is not 0:
            logging.getLogger(__name__).warning("Cannot remove file: " + url)

    def test_basic(self):
        if self.has_s3:
            for bucket in [self.standard_bucket, self.regional_bucket]:
                self._test_read_write_helper("s3://" + bucket + self.tempfile, 'hello,world,woof')
        else:
            logging.getLogger(__name__).info("No s3 bucket avaiable. Test pass.")

    def test_gzip(self):
        if self.has_s3:
            self._test_read_write_helper("s3://" + self.standard_bucket + self.tempfile + ".gz", 'hello,world,woof')
        else:
            logging.getLogger(__name__).info("No s3 bucket avaiable. Test pass.")

    def test_object_save_load(self):
        if self.has_s3:
            prefix = "s3://" + self.standard_bucket
            _test_save_load_object_helper(self, self.graph, prefix + self.tempfile)
            _test_save_load_object_helper(self, self.model, prefix + self.tempfile)
            _test_save_load_object_helper(self, self.sframe, prefix + self.tempfile)
        else:
            logging.getLogger(__name__).info("No s3 bucket avaiable. Test pass.")

    def test_exception(self):
        if self.has_s3:
            bad_bucket = "i_am_a_bad_bucket"
            prefix = "s3://" + bad_bucket
            self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("s3:///"))
            self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("s3://" + self.standard_bucket + "/somerandomfile"))
            self.assertRaises(IOError, lambda: glconnect.get_unity().__read__("s3://" + "/somerandomfile"))
            self.assertRaises(IOError, lambda: glconnect.get_unity().__write__("s3://" + "/somerandomfile", "somerandomcontent"))
            self.assertRaises(IOError, lambda: glconnect.get_unity().__write__("s3://" + self.standard_bucket + "I'amABadUrl/", "somerandomcontent"))
            self.assertRaises(IOError, lambda: self.graph.save(prefix + "/x.graph"))
            self.assertRaises(IOError, lambda: self.sframe.save(prefix + "/x.frame_idx"))
            self.assertRaises(IOError, lambda: self.model.save(prefix + "/x.model"))
            self.assertRaises(IOError, lambda: graphlab.load_graph(prefix + "/x.graph"))
            self.assertRaises(IOError, lambda: graphlab.load_sframe(prefix + "/x.frame_idx"))
            self.assertRaises(IOError, lambda: graphlab.load_model(prefix + "/x.model"))
        else:
            logging.getLogger(__name__).info("No s3 bucket avaiable. Test pass.")
