#!/usr/bin/env python3

# 
#   Copyright 2016 RIFT.IO 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 argparse
import asyncio
import contextlib
import io
import logging
import os
import sys
import tempfile
import time
import unittest
import uuid
import xmlrunner

from rift.mano import cloud
from rift.tasklets.rwimagemgr import upload
from rift.package import checksums
from rift.test.dts import async_test
import rw_status

import gi
gi.require_version('RwCal', '1.0')
gi.require_version('RwcalYang', '1.0')
gi.require_version('RwCloudYang', '1.0')
gi.require_version('RwLog', '1.0')
gi.require_version('RwTypes', '1.0')
from gi.repository import (
        RwCal,
        RwCloudYang,
        RwLog,
        RwTypes,
        RwcalYang,
        )

rwstatus = rw_status.rwstatus_from_exc_map({
    IndexError: RwTypes.RwStatus.NOTFOUND,
    KeyError: RwTypes.RwStatus.NOTFOUND,
    })


class CreateImageMock(object):
    def __init__(self, log):
        self._log = log
        self.image_name = None
        self.image_checksum = None

        self.do_fail = False
        self.do_read_slow = False

        self._image_msgs = []

    def add_existing_image(self, image_msg):
        self._log.debug("Appending existing image msg: %s", image_msg)
        self._image_msgs.append(image_msg)

    @rwstatus
    def do_create_image(self, _, image):
        if self.do_fail:
            self._log.debug("Simulating failure")
            raise ValueError("FAILED")

        if not image.has_field("fileno"):
            raise ValueError("Image must have fileno")

        self.image_name = image.name

        # Create a duplicate file descriptor to allow this code to own
        # its own descritor (and thus close it) and allow the client to
        # own and close its own descriptor.
        new_fileno = os.dup(image.fileno)
        with os.fdopen(new_fileno, 'rb') as hdl:
            bytes_hdl = io.BytesIO()
            if self.do_read_slow:
                self._log.debug("slow reading from mock cal")
                try:
                    num_bytes = 0
                    while True:
                        d = os.read(new_fileno, 1024)
                        num_bytes += len(d)
                        bytes_hdl.write(d)
                        if not d:
                            self._log.debug("read %s bytes", num_bytes)
                            return

                        time.sleep(.05)

                except Exception as e:
                    self._log.warning("caught exception when reading: %s",
                                      str(e))
                    raise

            else:
                bytes_hdl.write(hdl.read())

            bytes_hdl.seek(0)
            self.image_checksum = checksums.checksum(bytes_hdl)
            bytes_hdl.close()

        return str(uuid.uuid4())

    @rwstatus
    def do_get_image_list(self, account):
        boxed_image_list = RwcalYang.VimResources()
        for msg in self._image_msgs:
            boxed_image_list.imageinfo_list.append(msg)

        return boxed_image_list


def create_random_image_file():
    with open("/dev/urandom", "rb") as rand_hdl:
        file_hdl = tempfile.NamedTemporaryFile("r+b")
        file_hdl.write(rand_hdl.read(1 * 1024 * 1024))
        file_hdl.flush()
        file_hdl.seek(0)
        return file_hdl


def get_file_hdl_gen(file_hdl):
    while True:
        try:
            d = file_hdl.read(64)
        except ValueError:
            return

        if not d:
            return

        yield d


def get_image_checksum(image_hdl):
    image_checksum = checksums.checksum(image_hdl)
    image_hdl.seek(0)
    return image_checksum


def create_image_info(image_name, image_checksum):
    image = RwcalYang.ImageInfoItem()
    image.name = image_name
    image.checksum = image_checksum
    image.disk_format = os.path.splitext(image_name)[1][1:]
    image.container_format = "bare"

    return image


class UploadTaskMixin(object):
    def __init__(self, log, loop):
        self._log = log
        self._loop = loop

    def create_image_hdl(self):
        image_hdl = create_random_image_file()

        return image_hdl

    @contextlib.contextmanager
    def create_upload_task(self, account, image_name="test.qcow2",
                           image_checksum=None, image_info=None):

        with self.create_image_hdl() as image_hdl:

            image_checksum = get_image_checksum(image_hdl) \
                if image_checksum is None else image_checksum

            image_info = create_image_info(image_name, image_checksum) \
                if image_info is None else image_info

            iter_hdl = get_file_hdl_gen(image_hdl)
            pipe_gen = upload.GlanceImagePipeGen(self._log, self._loop, iter_hdl)

            upload_task = upload.AccountImageUploadTask(
                    self._log, self._loop, account, image_info, pipe_gen.read_hdl,
                    write_canceller=pipe_gen
                    )
            pipe_gen.start()

            yield upload_task


class ImageMockMixin(object):
    ACCOUNT_MSG = RwCloudYang.CloudAccount(
        name="mock",
        account_type="mock",
        )

    def __init__(self, log):
        self._log = log
        self._account = cloud.CloudAccount(
                self._log,
                RwLog.Ctx.new(__file__), ImageMockMixin.ACCOUNT_MSG
                )

        self._create_image_mock = CreateImageMock(self._log)

        # Mock the create_image call
        self._account.cal.create_image = self._create_image_mock.do_create_image
        self._account.cal.get_image_list = self._create_image_mock.do_get_image_list

    @property
    def account(self):
        return self._account

    @property
    def image_mock(self):
        return self._create_image_mock


