/*
 *  Copyright 2005-2009 Fabrice Colin
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */

#include <stdio.h>
#include <string.h>
#include <strings.h>
#include <stdarg.h>
#ifdef HAVE_REGEX_H
#include <regex.h>
#endif
#include <stdlib.h>
#include <iostream>

#include <curl/curl.h>

#include "Url.h"
#include "CurlDownloader.h"

using namespace std;

struct ContentInfo
{
	char *m_pContent;
	size_t m_contentLen;
	string m_lastModified;
};

static void freeContentInfo(struct ContentInfo *pInfo)
{
	if (pInfo == NULL)
	{
		return;
	}

	if (pInfo->m_pContent != NULL)
	{
		free(pInfo->m_pContent);
		pInfo->m_pContent = NULL;
		pInfo->m_contentLen = 0;
	}
}

static size_t writeCallback(void *pData, size_t dataSize, size_t elementsCount, void *pStream)
{
	ContentInfo *pInfo = NULL;
	size_t totalSize = elementsCount * dataSize;

	if (pStream == NULL)
	{
		return 0;
	}
	pInfo = (ContentInfo *)pStream;

	char *pNewContent = (char*)realloc(pInfo->m_pContent, pInfo->m_contentLen + totalSize + 1);
	if (pNewContent == NULL)
	{
#ifdef DEBUG
		clog << "writeCallback: failed to enlarge buffer" << endl;
#endif
		freeContentInfo(pInfo);
		return 0;
	}

	pInfo->m_pContent = pNewContent;
	memcpy(pInfo->m_pContent + pInfo->m_contentLen, pData, totalSize);
	pInfo->m_contentLen += totalSize;
	pInfo->m_pContent[pInfo->m_contentLen] = '\0';

	if (totalSize < strlen((const char*)pData))
	{
		void *pBadChar = NULL;

		// There's a NULL character in the buffer ? Replace it
		while ((pBadChar = memchr((void*)pInfo->m_pContent, '\0', pInfo->m_contentLen)) != NULL)
		{
			((char*)pBadChar)[0] = ' ';
#ifdef DEBUG
			clog << "writeCallback: bad character" << endl;
#endif
		}
	}

	return totalSize;
}

static size_t headerCallback(void *pData, size_t dataSize, size_t elementsCount, void *pStream)
{
	ContentInfo *pInfo = NULL;
	size_t totalSize = elementsCount * dataSize;

	if ((pStream == NULL) ||
		(pData == NULL) ||
		(totalSize == 0))
	{
		return 0;
	}
	pInfo = (ContentInfo *)pStream;

	string header((const char*)pData, totalSize);

	string::size_type pos = header.find("Last-Modified: ");
	if (pos != string::npos)
	{
		pInfo->m_lastModified = header.substr(15);
#ifdef DEBUG
		clog << "headerCallback: Last-Modified " << pInfo->m_lastModified << endl;
#endif
	}

	return totalSize;
}

unsigned int CurlDownloader::m_initialized = 0;

CurlDownloader::CurlDownloader() :
	DownloaderInterface()
{
	if (m_initialized == 0)
	{
		// Initialize
		curl_global_init(CURL_GLOBAL_ALL);

		++m_initialized;
	}
}

CurlDownloader::~CurlDownloader()
{
	--m_initialized;
	if (m_initialized == 0)
	{
		// Shutdown
		curl_global_cleanup();
	}
}

//
// Implementation of DownloaderInterface
//

