##
# Copyright (c) 2010-2017 Apple Inc. All rights reserved.
#
# 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.
##

from twext.enterprise.dal.syntax import Delete, Insert, Select, Count

from twisted.internet.defer import inlineCallbacks, succeed, returnValue
from twisted.trial.unittest import TestCase

from twistedcaldav.config import config

from txdav.caldav.datastore.sql import CalendarStoreFeatures
from txdav.caldav.datastore.test.util import CommonStoreTests
from txdav.common.datastore.sql_tables import schema
from txdav.common.datastore.test.util import theStoreBuilder, \
    StubNotifierFactory
from txdav.common.datastore.upgrade.sql.others import attachment_migration
from txdav.common.datastore.upgrade.sql.upgrade import UpgradeDatabaseOtherStep

import hashlib
import os

"""
Tests for L{txdav.common.datastore.upgrade.sql.upgrade}.
"""


class AttachmentMigrationModeTests(CommonStoreTests):
    """
    Tests for L{UpgradeDatabaseSchemaStep}.
    """

    @inlineCallbacks
    def setUp(self):
        # We need to skip the immediate base class since we are creating our own
        # store in each test
        yield TestCase.setUp(self)

    @inlineCallbacks
    def _initStore(self, enableManagedAttachments=True):
        """
        Build a store with certain bits cleaned out.
        """

        self.patch(config, "EnableManagedAttachments", enableManagedAttachments)

        store = yield theStoreBuilder.buildStore(
            self, {"push": StubNotifierFactory()}
        )
        store.enableManagedAttachments = enableManagedAttachments

        txn = store.newTransaction()
        cs = schema.CALENDARSERVER
        yield Delete(
            From=cs,
            Where=cs.NAME == "MANAGED-ATTACHMENTS"
        ).on(txn)
        yield txn.commit()

        returnValue(store)

    @inlineCallbacks
    def test_upgradeFromEmptyDropbox(self):
        """
        Test L{attachment_migration.doUpgrade} when managed attachments is enabled and dropbox items do not exist.
        """

        didUpgrade = [False, ]

        def _hasDropboxAttachments(_self, txn):
            return succeed(False)
        self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)

        def _upgradeToManagedAttachments(_self, batchSize=10):
            didUpgrade[0] = True
            return succeed(None)
        self.patch(CalendarStoreFeatures, "upgradeToManagedAttachments", _upgradeToManagedAttachments)

        store = (yield self._initStore())

        upgrader = UpgradeDatabaseOtherStep(store)
        yield attachment_migration.doUpgrade(upgrader)
        self.assertFalse(didUpgrade[0])

        txn = upgrader.sqlStore.newTransaction()
        managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
        yield txn.commit()
        self.assertNotEqual(managed, None)

    @inlineCallbacks
    def test_upgradeFromDropboxOK(self):
        """
        Test L{attachment_migration.doUpgrade} when managed attachments is enabled and dropbox items exist.
        """

        didUpgrade = [False, ]

        def _hasDropboxAttachments(_self, txn):
            return succeed(True)
        self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)

        def _upgradeToManagedAttachments(_self, batchSize=10):
            didUpgrade[0] = True
            return succeed(None)
        self.patch(CalendarStoreFeatures, "upgradeToManagedAttachments", _upgradeToManagedAttachments)

        store = (yield self._initStore())

        upgrader = UpgradeDatabaseOtherStep(store)
        yield attachment_migration.doUpgrade(upgrader)
        self.assertFalse(didUpgrade[0])

        txn = upgrader.sqlStore.newTransaction()
        managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
        yield txn.commit()
        self.assertNotEqual(managed, None)

    @inlineCallbacks
    def test_upgradeAlreadyDone(self):
        """
        Test L{attachment_migration.doUpgrade} when managed attachments is enabled and migration already done.
        """

        didUpgrade = [False, ]

        def _hasDropboxAttachments(_self, txn):
            return succeed(True)
        self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)

        def _upgradeToManagedAttachments(_self, batchSize=10):
            didUpgrade[0] = True
            return succeed(None)
        self.patch(CalendarStoreFeatures, "upgradeToManagedAttachments", _upgradeToManagedAttachments)

        store = (yield self._initStore())
        txn = store.newTransaction()
        yield txn.setCalendarserverValue("MANAGED-ATTACHMENTS", "1")
        yield txn.commit()

        upgrader = UpgradeDatabaseOtherStep(store)
        yield attachment_migration.doUpgrade(upgrader)
        self.assertFalse(didUpgrade[0])

        txn = upgrader.sqlStore.newTransaction()
        managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
        yield txn.commit()
        self.assertNotEqual(managed, None)

    @inlineCallbacks
    def test_upgradeNotEnabled(self):
        """
        Test L{attachment_migration.doUpgrade} when managed attachments is disabled.
        """

        didUpgrade = [False, ]

        def _hasDropboxAttachments(_self, txn):
            return succeed(True)
        self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)

        def _upgradeToManagedAttachments(_self, batchSize=10):
            didUpgrade[0] = True
            return succeed(None)
        self.patch(CalendarStoreFeatures, "upgradeToManagedAttachments", _upgradeToManagedAttachments)

        store = (yield self._initStore(False))

        upgrader = UpgradeDatabaseOtherStep(store)
        yield attachment_migration.doUpgrade(upgrader)
        self.assertFalse(didUpgrade[0])

        txn = upgrader.sqlStore.newTransaction()
        managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
        yield txn.commit()
        self.assertEqual(managed, None)


