##
# Copyright (c) 2015-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.python.log import Logger

from twisted.internet.defer import returnValue, inlineCallbacks
from twisted.python.failure import Failure

from twistedcaldav.accounting import emitAccounting

from txdav.caldav.datastore.sql import ManagedAttachment, CalendarBindRecord
from txdav.caldav.icalendarstore import ComponentUpdateState
from txdav.common.datastore.podding.migration.sync_metadata import CalendarMigrationRecord, \
    CalendarObjectMigrationRecord, AttachmentMigrationRecord
from txdav.common.datastore.podding.migration.work import HomeCleanupWork, MigrationCleanupWork
from txdav.common.datastore.sql_external import NotificationCollectionExternal
from txdav.common.datastore.sql_notification import NotificationCollection
from txdav.common.datastore.sql_tables import _HOME_STATUS_MIGRATING, _HOME_STATUS_DISABLED, \
    _HOME_STATUS_EXTERNAL, _HOME_STATUS_NORMAL
from txdav.common.idirectoryservice import DirectoryRecordNotFoundError

from functools import wraps
from uuid import uuid4
import datetime

log = Logger()

ACCOUNTING_TYPE = "migration"
ACCOUNTING_LOG = "migration.log"


def inTransactionWrapper(operation):
    """
    This wrapper converts an instance method that takes a transaction as its
    first parameter into one where the transaction parameter is an optional
    keyword argument. If the keyword argument is present and not None, then
    the instance method is called with that keyword as the first positional
    argument (i.e., almost a NoOp). If the keyword argument is not present,
    then a new transaction is created and the instance method called with
    it as the first positional argument, plus the call is wrapped with
    try/except/else to ensure proper commit and abort of the internally
    created transaction is done.

    So this wrapper allows for a method that requires a transaction to be run
    with either an existing transaction or one created just for the purpose
    of running it.

    @param operation: a callable that takes an L{IAsyncTransaction} as its first
        argument, and returns a value.
    """

    @wraps(operation)
    @inlineCallbacks
    def _inTxn(self, *args, **kwargs):
        label = self.label(operation.__name__)
        if "txn" in kwargs:
            txn = kwargs["txn"]
            del kwargs["txn"]
            result = yield operation(self, txn, *args, **kwargs)
            returnValue(result)
        else:
            txn = self.store.newTransaction(label=label)
            try:
                result = yield operation(self, txn, *args, **kwargs)
            except Exception as ex:
                f = Failure()
                yield txn.abort()
                log.error("{label} failed: {e}", label=label, e=str(ex))
                returnValue(f)
            else:
                yield txn.commit()
                returnValue(result)

    return _inTxn


