/*
 * Copyright 2011-2016, Axel Dörfler, axeld@pinc-software.de.
 * Distributed under the terms of the MIT License.
 */


#include "IMAPConnectionWorker.h"

#include <Autolock.h>
#include <Messenger.h>

#include <AutoDeleter.h>

#include "IMAPFolder.h"
#include "IMAPMailbox.h"
#include "IMAPProtocol.h"


using IMAP::MessageUIDList;


static const uint32 kMaxFetchEntries = 500;
static const uint32 kMaxDirectDownloadSize = 4096;


class WorkerPrivate {
public:
	WorkerPrivate(IMAPConnectionWorker& worker)
		:
		fWorker(worker)
	{
	}

	IMAP::Protocol& Protocol()
	{
		return fWorker.fProtocol;
	}

	status_t AddFolders(BObjectList<IMAPFolder>& folders)
	{
		IMAPConnectionWorker::MailboxMap::iterator iterator
			= fWorker.fMailboxes.begin();
		for (; iterator != fWorker.fMailboxes.end(); iterator++) {
			IMAPFolder* folder = iterator->first;
			if (!folders.AddItem(folder))
				return B_NO_MEMORY;
		}
		return B_OK;
	}

	status_t SelectMailbox(IMAPFolder& folder)
	{
		return fWorker._SelectMailbox(folder, NULL);
	}

	status_t SelectMailbox(IMAPFolder& folder, uint32& nextUID)
	{
		return fWorker._SelectMailbox(folder, &nextUID);
	}

	IMAPMailbox* MailboxFor(IMAPFolder& folder)
	{
		return fWorker._MailboxFor(folder);
	}

	int32 BodyFetchLimit() const
	{
		return fWorker.fSettings.BodyFetchLimit();
	}

	uint32 MessagesExist() const
	{
		return fWorker._MessagesExist();
	}

	status_t EnqueueCommand(WorkerCommand* command)
	{
		return fWorker._EnqueueCommand(command);
	}

	void SyncCommandDone()
	{
		fWorker._SyncCommandDone();
	}

	void Quit()
	{
		fWorker.fStopped = true;
	}

private:
	IMAPConnectionWorker&	fWorker;
};


class WorkerCommand {
public:
	WorkerCommand()
		:
		fContinuation(false)
	{
	}

	virtual ~WorkerCommand()
	{
	}

	virtual	status_t Process(IMAPConnectionWorker& worker) = 0;

	virtual bool IsDone() const
	{
		return true;
	}

	bool IsContinuation() const
	{
		return fContinuation;
	}

	void SetContinuation()
	{
		fContinuation = true;
	}

private:
			bool				fContinuation;
};


/*!	All commands that inherit from this class will automatically maintain the
	worker's fSyncPending member, and will thus prevent syncing more than once
	concurrently.
*/
class SyncCommand : public WorkerCommand {
};


class QuitCommand : public WorkerCommand {
public:
	QuitCommand()
	{
	}

	virtual status_t Process(IMAPConnectionWorker& worker)
	{
		WorkerPrivate(worker).Quit();
		return B_OK;
	}
};


class CheckSubscribedFoldersCommand : public WorkerCommand {
public:
	virtual status_t Process(IMAPConnectionWorker& worker)
	{
		IMAP::Protocol& protocol = WorkerPrivate(worker).Protocol();

		// The main worker checks the subscribed folders, and creates
		// other workers as needed
		return worker.Owner().CheckSubscribedFolders(protocol,
			worker.UsesIdle());
	}
};


class FetchBodiesCommand : public SyncCommand, public IMAP::FetchListener {
public:
	FetchBodiesCommand(IMAPFolder& folder, IMAPMailbox& mailbox,
		MessageUIDList& entries, const BMessenger* replyTo = NULL)
		:
		fFolder(folder),
		fMailbox(mailbox),
		fEntries(entries)
	{
		folder.RegisterPendingBodies(entries, replyTo);
	}