/// Retrieves the specified document; NULL if error.
Document *CurlDownloader::retrieveUrl(const DocumentInfo &docInfo)
{
	Document *pDocument = NULL;
	string url(Url::escapeUrl(docInfo.getLocation()));
	unsigned int redirectionsCount = 0;

	if (url.empty() == true)
	{
#ifdef DEBUG
		clog << "CurlDownloader::retrieveUrl: no URL specified !" << endl;
#endif
		return NULL;
	}

	// Create a session
	CURL *pCurlHandler = curl_easy_init();
	if (pCurlHandler != NULL)
	{
		ContentInfo *pContentInfo = new ContentInfo;

		pContentInfo->m_pContent = NULL;
		pContentInfo->m_contentLen = 0;

		curl_easy_setopt(pCurlHandler, CURLOPT_AUTOREFERER, 1);
		curl_easy_setopt(pCurlHandler, CURLOPT_FOLLOWLOCATION, 1);
		curl_easy_setopt(pCurlHandler, CURLOPT_MAXREDIRS, 10);
		curl_easy_setopt(pCurlHandler, CURLOPT_USERAGENT, m_userAgent.c_str());
		curl_easy_setopt(pCurlHandler, CURLOPT_NOSIGNAL, (long)1);
		curl_easy_setopt(pCurlHandler, CURLOPT_TIMEOUT, (long)m_timeout);
#ifndef DEBUG
		curl_easy_setopt(pCurlHandler, CURLOPT_NOPROGRESS, 1);
#endif
		curl_easy_setopt(pCurlHandler, CURLOPT_WRITEFUNCTION, writeCallback);
		curl_easy_setopt(pCurlHandler, CURLOPT_WRITEDATA, pContentInfo);
		curl_easy_setopt(pCurlHandler, CURLOPT_HEADERFUNCTION, headerCallback);
		curl_easy_setopt(pCurlHandler, CURLOPT_HEADERDATA, pContentInfo);

		// Is a proxy defined ?
		// Curl automatically checks and makes use of the *_proxy environment variables 
		if ((m_proxyAddress.empty() == false) &&
			(m_proxyPort > 0))
		{
			curl_proxytype proxyType = CURLPROXY_HTTP;

			curl_easy_setopt(pCurlHandler, CURLOPT_PROXY, m_proxyAddress.c_str());
			curl_easy_setopt(pCurlHandler, CURLOPT_PROXYPORT, m_proxyPort);
			// Type defaults to HTTP
			if (m_proxyType.empty() == false)
			{
				if (m_proxyType == "SOCKS4")
				{
					proxyType = CURLPROXY_SOCKS4;
				}
				else if (m_proxyType == "SOCKS5")
				{
					proxyType = CURLPROXY_SOCKS5;
				}
			}
			curl_easy_setopt(pCurlHandler, CURLOPT_PROXYTYPE, proxyType);
		}

#ifdef DEBUG
		clog << "CurlDownloader::retrieveUrl: URL is " << url << endl;
#endif
		while (redirectionsCount < 10)
		{
			curl_easy_setopt(pCurlHandler, CURLOPT_URL, url.c_str());

			if (m_method == "POST")
			{
				curl_easy_setopt(pCurlHandler, CURLOPT_POST, 1);
				if (m_postFields.empty() == false)
				{
					curl_easy_setopt(pCurlHandler, CURLOPT_POSTFIELDS, m_postFields.c_str());
				}
			}

			CURLcode res = curl_easy_perform(pCurlHandler);
			if ((res == CURLE_OK) &&
				(pContentInfo->m_pContent != NULL) &&
				(pContentInfo->m_contentLen > 0))
			{
				char *pContentType = NULL;

				// Copy the document content
				pDocument = new Document(docInfo);
				pDocument->setData(pContentInfo->m_pContent, pContentInfo->m_contentLen);
				pDocument->setLocation(url);
				pDocument->setSize((off_t )pContentInfo->m_contentLen);

				// What's the Content-Type ?
				res = curl_easy_getinfo(pCurlHandler, CURLINFO_CONTENT_TYPE, &pContentType);
				if ((res == CURLE_OK) &&
					(pContentType != NULL))
				{
					pDocument->setType(pContentType);
				}

				// The Last-Modified date ?
				if (pContentInfo->m_lastModified.empty() == false)
				{
					pDocument->setTimestamp(pContentInfo->m_lastModified);
				}

#ifdef HAVE_REGEX_H
				regex_t refreshRegex;
				regmatch_t pMatches[2];

				// Any REFRESH META tag ?
				// Look for <meta http-equiv="refresh" content="SECS;url=URL">
				if (regcomp(&refreshRegex,
					"<meta http-equiv=\"refresh\" content=\"([0-9]*);url=([^\"]*)\">"
					REG_EXTENDED|REG_ICASE) == 0)
				{
					if (regexec(&refreshRegex, pContentInfo->m_pContent, 2,
						pMatches, REG_NOTBOL|REG_NOTEOL) == 0)
					{
						url = pMatches[1];
#ifdef DEBUG
						clog << "CurlDownloader::retrieveUrl: redirected to URL " << url << endl;
#endif
						delete pDocument;
						pDocument = NULL;
						freeContentInfo(pContentInfo);
						++redirectionsCount;
						continue;
					}
#ifdef DEBUG
					else clog << "CurlDownloader::retrieveUrl: no REFRESH META tag" << endl;
#endif

					regfree(&refreshRegex);
				}
#ifdef DEBUG
				else clog << "CurlDownloader::retrieveUrl: couldn't look for a REFRESH META tag" << endl;
#endif
#endif
			}
			else
			{
				clog << "Couldn't download " << url << ": " << curl_easy_strerror(res) << endl;
			}

			break;
		}

		freeContentInfo(pContentInfo);
		delete pContentInfo;

		// Cleanup
		curl_easy_cleanup(pCurlHandler);
	}

	return pDocument;
}