class AttachmentMigrationTests(CommonStoreTests):
    """
    Tests for L{UpgradeDatabaseSchemaStep}.
    """

    @inlineCallbacks
    def setUp(self):
        self.patch(config, "EnableManagedAttachments", True)

        yield super(AttachmentMigrationTests, self).setUp()

        self._sqlCalendarStore.enableManagedAttachments = True

        txn = self.transactionUnderTest()
        cs = schema.CALENDARSERVER
        yield Delete(
            From=cs,
            Where=cs.NAME == "MANAGED-ATTACHMENTS"
        ).on(txn)
        yield self.commit()

    @inlineCallbacks
    def test_upgradeOrphanedAttachment(self):
        """
        Test L{attachment_migration.doUpgrade} when an orphaned attachment is present.
        """

        def _hasDropboxAttachments(_self, txn):
            return succeed(True)
        self.patch(CalendarStoreFeatures, "hasDropboxAttachments", _hasDropboxAttachments)

        # Create orphaned attachment
        dropboxID = "ABCD.dropbox"
        attachmentName = "test.txt"
        home = yield self.homeUnderTest(name="user01")
        at = schema.ATTACHMENT
        yield Insert(
            {
                at.CALENDAR_HOME_RESOURCE_ID: home._resourceID,
                at.DROPBOX_ID: dropboxID,
                at.CONTENT_TYPE: "text/plain",
                at.SIZE: 10,
                at.MD5: "abcd",
                at.PATH: attachmentName,
            }
        ).on(self.transactionUnderTest())
        yield self.commit()

        hasheduid = hashlib.md5(dropboxID).hexdigest()
        fp = self._sqlCalendarStore.attachmentsPath.child(hasheduid[0:2]).child(hasheduid[2:4]).child(hasheduid)
        fp.makedirs()
        fp = fp.child(attachmentName)
        fp.setContent("1234567890")

        self.assertTrue(os.path.exists(fp.path))

        upgrader = UpgradeDatabaseOtherStep(self._sqlCalendarStore)
        yield attachment_migration.doUpgrade(upgrader)

        txn = upgrader.sqlStore.newTransaction()
        managed = (yield txn.calendarserverValue("MANAGED-ATTACHMENTS", raiseIfMissing=False))
        count = (yield Select(
            [Count(at.DROPBOX_ID), ],
            From=at,
        ).on(txn))[0][0]
        yield txn.commit()
        self.assertEqual(count, 1)
        self.assertNotEqual(managed, None)

        self.assertTrue(os.path.exists(fp.path))