	virtual status_t Process(IMAPConnectionWorker& worker)
	{
		IMAP::Protocol& protocol = WorkerPrivate(worker).Protocol();

		if (fEntries.empty())
			return B_OK;

		fUID = *fEntries.begin();
		fEntries.erase(fEntries.begin());

		status_t status = WorkerPrivate(worker).SelectMailbox(fFolder);
		if (status == B_OK) {
			printf("IMAP: fetch body for %" B_PRIu32 "\n", fUID);
			// Since RFC3501 does not specify whether the FETCH response may
			// alter the order of the message data items we request, we cannot
			// request more than a single UID at a time, or else we may not be
			// able to assign the data to the correct message beforehand.
			IMAP::FetchCommand fetch(fUID, fUID, IMAP::kFetchBody);
			fetch.SetListener(this);

			status = protocol.ProcessCommand(fetch);
		}
		if (status == B_OK)
			status = fFetchStatus;
		if (status != B_OK)
			fFolder.StoringBodyFailed(fRef, fUID, status);

		return status;
	}

	virtual bool IsDone() const
	{
		return fEntries.empty();
	}

	virtual bool FetchData(uint32 fetchFlags, BDataIO& stream, size_t& length)
	{
		fFetchStatus = fFolder.StoreBody(fUID, stream, length, fRef, fFile);
		return true;
	}

	virtual void FetchedData(uint32 fetchFlags, uint32 uid, uint32 flags)
	{
		if (fFetchStatus == B_OK)
			fFolder.BodyStored(fRef, fFile, uid);
	}

private:
	IMAPFolder&				fFolder;
	IMAPMailbox&			fMailbox;
	MessageUIDList			fEntries;
	uint32					fUID;
	entry_ref				fRef;
	BFile					fFile;
	status_t				fFetchStatus;
};


class FetchHeadersCommand : public SyncCommand, public IMAP::FetchListener {
public:
	FetchHeadersCommand(IMAPFolder& folder, IMAPMailbox& mailbox,
		MessageUIDList& uids, int32 bodyFetchLimit)
		:
		fFolder(folder),
		fMailbox(mailbox),
		fUIDs(uids),
		fBodyFetchLimit(bodyFetchLimit)
	{
	}

	virtual status_t Process(IMAPConnectionWorker& worker)
	{
		IMAP::Protocol& protocol = WorkerPrivate(worker).Protocol();

		status_t status = WorkerPrivate(worker).SelectMailbox(fFolder);
		if (status != B_OK)
			return status;

		printf("IMAP: fetch %" B_PRIuSIZE "u headers\n", fUIDs.size());

		IMAP::FetchCommand fetch(fUIDs, kMaxFetchEntries,
			IMAP::kFetchHeader | IMAP::kFetchFlags);
		fetch.SetListener(this);

		status = protocol.ProcessCommand(fetch);
		if (status != B_OK)
			return status;

		if (IsDone() && !fFetchBodies.empty()) {
			// Enqueue command to fetch the message bodies
			WorkerPrivate(worker).EnqueueCommand(new FetchBodiesCommand(fFolder,
				fMailbox, fFetchBodies));
		}

		return B_OK;
	}

	virtual bool IsDone() const
	{
		return fUIDs.empty();
	}

	virtual bool FetchData(uint32 fetchFlags, BDataIO& stream, size_t& length)
	{
		fFetchStatus = fFolder.StoreMessage(fetchFlags, stream, length,
			fRef, fFile);
		return true;
	}

	virtual void FetchedData(uint32 fetchFlags, uint32 uid, uint32 flags)
	{
		if (fFetchStatus == B_OK) {
			fFolder.MessageStored(fRef, fFile, fetchFlags, uid, flags);

			uint32 size = fMailbox.MessageSize(uid);
			if (fBodyFetchLimit < 0 || size < fBodyFetchLimit)
				fFetchBodies.push_back(uid);
		}
	}

private:
	IMAPFolder&				fFolder;
	IMAPMailbox&			fMailbox;
	MessageUIDList			fUIDs;
	MessageUIDList			fFetchBodies;
	uint32					fBodyFetchLimit;
	entry_ref				fRef;
	BFile					fFile;
	status_t				fFetchStatus;
};