# Cross-pod synchronization of an entire calendar home
class CrossPodHomeSync(object):

    BATCH_SIZE = 50

    def __init__(self, store, diruid, final=False, uselog=None):
        """
        @param store: the data store
        @type store: L{CommonDataStore}
        @param diruid: directory uid of the user whose home is to be sync'd
        @type diruid: L{str}
        @param final: indicates whether this is in the final sync stage with the remote home
            already disabled
        @type final: L{bool}
        @param uselog: additional logging written to this object
        @type: L{File}
        """

        self.store = store
        self.diruid = diruid
        self.disabledRemote = final
        self.uselog = uselog
        self.record = None
        self.homeId = None

    def label(self, detail):
        return "Cross-pod Migration Sync for {}: {}".format(self.diruid, detail)

    def accounting(self, logstr):
        emitAccounting(ACCOUNTING_TYPE, self.record if self.record is None else self.diruid, "{} {}\n".format(datetime.datetime.now().isoformat(), logstr), filename=ACCOUNTING_LOG)
        if self.uselog is not None:
            self.uselog.write("CrossPodHomeSync: {}\n".format(logstr))

    @inlineCallbacks
    def migrateHere(self):
        """
        This is a full, serialized version of a data migration (minus any directory
        update) that can be triggered via a command line tool. It is designed to
        minimize down time for the migrating user.
        """

        # Step 1 - initial full sync
        yield self.sync()

        # Step 2 - increment sync (since the initial sync may take a long time
        # to run we should do one incremental sync before bringing down the
        # account being migrated)
        yield self.sync()

        # Step 3 - disable remote home
        # NB Any failure from this point on will need to be caught and
        # handled by re-enabling the old home (and fixing any sharing state
        # that may have been changed)
        yield self.disableRemoteHome()

        # Step 4 - final incremental sync
        yield self.sync()

        # Step 5 - final overall sync of meta-data (including sharing re-linking)
        yield self.finalSync()

        # Step 6 - enable new home
        yield self.enableLocalHome()

        # Step 7 - remove remote home
        yield self.removeRemoteHome()

        # Step 8 - say phew! TODO: Actually alert everyone else
        pass

    @inlineCallbacks
    def sync(self):
        """
        Initiate a sync of the home. This is a simple data sync that does not
        reconcile sharing state etc. The L{finalSync} method will do a full
        sharing reconcile as well as disable the migration source home.
        """

        yield self.loadRecord()
        self.accounting("Starting: sync...")
        yield self.prepareCalendarHome()

        # Calendar list and calendar data
        yield self.syncCalendarList()

        # Sync home metadata such as alarms, default calendars, etc
        yield self.syncCalendarHomeMetaData()

        # Sync attachments
        yield self.syncAttachments()

        self.accounting("Completed: sync.\n")

    @inlineCallbacks
    def finalSync(self):
        """
        Do the final sync up of any additional data, re-link sharing bind
        rows, recalculate quota etc.
        """

        yield self.loadRecord()
        self.accounting("Starting: finalSync...")
        yield self.prepareCalendarHome()

        # Link attachments to resources: ATTACHMENT_CALENDAR_OBJECT table
        yield self.linkAttachments()

        # TODO: Re-write attachment URIs - not sure if we need this as reverse proxy may take care of it
        pass

        # Group attendee reconcile
        yield self.groupAttendeeReconcile()

        # Delegates reconcile
        yield self.delegateReconcile()

        # Shared collections reconcile (including group sharees)
        yield self.sharedByCollectionsReconcile()
        yield self.sharedToCollectionsReconcile()

        # Notifications
        yield self.notificationsReconcile()

        # iMIP tokens
        yield self.iMIPTokensReconcile()

        # Work items
        yield self.workItemsReconcile()

        self.accounting("Completed: finalSync.\n")

    @inTransactionWrapper
    @inlineCallbacks
    def disableRemoteHome(self, txn):
        """
        Mark the remote home as disabled. Also, prevent any scheduling jobs for the corresponding user
        from being run.
        """

        yield self.loadRecord()
        self.accounting("Starting: disableRemoteHome...")
        yield self.prepareCalendarHome()

        # Stop any work first
        remote_home = yield self._remoteHome(txn)
        yield remote_home.pauseWork()

        # Calendar home
        yield remote_home.setStatus(_HOME_STATUS_DISABLED)

        # Notification home
        notifications = yield self._remoteNotificationsHome(txn)
        yield notifications.setStatus(_HOME_STATUS_DISABLED)

        self.disabledRemote = True

        self.accounting("Completed: disableRemoteHome.\n")

    @inTransactionWrapper
    @inlineCallbacks
    def enableLocalHome(self, txn):
        """
        Mark the local home as enabled and remove any previously existing external home.
        """

        yield self.loadRecord()
        self.accounting("Starting: enableLocalHome...")
        yield self.prepareCalendarHome()

        # Disable any local external homes
        oldhome = yield txn.calendarHomeWithUID(self.diruid, status=_HOME_STATUS_EXTERNAL)
        if oldhome is not None:
            yield oldhome.setLocalStatus(_HOME_STATUS_DISABLED)
        oldnotifications = yield txn.notificationsWithUID(self.diruid, status=_HOME_STATUS_EXTERNAL)
        if oldnotifications:
            yield oldnotifications.setLocalStatus(_HOME_STATUS_DISABLED)

        # Enable the migrating ones
        newhome = yield txn.calendarHomeWithUID(self.diruid, status=_HOME_STATUS_MIGRATING)
        if newhome is not None:
            yield newhome.setStatus(_HOME_STATUS_NORMAL)
        newnotifications = yield txn.notificationsWithUID(self.diruid, status=_HOME_STATUS_MIGRATING)
        if newnotifications:
            yield newnotifications.setStatus(_HOME_STATUS_NORMAL)

        # Unpause work items
        yield newhome.unpauseWork()

        # Remove migration state
        yield MigrationCleanupWork.reschedule(
            txn,
            MigrationCleanupWork.notBeforeDelay,
            homeResourceID=newhome.id(),
        )

        # Purge the old ones
        yield HomeCleanupWork.reschedule(
            txn,
            HomeCleanupWork.notBeforeDelay,
            ownerUID=newhome.uid(),
        )

        self.accounting("Completed: enableLocalHome.\n")

    @inlineCallbacks
    def removeRemoteHome(self):
        """
        Remove all the old data on the remote pod.
        """

        # TODO: implement API on CommonHome to purge the old data without
        # any side-effects (scheduling, sharing etc). Also purge associated
        # data such as iMIP tokens, delegates, work items, etc
        yield self.loadRecord()
        self.accounting("Starting: removeRemoteHome...")
        yield self.prepareCalendarHome()
        yield self._migratedHome()

        self.accounting("Completed: removeRemoteHome.\n")

    @inTransactionWrapper
    def _migratedHome(self, txn):
        """
        Send cross-pod message to tell the old pod to remove the migrated data.
        """
        return txn.store().conduit.send_migrated_home(txn, self.diruid)

    @inlineCallbacks
    def loadRecord(self):
        """
        Initiate a sync of the home.
        """

        if self.record is None:
            self.record = yield self.store.directoryService().recordWithUID(self.diruid)
            if self.record is None:
                raise DirectoryRecordNotFoundError("Cross-pod Migration Sync missing directory record for {}".format(self.diruid))
            if self.record.thisServer():
                raise ValueError("Cross-pod Migration Sync cannot sync with user already on this server: {}".format(self.diruid))

    @inTransactionWrapper
    @inlineCallbacks
    def prepareCalendarHome(self, txn):
        """
        Make sure the inactive home to migrate into is present on this pod.
        """

        if self.homeId is None:
            home = yield self._localHome(txn)
            if home is None:
                if self.disabledRemote:
                    self.homeId = None
                else:
                    home = yield txn.calendarHomeWithUID(self.diruid, status=_HOME_STATUS_MIGRATING, create=True)
                    self.accounting("  Created new home collection to migrate into.")
            self.homeId = home.id() if home is not None else None

    @inTransactionWrapper
    @inlineCallbacks
    def syncCalendarHomeMetaData(self, txn):
        """
        Make sure the home meta-data (alarms, default calendars) is properly sync'd
        """

        self.accounting("Starting: syncCalendarHomeMetaData...")
        remote_home = yield self._remoteHome(txn)
        yield remote_home.readMetaData()

        calendars = yield CalendarMigrationRecord.querysimple(txn, calendarHomeResourceID=self.homeId)
        calendarIDMap = dict((item.remoteResourceID, item.localResourceID) for item in calendars)

        local_home = yield self._localHome(txn)
        yield local_home.copyMetadata(remote_home, calendarIDMap)

        self.accounting("Completed: syncCalendarHomeMetaData.")

    @inlineCallbacks
    def _remoteHome(self, txn):
        """
        Create a synthetic external home object that maps to the actual remote home.
        """

        from txdav.caldav.datastore.sql_external import CalendarHomeExternal
        resourceID = yield txn.store().conduit.send_home_resource_id(txn, self.record, migrating=True)
        home = CalendarHomeExternal.makeSyntheticExternalHome(txn, self.record.uid, resourceID) if resourceID is not None else None
        if self.disabledRemote:
            home._migratingHome = True
        returnValue(home)

    @inlineCallbacks
    def _remoteNotificationsHome(self, txn):
        """
        Create a synthetic external home object that maps to the actual remote home.
        """

        notifications = yield NotificationCollectionExternal.notificationsWithUID(txn, self.diruid, create=True)
        if self.disabledRemote:
            notifications._migratingHome = True
        returnValue(notifications)

    def _localHome(self, txn):
        """
        Get the home on this pod that will have data migrated to it.
        """

        return txn.calendarHomeWithUID(self.diruid, status=_HOME_STATUS_MIGRATING)

    @inlineCallbacks
    def syncCalendarList(self):
        """
        Synchronize each owned calendar.
        """

        self.accounting("Starting: syncCalendarList...")

        # Remote sync details
        remote_sync_state = yield self.getCalendarSyncList()
        self.accounting("  Found {} remote calendars to sync.".format(len(remote_sync_state)))

        # Get local sync details from local DB
        local_sync_state = yield self.getSyncState()
        self.accounting("  Found {} local calendars to sync.".format(len(local_sync_state)))

        # Remove local calendars no longer on the remote side
        yield self.purgeLocal(local_sync_state, remote_sync_state)

        # Sync each calendar that matches on both sides
        for remoteID in remote_sync_state.keys():
            yield self.syncCalendar(remoteID, local_sync_state, remote_sync_state)

        self.accounting("Completed: syncCalendarList.")

    @inTransactionWrapper
    @inlineCallbacks
    def getCalendarSyncList(self, txn):
        """
        Get the names and sync-tokens for each remote owned calendar.
        """

        # List of calendars from the remote side
        home = yield self._remoteHome(txn)
        if home is None:
            returnValue(None)
        calendars = yield home.loadChildren()
        results = {}
        for calendar in calendars:
            if calendar.owned():
                sync_token = yield calendar.syncToken()
                results[calendar.id()] = CalendarMigrationRecord.make(
                    calendarHomeResourceID=home.id(),
                    remoteResourceID=calendar.id(),
                    localResourceID=0,
                    lastSyncToken=sync_token,
                )

        returnValue(results)

    @inTransactionWrapper
    @inlineCallbacks
    def getSyncState(self, txn):
        """
        Get local synchronization state for the home being migrated.
        """
        records = yield CalendarMigrationRecord.querysimple(
            txn, calendarHomeResourceID=self.homeId
        )
        returnValue(dict([(record.remoteResourceID, record) for record in records]))

    @inTransactionWrapper
    @inlineCallbacks
    def updateSyncState(self, txn, stateRecord, newSyncToken):
        """
        Update or insert an L{CalendarMigrationRecord} with the new specified sync token.
        """
        if stateRecord.isnew():
            stateRecord.lastSyncToken = newSyncToken
            yield stateRecord.insert(txn)
        else:
            # The existing stateRecord has a stale txn, but valid column values. We have
            # to duplicate it before we can give it a different txn.
            stateRecord = stateRecord.duplicate()
            stateRecord.transaction = txn
            yield stateRecord.update(lastSyncToken=newSyncToken)

    @inTransactionWrapper
    @inlineCallbacks
    def purgeLocal(self, txn, local_sync_state, remote_sync_state):
        """
        Remove (silently - i.e., no scheduling) local calendars that are no longer on the remote side.

        @param txn: transaction to use
        @type txn: L{CommonStoreTransaction}
        @param local_sync_state: local sync state
        @type local_sync_state: L{dict}
        @param remote_sync_state: remote sync state
        @type remote_sync_state: L{dict}
        """
        home = yield self._localHome(txn)
        for localID in set(local_sync_state.keys()) - set(remote_sync_state.keys()):
            calendar = yield home.childWithID(local_sync_state[localID].localResourceID)
            if calendar is not None:
                yield calendar.purge()
            del local_sync_state[localID]
            self.accounting("  Purged calendar local-id={} that no longer exists on the remote pod.".format(localID))

    @inlineCallbacks
    def syncCalendar(self, remoteID, local_sync_state, remote_sync_state):
        """
        Sync the contents of a calendar from the remote side. The local calendar may need to be created
        on initial sync. Make use of sync tokens to avoid unnecessary work.

        @param remoteID: id of the remote calendar to sync
        @type remoteID: L{int}
        @param local_sync_state: local sync state
        @type local_sync_state: L{dict}
        @param remote_sync_state: remote sync state
        @type remote_sync_state: L{dict}
        """

        self.accounting("Starting: syncCalendar.")

        # See if we need to create the local one first
        if remoteID not in local_sync_state:
            localID = yield self.newCalendar()
            local_sync_state[remoteID] = CalendarMigrationRecord.make(
                calendarHomeResourceID=self.homeId,
                remoteResourceID=remoteID,
                localResourceID=localID,
                lastSyncToken=None,
            )
            self.accounting("  Created new calendar local-id={}, remote-id={}.".format(localID, remoteID))
        else:
            localID = local_sync_state.get(remoteID).localResourceID
            self.accounting("  Updating calendar local-id={}, remote-id={}.".format(localID, remoteID))
        local_record = local_sync_state.get(remoteID)

        remote_token = remote_sync_state[remoteID].lastSyncToken
        if local_record.lastSyncToken != remote_token:
            # Sync meta-data such as name, alarms, supported-components, transp, etc
            yield self.syncCalendarMetaData(local_record)

            # Sync object resources
            changed, removed = yield self.findObjectsToSync(local_record)
            self.accounting("  Calendar objects changed={}, removed={}.".format(len(changed), len(removed)))
            yield self.purgeDeletedObjectsInBatches(local_record, removed)
            yield self.updateChangedObjectsInBatches(local_record, changed)

        yield self.updateSyncState(local_record, remote_token)
        self.accounting("Completed: syncCalendar.")

    @inTransactionWrapper
    @inlineCallbacks
    def newCalendar(self, txn):
        """
        Create a new local calendar to sync remote data to. We don't care about the name
        of the calendar right now - it will be sync'd later.
        """

        home = yield self._localHome(txn)
        calendar = yield home.createChildWithName(str(uuid4()))
        returnValue(calendar.id())

    @inTransactionWrapper
    @inlineCallbacks
    def syncCalendarMetaData(self, txn, migrationRecord):
        """
        Sync the metadata of a calendar from the remote side.

        @param migrationRecord: current migration record
        @type localID: L{CalendarMigrationRecord}
        """

        # Remote changes
        remote_home = yield self._remoteHome(txn)
        remote_calendar = yield remote_home.childWithID(migrationRecord.remoteResourceID)
        if remote_calendar is None:
            returnValue(None)

        # Check whether the deleted set items
        local_home = yield self._localHome(txn)
        local_calendar = yield local_home.childWithID(migrationRecord.localResourceID)
        yield local_calendar.copyMetadata(remote_calendar)
        self.accounting("  Copied calendar meta-data for calendar local-id={0.localResourceID}, remote-id={0.remoteResourceID}.".format(migrationRecord))

    @inTransactionWrapper
    @inlineCallbacks
    def findObjectsToSync(self, txn, migrationRecord):
        """
        Find the set of object resources that need to be sync'd from the remote