/*
 * BEGIN FLOCK GPL
 *
 * Copyright Flock Inc. 2005-2009
 * http://flock.com
 *
 * This file may be used under the terms of the
 * GNU General Public License Version 2 or later (the "GPL"),
 * http://www.gnu.org/licenses/gpl.html
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
 * for the specific language governing rights and limitations under the
 * License.
 *
 * END FLOCK GPL
 */

#include "flockFaviconDataService.h"

#include "nsNetUtil.h"
#include "nsServiceManagerUtils.h"
#include "nsICategoryManager.h"
#include "nsIChannel.h"
#include "nsIHttpChannel.h"
#include "nsIChannelEventSink.h"
#include "nsIContentSniffer.h"
#include "nsIInterfaceRequestor.h"
#include "nsIStreamListener.h"
#include "nsISupportsPrimitives.h"
#include "nsToolkitCompsCID.h"
#include "plbase64.h"

#define CONTENT_SNIFFING_SERVICES "content-sniffing-services"

NS_IMPL_ISUPPORTS1(flockFaviconDataService, flockIFaviconDataService)

flockFaviconDataService::flockFaviconDataService()
{
  nsresult rv;
  mFaviconService = do_GetService(NS_FAVICONSERVICE_CONTRACTID, &rv);
  NS_WARN_IF_FALSE(NS_SUCCEEDED(rv), "could not get favicon service");
}

flockFaviconDataService::~flockFaviconDataService()
{
}

NS_IMETHODIMP
flockFaviconDataService::SetFaviconDataFromDataURL(nsIURI* aFaviconURI,
                                                   const nsAString& aDataURL,
                                                   PRTime aExpiration)
{
  nsresult rv;

  nsCOMPtr<nsIURI> dataURI;
  rv = NS_NewURI(getter_AddRefs(dataURI), aDataURL);
  NS_ENSURE_SUCCESS(rv, rv);

  // use the data: protocol handler to convert the data
  nsCOMPtr<nsIIOService> ioService = do_GetIOService(&rv);
  NS_ENSURE_SUCCESS(rv, rv);
  nsCOMPtr<nsIProtocolHandler> protocolHandler;
  rv = ioService->GetProtocolHandler("data", getter_AddRefs(protocolHandler));
  NS_ENSURE_SUCCESS(rv, rv);

  nsCOMPtr<nsIChannel> channel;
  rv = protocolHandler->NewChannel(dataURI, getter_AddRefs(channel));
  NS_ENSURE_SUCCESS(rv, rv);

  // blocking stream is OK for data URIs
  nsCOMPtr<nsIInputStream> stream;
  rv = channel->Open(getter_AddRefs(stream));
  NS_ENSURE_SUCCESS(rv, rv);

  PRUint32 available;
  rv = stream->Available(&available);
  NS_ENSURE_SUCCESS(rv, rv);
  if (available == 0)
    return NS_ERROR_FAILURE;

  // read all the decoded data
  PRUint8* buffer = static_cast<PRUint8*>
                               (nsMemory::Alloc(sizeof(PRUint8) * available));
  if (!buffer)
    return NS_ERROR_OUT_OF_MEMORY;
  PRUint32 numRead;
  rv = stream->Read(reinterpret_cast<char*>(buffer), available, &numRead);
  if (NS_FAILED(rv) || numRead != available) {
    nsMemory::Free(buffer);
    return rv;
  }

  nsCAutoString mimeType;
  rv = channel->GetContentType(mimeType);
  NS_ENSURE_SUCCESS(rv, rv);

  // SetFaviconData can now do the dirty work
  rv = mFaviconService->SetFaviconData(aFaviconURI, buffer, available,
                                       mimeType, aExpiration);
  nsMemory::Free(buffer);
  NS_ENSURE_SUCCESS(rv, rv);

  return NS_OK;
}

NS_IMETHODIMP
flockFaviconDataService::GetFaviconDataAsDataURL(nsIURI* aFaviconURI,
                                                 nsAString& aDataURL)
{
  nsresult rv;

  PRUint8* data;
  PRUint32 dataLen;
  nsCAutoString mimeType;

  rv = mFaviconService->GetFaviconData(aFaviconURI, mimeType, &dataLen, &data);
  if (NS_FAILED(rv))
    return rv;

  if (!data) {
    aDataURL.SetIsVoid(PR_TRUE);
    return NS_OK;
  }

  char* encoded = PL_Base64Encode(reinterpret_cast<const char*>(data),
                                  dataLen, nsnull);
  nsMemory::Free(data);

  if (!encoded)
    return NS_ERROR_OUT_OF_MEMORY;

  aDataURL.AssignLiteral("data:");
  aDataURL.Append(NS_ConvertUTF8toUTF16(mimeType));
  aDataURL.AppendLiteral(";base64,");
  aDataURL.Append(NS_ConvertUTF8toUTF16(encoded));

  nsMemory::Free(encoded);
  return NS_OK;
}