class CheckMailboxesCommand : public SyncCommand {
public:
	CheckMailboxesCommand(IMAPConnectionWorker& worker)
		:
		fWorker(worker),
		fFolders(5, false),
		fState(INIT),
		fFolder(NULL),
		fMailbox(NULL)
	{
	}

	virtual status_t Process(IMAPConnectionWorker& worker)
	{
		IMAP::Protocol& protocol = WorkerPrivate(worker).Protocol();

		if (fState == INIT) {
			// Collect folders
			status_t status = WorkerPrivate(worker).AddFolders(fFolders);
			if (status != B_OK || fFolders.IsEmpty()) {
				fState = DONE;
				return status;
			}

			fState = SELECT;
		}

		if (fState == SELECT) {
			// Get next mailbox from list, and select it
			fFolder = fFolders.RemoveItemAt(fFolders.CountItems() - 1);
			if (fFolder == NULL) {
				for (int32 i = 0; i < fFetchCommands.CountItems(); i++) {
					WorkerPrivate(worker).EnqueueCommand(
						fFetchCommands.ItemAt(i));
				}

				fState = DONE;
				return B_OK;
			}

			fMailbox = WorkerPrivate(worker).MailboxFor(*fFolder);

			status_t status = WorkerPrivate(worker).SelectMailbox(*fFolder);
			if (status != B_OK)
				return status;

			fLastIndex = WorkerPrivate(worker).MessagesExist();
			fFirstIndex = fMailbox->CountMessages() + 1;
			if (fLastIndex > 0)
				fState = FETCH_ENTRIES;
		}

		if (fState == FETCH_ENTRIES) {
			status_t status = WorkerPrivate(worker).SelectMailbox(*fFolder);
			if (status != B_OK)
				return status;

			uint32 to = fLastIndex;
			uint32 from = fFirstIndex + kMaxFetchEntries < to
				? fLastIndex - kMaxFetchEntries : fFirstIndex;

			printf("IMAP: get entries from %" B_PRIu32 " to %" B_PRIu32 "\n",
				from, to);

			IMAP::MessageEntryList entries;
			IMAP::FetchMessageEntriesCommand fetch(entries, from, to, false);
			status = protocol.ProcessCommand(fetch);
			if (status != B_OK)
				return status;

			// Determine how much we need to download
			// TODO: also retrieve the header size, and only take the body
			// size into account if it's below the limit -- that does not
			// seem to be possible, though
			for (size_t i = 0; i < entries.size(); i++) {
				printf("%10" B_PRIu32 " %8" B_PRIu32 " bytes, flags: %#"
					B_PRIx32 "\n", entries[i].uid, entries[i].size,
					entries[i].flags);
				fMailbox->AddMessageEntry(from + i, entries[i].uid,
					entries[i].flags, entries[i].size);

				if (entries[i].uid > fFolder->LastUID()) {
					fTotalBytes += entries[i].size;
					fUIDsToFetch.push_back(entries[i].uid);
				} else {
					fFolder->SyncMessageFlags(entries[i].uid, entries[i].flags);
				}
			}

			fTotalEntries += fUIDsToFetch.size();
			fLastIndex = from - 1;

			if (from == 1) {
				fFolder->MessageEntriesFetched();

				if (fUIDsToFetch.size() > 0) {
					// Add pending command to fetch the message headers
					WorkerCommand* command = new FetchHeadersCommand(*fFolder,
						*fMailbox, fUIDsToFetch,
						WorkerPrivate(worker).BodyFetchLimit());
					if (!fFetchCommands.AddItem(command))
						delete command;

					fUIDsToFetch.clear();
				}
				fState = SELECT;
			}
		}

		return B_OK;
	}

	virtual bool IsDone() const
	{
		return fState == DONE;
	}

private:
	enum State {
		INIT,
		SELECT,
		FETCH_ENTRIES,
		DONE
	};