class TestImageUploadTask(unittest.TestCase, UploadTaskMixin, ImageMockMixin):
    def __init__(self, *args, **kwargs):
        self._loop = asyncio.get_event_loop()
        self._log = logging.getLogger(__file__)

        ImageMockMixin.__init__(self, self._log)
        UploadTaskMixin.__init__(self, self._log, self._loop)
        unittest.TestCase.__init__(self, *args, **kwargs)

    @async_test
    def test_upload_image_task(self):
        with self.create_upload_task(self.account) as upload_task:
            yield from upload_task.start()

        self.assertIn("QUEUED", upload_task.state_stack)
        self.assertIn("CHECK_IMAGE_EXISTS", upload_task.state_stack)
        self.assertIn("UPLOADING", upload_task.state_stack)
        self.assertIn("COMPLETED", upload_task.state_stack)

        self.assertEqual("COMPLETED", upload_task.state)

        self.assertEqual(self.image_mock.image_name, upload_task.image_name)
        self.assertEqual(self.image_mock.image_checksum, upload_task.image_checksum)

        task_pb_msg = upload_task.pb_msg
        self.assertEqual(upload_task.image_name, task_pb_msg.image_name)

    @async_test
    def test_cancel_image_task(self):
        @asyncio.coroutine
        def wait_for_task_state(upload_task, state, timeout=10):
            start_time = time.time()
            while (time.time() - start_time) < timeout:
                if upload_task.state == state:
                    return

                yield from asyncio.sleep(.01)

            raise asyncio.TimeoutError()

        self.image_mock.do_read_slow = True

        with self.create_upload_task(self.account) as upload_task:
            upload_task.start()
            yield from wait_for_task_state(upload_task, "UPLOADING")
            upload_task.stop()
            self.assertEqual("CANCELLING", upload_task.state)
            yield from wait_for_task_state(upload_task, "CANCELLED")

    @async_test
    def test_create_image_failed(self):
        self.image_mock.do_fail = True

        with self.create_upload_task(self.account) as upload_task:
            yield from upload_task.start()

        self.assertEqual("FAILED", upload_task.state)

    @async_test
    def test_create_image_name_and_checksum_exists(self):
        with self.create_upload_task(self.account) as upload_task:
            image_entry = RwcalYang.ImageInfoItem(
                    id="asdf",
                    name=upload_task.image_name,
                    checksum=upload_task.image_checksum
                    )
            self.image_mock.add_existing_image(image_entry)

            yield from upload_task.start()

        # No image should have been uploaded, since the name and checksum
        self.assertEqual(self.image_mock.image_checksum, None)

        self.assertEqual("COMPLETED", upload_task.state)
        self.assertTrue("UPLOADING" not in upload_task.state_stack)


class TestUploadJob(unittest.TestCase, UploadTaskMixin, ImageMockMixin):
    def __init__(self, *args, **kwargs):
        self._loop = asyncio.get_event_loop()
        self._log = logging.getLogger(__file__)

        ImageMockMixin.__init__(self, self._log)
        UploadTaskMixin.__init__(self, self._log, self._loop)
        unittest.TestCase.__init__(self, *args, **kwargs)

    @async_test
    def test_single_task_upload_job(self):
        with self.create_upload_task(self.account) as upload_task:
            job = upload.ImageUploadJob(self._log, self._loop, [upload_task])
            self.assertEqual("QUEUED", job.state)
            yield from job.start()

        self.assertIn("QUEUED", job.state_stack)
        self.assertIn("IN_PROGRESS", job.state_stack)
        self.assertIn("COMPLETED", job.state_stack)

        self.assertEqual("COMPLETED", job.state)

        job_pb_msg = job.pb_msg
        self.assertEqual("COMPLETED", job_pb_msg.status)

    @async_test
    def test_multiple_tasks_upload_job(self):
        with self.create_upload_task(self.account) as upload_task1:
            with self.create_upload_task(self.account) as upload_task2:
                job = upload.ImageUploadJob(
                        self._log, self._loop, [upload_task1, upload_task2])
                yield from job.start()

        self.assertEqual("COMPLETED", job.state)

    @async_test
    def test_failed_task_in_job(self):
        self.image_mock.do_fail = True

        with self.create_upload_task(self.account) as upload_task:
            job = upload.ImageUploadJob(
                    self._log, self._loop, [upload_task])
            yield from job.start()

        self.assertEqual("FAILED", job.state)

    @async_test
    def test_cancel_job(self):
        @asyncio.coroutine
        def wait_for_job_state(upload_job, state, timeout=10):
            start_time = time.time()
            while (time.time() - start_time) < timeout:
                if upload_job.state == state:
                    return

                yield from asyncio.sleep(.01)

            raise asyncio.TimeoutError()

        self.image_mock.do_read_slow = True

        with self.create_upload_task(self.account) as upload_task:
            job = upload.ImageUploadJob(
                    self._log, self._loop, [upload_task])
            job.start()
            yield from wait_for_job_state(job, "IN_PROGRESS")
            job.stop()
            self.assertEqual("CANCELLING", job.state)
            yield from wait_for_job_state(job, "CANCELLED")

        self.assertEqual("CANCELLED", job.state)