// A copy of NS_ConsumeStream from nsStreamUtils.cpp since that can't be
// used from external linkage
nsresult
ConsumeStream(nsIInputStream *stream, PRUint32 maxCount, nsACString &result)
{
  nsresult rv = NS_OK;
  result.Truncate();

  while (maxCount) {
    PRUint32 avail;
    rv = stream->Available(&avail);
    if (NS_FAILED(rv)) {
      if (rv == NS_BASE_STREAM_CLOSED)
        rv = NS_OK;
        break;
      }
      if (avail == 0)
        break;
      if (avail > maxCount)
          avail = maxCount;

      // resize result buffer
      PRUint32 length = result.Length();
      result.SetLength(length + avail);
      if (result.Length() != (length + avail))
        return NS_ERROR_OUT_OF_MEMORY;
      char *buf = result.BeginWriting() + length;

      PRUint32 n;
      rv = stream->Read(buf, avail, &n);
      if (NS_FAILED(rv))
        break;
      if (n != avail)
        result.SetLength(length + n);
      if (n == 0)
        break;
      maxCount -= n;
  }

  return rv;
}

// Swiped from nsFaviconService.cpp, and then tweaked to support a listener
class FaviconLoadListener : public nsIStreamListener,
                            public nsIInterfaceRequestor,
                            public nsIChannelEventSink
{
public:
  FaviconLoadListener(nsIFaviconService* aFaviconService,
                      flockIFaviconDataService* aFlockFaviconService,
                      nsIURI* aFaviconURI,
                      nsIChannel* aChannel,
                      flockIFaviconFetchListener* aFetchListener,
                      const nsAString& aTopic);

  NS_DECL_ISUPPORTS
  NS_DECL_NSIREQUESTOBSERVER
  NS_DECL_NSISTREAMLISTENER
  NS_DECL_NSIINTERFACEREQUESTOR
  NS_DECL_NSICHANNELEVENTSINK

private:
  ~FaviconLoadListener();

  nsCOMPtr<nsIFaviconService> mFaviconService;
  nsCOMPtr<flockIFaviconDataService> mFlockFaviconService;
  nsCOMPtr<nsIChannel> mChannel;
  nsCOMPtr<nsIURI> mPageURI;
  nsCOMPtr<nsIURI> mFaviconURI;

  nsCString mData;

  nsCOMPtr<flockIFaviconFetchListener> mFetchListener;
  nsString mTopic;
};

NS_IMPL_ISUPPORTS4(FaviconLoadListener,
                   nsIRequestObserver,
                   nsIStreamListener,
                   nsIInterfaceRequestor,
                   nsIChannelEventSink)

FaviconLoadListener::FaviconLoadListener(nsIFaviconService* aFaviconService,
                                         flockIFaviconDataService* aFlockFaviconService,
                                         nsIURI* aFaviconURI,
                                         nsIChannel* aChannel,
                                         flockIFaviconFetchListener* aFetchListener,
                                         const nsAString& aTopic) :
  mFaviconService(aFaviconService),
  mFlockFaviconService(aFlockFaviconService),
  mChannel(aChannel),
  mFaviconURI(aFaviconURI),
  mFetchListener(aFetchListener),
  mTopic(aTopic)
{
}

FaviconLoadListener::~FaviconLoadListener()
{
}

NS_IMETHODIMP
FaviconLoadListener::OnStartRequest(nsIRequest *aRequest, nsISupports *aContext)
{
  return NS_OK;
}