	IMAPConnectionWorker&	fWorker;
	BObjectList<IMAPFolder>	fFolders;
	State					fState;
	IMAPFolder*				fFolder;
	IMAPMailbox*			fMailbox;
	uint32					fFirstIndex;
	uint32					fLastIndex;
	uint64					fTotalEntries;
	uint64					fTotalBytes;
	WorkerCommandList		fFetchCommands;
	MessageUIDList			fUIDsToFetch;
};


class UpdateFlagsCommand : public WorkerCommand {
public:
	UpdateFlagsCommand(IMAPFolder& folder, IMAPMailbox& mailbox,
		MessageUIDList& entries, uint32 flags)
		:
		fFolder(folder),
		fMailbox(mailbox),
		fEntries(entries),
		fFlags(flags)
	{
	}

	virtual status_t Process(IMAPConnectionWorker& worker)
	{
		if (fEntries.empty())
			return B_OK;

		fUID = *fEntries.begin();
		fEntries.erase(fEntries.begin());

		status_t status = WorkerPrivate(worker).SelectMailbox(fFolder);
		if (status == B_OK) {
			IMAP::Protocol& protocol = WorkerPrivate(worker).Protocol();
			IMAP::SetFlagsCommand set(fUID, fFlags);
			status = protocol.ProcessCommand(set);
		}

		return status;
	}

	virtual bool IsDone() const
	{
		return fEntries.empty();
	}

private:
	IMAPFolder&				fFolder;
	IMAPMailbox&			fMailbox;
	MessageUIDList			fEntries;
	uint32					fUID;
	uint32					fFlags;
};


struct CommandDelete
{
	inline void operator()(WorkerCommand* command)
	{
		delete command;
	}
};


/*!	An auto deleter similar to ObjectDeleter that calls SyncCommandDone()
	for all SyncCommands.
*/
struct CommandDeleter : BPrivate::AutoDeleter<WorkerCommand, CommandDelete>
{
	CommandDeleter(IMAPConnectionWorker& worker, WorkerCommand* command)
		:
		BPrivate::AutoDeleter<WorkerCommand, CommandDelete>(command),
		fWorker(worker)
	{
	}

	~CommandDeleter()
	{
		if (dynamic_cast<SyncCommand*>(Get()) != NULL)
			WorkerPrivate(fWorker).SyncCommandDone();
	}

private:
	IMAPConnectionWorker&	fWorker;
};


// #pragma mark -


IMAPConnectionWorker::IMAPConnectionWorker(IMAPProtocol& owner,
	const Settings& settings, bool main)
	:
	fOwner(owner),
	fSettings(settings),
	fPendingCommandsSemaphore(-1),
	fIdleBox(NULL),
	fMain(main),
	fStopped(false)
{
	fExistsHandler.SetListener(this);
	fProtocol.AddHandler(fExistsHandler);

	fExpungeHandler.SetListener(this);
	fProtocol.AddHandler(fExpungeHandler);
}


IMAPConnectionWorker::~IMAPConnectionWorker()
{
	puts("worker quit");
	delete_sem(fPendingCommandsSemaphore);
	_Disconnect();
}


bool
IMAPConnectionWorker::HasMailboxes() const
{
	BAutolock locker(const_cast<IMAPConnectionWorker*>(this)->fLocker);
	return !fMailboxes.empty();
}


uint32
IMAPConnectionWorker::CountMailboxes() const
{
	BAutolock locker(const_cast<IMAPConnectionWorker*>(this)->fLocker);
	return fMailboxes.size();
}


void
IMAPConnectionWorker::AddMailbox(IMAPFolder* folder)
{
	BAutolock locker(fLocker);

	fMailboxes.insert(std::make_pair(folder, (IMAPMailbox*)NULL));

	// Prefer to have the INBOX in idle mode over other mail boxes
	if (fIdleBox == NULL || folder->MailboxName().ICompare("INBOX") == 0)
		fIdleBox = folder;
}


void
IMAPConnectionWorker::RemoveAllMailboxes()
{
	BAutolock locker(fLocker);

	// Reset listeners, and delete the mailboxes
	MailboxMap::iterator iterator = fMailboxes.begin();
	for (; iterator != fMailboxes.end(); iterator++) {
		iterator->first->SetListener(NULL);
		delete iterator->second;
	}

	fIdleBox = NULL;
	fMailboxes.clear();
}