class TestUploadJobController(unittest.TestCase, UploadTaskMixin, ImageMockMixin):
    def __init__(self, *args, **kwargs):
        self._loop = asyncio.get_event_loop()
        self._log = logging.getLogger(__file__)

        ImageMockMixin.__init__(self, self._log)
        unittest.TestCase.__init__(self, *args, **kwargs)

    @async_test
    def test_controller_single_task_job(self):
        controller = upload.ImageUploadJobController(
                self._log, self._loop
                )

        with self.create_upload_task(self.account) as upload_task:
            job_id = controller.create_job([upload_task])
            self.assertEqual(len(controller.active_jobs), 1)
            self.assertEqual(len(controller.completed_jobs), 0)

            job = controller.get_job(job_id)
            yield from job.wait()

            self.assertEqual(len(controller.active_jobs), 0)
            self.assertEqual(len(controller.completed_jobs), 1)

            upload_jobs_pb_msg = controller.pb_msg
            self.assertEqual(len(upload_jobs_pb_msg.job), 1)

    @async_test
    def test_controller_multi_task_job(self):
        controller = upload.ImageUploadJobController(
                self._log, self._loop
                )

        with self.create_upload_task(self.account) as upload_task1:
            with self.create_upload_task(self.account) as upload_task2:
                job_id = controller.create_job([upload_task1, upload_task2])
                self.assertEqual(len(controller.active_jobs), 1)
                self.assertEqual(len(controller.completed_jobs), 0)

                job = controller.get_job(job_id)
                yield from job.wait()
                self.assertEqual(len(controller.active_jobs), 0)
                self.assertEqual(len(controller.completed_jobs), 1)

    @async_test
    def test_controller_multi_jobs(self):
        controller = upload.ImageUploadJobController(
                self._log, self._loop
                )

        with self.create_upload_task(self.account) as upload_task1:
            with self.create_upload_task(self.account) as upload_task2:
                job1_id = controller.create_job([upload_task1])
                job2_id = controller.create_job([upload_task2])
                self.assertEqual(len(controller.active_jobs), 2)
                self.assertEqual(len(controller.completed_jobs), 0)

                job1 = controller.get_job(job1_id)
                job2 = controller.get_job(job2_id)

                yield from asyncio.wait(
                        [job1.wait(), job2.wait()],
                        loop=self._loop
                        )

                self.assertEqual(len(controller.active_jobs), 0)
                self.assertEqual(len(controller.completed_jobs), 2)


class TestRateCalc(unittest.TestCase):
    def test_no_smoothing(self):
        calc = upload.ByteRateCalculator(1)
        self.assertEqual(0, calc.rate)
        calc.add_measurement(100, 1)
        self.assertEqual(100, calc.rate)
        calc.add_measurement(400, 2)
        self.assertEqual(200, calc.rate)

    def test_smoothing(self):
        calc = upload.ByteRateCalculator(2)
        calc.add_measurement(100, 1)
        self.assertEqual(100, calc.rate)

        calc.add_measurement(400, 2)
        self.assertEqual(150, calc.rate)

        calc.add_measurement(400, 2)
        self.assertEqual(175, calc.rate)


class TestUploadProgress(unittest.TestCase):
    def setUp(self):
        self._loop = asyncio.get_event_loop()
        self._log = logging.getLogger(__file__)

    def test_write_proxy(self):
        mem_hdl = io.BytesIO()
        proxy = upload.UploadProgressWriteProxy(self._log, self._loop, 1000, mem_hdl)

        data = b'test_bytes'

        proxy.write(data)
        self.assertEqual(data, mem_hdl.getvalue())
        self.assertEqual(len(data), proxy.bytes_written)
        self.assertEqual(1000, proxy.bytes_total)
        self.assertEqual(1, proxy.progress_percent)

        proxy.close()
        self.assertTrue(mem_hdl.closed)


def main(argv=sys.argv[1:]):
    logging.basicConfig(format='TEST %(message)s')

    runner = xmlrunner.XMLTestRunner(output=os.environ["RIFT_MODULE_TEST"])
    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--verbose', action='store_true')
    parser.add_argument('-n', '--no-runner', action='store_true')

    args, unknown = parser.parse_known_args(argv)
    if args.no_runner:
        runner = None

    # Set the global logging level
    logging.getLogger().setLevel(logging.DEBUG if args.verbose else logging.ERROR)

    # The unittest framework requires a program name, so use the name of this
    # file instead (we do not want to have to pass a fake program name to main
    # when this is called from the interpreter).
    unittest.main(argv=[__file__] + unknown + ["-v"], testRunner=runner)

if __name__ == '__main__':
    main()