NS_IMETHODIMP
FaviconLoadListener::OnStopRequest(nsIRequest *aRequest, nsISupports *aContext,
                                   nsresult aStatusCode)
{
  nsresult rv;

  if (NS_FAILED(aStatusCode)) {
    if (mFetchListener) {
      mFetchListener->OnFaviconFetchError(mFaviconURI, mTopic);
    }
    return NS_OK;
  }

  nsCOMPtr<nsIHttpChannel> httpChannel = do_QueryInterface(aRequest);
  if (httpChannel) {
    PRUint32 responseStatus = 0;
    rv = httpChannel->GetResponseStatus(&responseStatus);
    NS_ENSURE_SUCCESS(rv, rv);

    if (responseStatus != 200) {
      if (mFetchListener) {
        mFetchListener->OnFaviconFetchStatusError(mFaviconURI,
                                                  responseStatus,
                                                  mTopic);
      }
      return NS_OK;
    }
  }

  if (mData.Length() == 0) {
    if (mFetchListener) {
      mFetchListener->OnFaviconFetchError(mFaviconURI, mTopic);
    }
    return NS_OK;
  }

  // sniff the MIME type
  nsCOMPtr<nsICategoryManager> categoryManager =
    do_GetService(NS_CATEGORYMANAGER_CONTRACTID, &rv);
  NS_ENSURE_SUCCESS(rv, rv);
  nsCOMPtr<nsISimpleEnumerator> sniffers;
  rv = categoryManager->EnumerateCategory(CONTENT_SNIFFING_SERVICES,
                                          getter_AddRefs(sniffers));
  NS_ENSURE_SUCCESS(rv, rv);

  nsCString mimeType;
  PRBool hasMore = PR_FALSE;
  while (mimeType.IsEmpty() && NS_SUCCEEDED(sniffers->HasMoreElements(&hasMore))
         && hasMore) {
    nsCOMPtr<nsISupports> snifferCIDSupports;
    rv = sniffers->GetNext(getter_AddRefs(snifferCIDSupports));
    NS_ENSURE_SUCCESS(rv, rv);

    nsCOMPtr<nsISupportsCString> snifferCIDSupportsCString =
      do_QueryInterface(snifferCIDSupports, &rv);
    NS_ENSURE_SUCCESS(rv, rv);

    nsCAutoString snifferCID;
    rv = snifferCIDSupportsCString->GetData(snifferCID);
    NS_ENSURE_SUCCESS(rv, rv);

    nsCOMPtr<nsIContentSniffer> sniffer = do_GetService(snifferCID.get(), &rv);
    NS_ENSURE_SUCCESS(rv, rv);

    sniffer->GetMIMETypeFromContent(
               aRequest,
               reinterpret_cast<PRUint8*>(const_cast<char*>(mData.get())),
               mData.Length(), mimeType);
    // ignore errors: mime type will be left empty and we'll try the next sniffer
  }

  if (mimeType.IsEmpty()) {
    // we can not handle favicons that do not have a recognisable MIME type
    if (mFetchListener) {
      mFetchListener->OnFaviconFetchError(mFaviconURI, mTopic);
    }
    return NS_OK;
  }

  // expire this favicon in one day
  PRTime expiration = PR_Now() +
                      (PRInt64)(24 * 60 * 60) * (PRInt64)PR_USEC_PER_SEC;

  // save the favicon data
  rv = mFaviconService->SetFaviconData(mFaviconURI,
               reinterpret_cast<PRUint8*>(const_cast<char*>(mData.get())),
               mData.Length(), mimeType, expiration);
  NS_ENSURE_SUCCESS(rv, rv);

  if (mFetchListener) {
    // now pull it out again, after any transformations that are done
    nsAutoString dataURL;
    rv = mFlockFaviconService->GetFaviconDataAsDataURL(mFaviconURI, dataURL);
    NS_ENSURE_SUCCESS(rv, rv);

    // and tell our listener about it
    mFetchListener->OnFaviconFetchSuccess(mFaviconURI, dataURL, mTopic);
  }

  return NS_OK;
}

NS_IMETHODIMP
FaviconLoadListener::OnDataAvailable(nsIRequest *aRequest,
                                     nsISupports *aContext,
                                     nsIInputStream *aInputStream,
                                     PRUint32 aOffset,
                                     PRUint32 aCount)
{
  nsCString buffer;
  nsresult rv = ConsumeStream(aInputStream, aCount, buffer);
  if (rv != NS_BASE_STREAM_WOULD_BLOCK && NS_FAILED(rv))
    return rv;

  mData.Append(buffer);
  return NS_OK;
}

NS_IMETHODIMP
FaviconLoadListener::GetInterface(const nsIID& uuid, void** aResult)
{
  return QueryInterface(uuid, aResult);
}

NS_IMETHODIMP
FaviconLoadListener::OnChannelRedirect(nsIChannel* oldChannel,
                                       nsIChannel* newChannel,
                                       PRUint32 flags)
{
  mChannel = newChannel;
  return NS_OK;
}

NS_IMETHODIMP
flockFaviconDataService::FetchFaviconData(nsIURI* aFaviconURI,
                                          flockIFaviconFetchListener* aFetchListener,
                                          const nsAString& aTopic)
{
  nsresult rv;

  nsCOMPtr<nsIIOService> ioservice = do_GetIOService(&rv);
  NS_ENSURE_SUCCESS(rv, rv);

  nsCOMPtr<nsIChannel> channel;
  rv = ioservice->NewChannelFromURI(aFaviconURI, getter_AddRefs(channel));
  NS_ENSURE_SUCCESS(rv, rv);

  nsCOMPtr<nsIStreamListener> listener =
    new FaviconLoadListener(mFaviconService, this, aFaviconURI, channel,
                            aFetchListener, aTopic);
  NS_ENSURE_TRUE(listener, NS_ERROR_OUT_OF_MEMORY);

  nsCOMPtr<nsIInterfaceRequestor> listenerRequestor =
    do_QueryInterface(listener, &rv);
  NS_ENSURE_SUCCESS(rv, rv);

  rv = channel->SetNotificationCallbacks(listenerRequestor);
  NS_ENSURE_SUCCESS(rv, rv);
  rv = channel->AsyncOpen(listener, nsnull);
  NS_ENSURE_SUCCESS(rv, rv);

  return NS_OK;
}