status_t
IMAPConnectionWorker::Run()
{
	fPendingCommandsSemaphore = create_sem(0, "imap pending commands");
	if (fPendingCommandsSemaphore < 0)
		return fPendingCommandsSemaphore;

	fThread = spawn_thread(&_Worker, "imap connection worker",
		B_NORMAL_PRIORITY, this);
	if (fThread < 0)
		return fThread;

	resume_thread(fThread);
	return B_OK;
}


void
IMAPConnectionWorker::Quit()
{
	printf("IMAP: worker %p: enqueue quit\n", this);
	BAutolock qlocker(fQueueLocker);
	while (!fPendingCommands.IsEmpty())
		delete(fPendingCommands.RemoveItemAt(0));
	_EnqueueCommand(new QuitCommand());
}


status_t
IMAPConnectionWorker::EnqueueCheckSubscribedFolders()
{
	printf("IMAP: worker %p: enqueue check subscribed folders\n", this);
	return _EnqueueCommand(new CheckSubscribedFoldersCommand());
}


status_t
IMAPConnectionWorker::EnqueueCheckMailboxes()
{
	// Do not schedule checking mailboxes again if we're still working on
	// those.
	if (fSyncPending > 0)
		return B_OK;

	printf("IMAP: worker %p: enqueue check mailboxes\n", this);
	return _EnqueueCommand(new CheckMailboxesCommand(*this));
}


status_t
IMAPConnectionWorker::EnqueueFetchBody(IMAPFolder& folder, uint32 uid,
	const BMessenger& replyTo)
{
	IMAPMailbox* mailbox = _MailboxFor(folder);
	if (mailbox == NULL)
		return B_ENTRY_NOT_FOUND;

	std::vector<uint32> uids;
	uids.push_back(uid);

	return _EnqueueCommand(new FetchBodiesCommand(folder, *mailbox, uids,
		&replyTo));
}


status_t
IMAPConnectionWorker::EnqueueUpdateFlags(IMAPFolder& folder, uint32 uid,
	uint32 flags)
{
	IMAPMailbox* mailbox = _MailboxFor(folder);
	if (mailbox == NULL)
		return B_ENTRY_NOT_FOUND;

	std::vector<uint32> uids;
	uids.push_back(uid);

	return _EnqueueCommand(new UpdateFlagsCommand(folder, *mailbox, uids,
		flags));
}


// #pragma mark - Handler listener


void
IMAPConnectionWorker::MessageExistsReceived(uint32 count)
{
	printf("Message exists: %" B_PRIu32 "\n", count);
	fMessagesExist = count;

	// TODO: We might want to trigger another check even during sync
	// (but only one), if this isn't the result of a SELECT
	EnqueueCheckMailboxes();
}


void
IMAPConnectionWorker::MessageExpungeReceived(uint32 index)
{
	printf("Message expunge: %" B_PRIu32 "\n", index);
	if (fSelectedBox == NULL)
		return;

	BAutolock locker(fLocker);

	IMAPMailbox* mailbox = _MailboxFor(*fSelectedBox);
	if (mailbox != NULL) {
		mailbox->RemoveMessageEntry(index);
		// TODO: remove message from folder
	}
}


// #pragma mark - private


status_t
IMAPConnectionWorker::_Worker()
{
	status_t status = B_OK;

	while (!fStopped) {
		BAutolock qlocker(fQueueLocker);

		if (fPendingCommands.IsEmpty()) {
			if (!fIdle)
				_Disconnect();
			qlocker.Unlock();

			// TODO: in idle mode, we'd need to parse any incoming message here
			_WaitForCommands();
			continue;
		}

		WorkerCommand* command = fPendingCommands.RemoveItemAt(0);
		if (command == NULL)
			continue;

		qlocker.Unlock();
		BAutolock locker(fLocker);
		CommandDeleter deleter(*this, command);

		if (dynamic_cast<QuitCommand*>(command) == NULL) { // do not connect on QuitCommand
			status = _Connect();
			if (status != B_OK)
				break;
		}

		status = command->Process(*this);
		if (status != B_OK)
			break;

		if (!command->IsDone()) {
			deleter.Detach();
			command->SetContinuation();
			locker.Unlock();
			_EnqueueCommand(command);
		}
	}

	fOwner.WorkerQuit(this);
	return status;
}


/*!	Enqueues the given command to the worker queue. This method will take
	over ownership of the given command even in the error case.
*/
status_t
IMAPConnectionWorker::_EnqueueCommand(WorkerCommand* command)
{
	BAutolock qlocker(fQueueLocker);

	if (!fPendingCommands.AddItem(command)) {
		delete command;
		return B_NO_MEMORY;
	}

	if (dynamic_cast<SyncCommand*>(command) != NULL
		&& !command->IsContinuation())
		fSyncPending++;

	qlocker.Unlock();
	release_sem(fPendingCommandsSemaphore);
	return B_OK;
}


void
IMAPConnectionWorker::_WaitForCommands()
{
	int32 count = 1;
	get_sem_count(fPendingCommandsSemaphore, &count);
	if (count < 1)
		count = 1;

	while (acquire_sem_etc(fPendingCommandsSemaphore, count, 0,
			B_INFINITE_TIMEOUT) == B_INTERRUPTED);
}


status_t
IMAPConnectionWorker::_SelectMailbox(IMAPFolder& folder, uint32* _nextUID)
{
	if (fSelectedBox == &folder && _nextUID == NULL)
		return B_OK;

	IMAP::SelectCommand select(folder.MailboxName().String());

	status_t status = fProtocol.ProcessCommand(select);
	if (status == B_OK) {
		folder.SetUIDValidity(select.UIDValidity());
		if (_nextUID != NULL)
			*_nextUID = select.NextUID();
		fSelectedBox = &folder;
	}

	return status;
}


IMAPMailbox*
IMAPConnectionWorker::_MailboxFor(IMAPFolder& folder)
{
	MailboxMap::iterator found = fMailboxes.find(&folder);
	if (found == fMailboxes.end())
		return NULL;

	IMAPMailbox* mailbox = found->second;
	if (mailbox == NULL) {
		mailbox = new IMAPMailbox(fProtocol, folder.MailboxName());
		folder.SetListener(mailbox);
		found->second = mailbox;
	}
	return mailbox;
}


void
IMAPConnectionWorker::_SyncCommandDone()
{
	fSyncPending--;
}


bool
IMAPConnectionWorker::_IsQuitPending()
{
	BAutolock locker(fQueueLocker);
	WorkerCommand* nextCommand = fPendingCommands.ItemAt(0);
	return dynamic_cast<QuitCommand*>(nextCommand) != NULL;
}


status_t
IMAPConnectionWorker::_Connect()
{
	if (fProtocol.IsConnected())
		return B_OK;

	status_t status = B_INTERRUPTED;
	int tries = 10;
	while (tries-- > 0) {
		if (_IsQuitPending())
			break;
		status = fProtocol.Connect(fSettings.ServerAddress(),
			fSettings.Username(), fSettings.Password(), fSettings.UseSSL());
		if (status == B_OK)
			break;

		// Wait for 1 second, and try again
		snooze(1000000);
	}
	// TODO: if other workers are connected, but it fails for us, we need to
	// remove this worker, and reduce the number of concurrent connections
	if (status != B_OK)
		return status;

	//fIdle = fSettings.IdleMode() && fProtocol.Capabilities().Contains("IDLE");
	// TODO: Idle mode is not yet implemented!
	fIdle = false;
	return B_OK;
}


void
IMAPConnectionWorker::_Disconnect()
{
	fProtocol.Disconnect();
	fSelectedBox = NULL;
}


/*static*/ status_t
IMAPConnectionWorker::_Worker(void* _self)
{
	IMAPConnectionWorker* self = (IMAPConnectionWorker*)_self;
	status_t status = self->_Worker();

	delete self;
	return status;
}