/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

#include <SearchAPI.h>
#include <shellapi.h>
#include <objbase.h>
#include <string>

static const CLSID CLSID_CSearchManager = {
    0x7d096c5f,
    0xac08,
    0x4f1f,
    {0xbe, 0xb7, 0x5c, 0x22, 0xc5, 0x17, 0xce, 0x39}};
static const IID IID_ISearchManager = {
    0xab310581,
    0xac80,
    0x11d1,
    {0x8d, 0xf3, 0x00, 0xc0, 0x4f, 0xb6, 0xef, 0x69}};

static const WCHAR* const sFoldersToIndex[] = {L"\\Mail\\", L"\\ImapMail\\",
                                               L"\\News\\"};

struct RegKey {
  HKEY mRoot;
  LPCWSTR mSubKey;
  LPCWSTR mName;
  LPCWSTR mValue;

  RegKey(HKEY aRoot, LPCWSTR aSubKey, LPCWSTR aName, LPCWSTR aValue)
      : mRoot(aRoot), mSubKey(aSubKey), mName(aName), mValue(aValue) {}
};

static const RegKey* const sRegKeys[] = {
    new RegKey(HKEY_LOCAL_MACHINE,
               L"SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\PropertySystem\\"
               L"PropertyHandlers\\.wdseml",
               L"", L"{5FA29220-36A1-40f9-89C6-F4B384B7642E}"),
    new RegKey(HKEY_CLASSES_ROOT, L".wdseml", L"Content Type",
               L"message/rfc822"),
    new RegKey(HKEY_CLASSES_ROOT, L".wdseml\\PersistentHandler", L"",
               L"{5645c8c4-e277-11cf-8fda-00aa00a14f93}"),
    new RegKey(HKEY_CLASSES_ROOT,
               L".wdseml\\shellex\\{8895B1C6-B41F-4C1C-A562-0D564250836F}", L"",
               L"{b9815375-5d7f-4ce2-9245-c9d4da436930}"),
    new RegKey(
        HKEY_LOCAL_MACHINE,
        L"SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\explorer\\KindMap",
        L".wdseml", L"email;communication")};

HRESULT GetCrawlScopeManager(ISearchCrawlScopeManager** aCrawlScopeManager) {
  *aCrawlScopeManager = NULL;

  ISearchManager* searchManager;
  HRESULT hr = CoCreateInstance(CLSID_CSearchManager, NULL, CLSCTX_ALL,
                                IID_ISearchManager, (void**)&searchManager);
  if (SUCCEEDED(hr)) {
    ISearchCatalogManager* catalogManager;
    hr = searchManager->GetCatalog(L"SystemIndex", &catalogManager);
    if (SUCCEEDED(hr)) {
      hr = catalogManager->GetCrawlScopeManager(aCrawlScopeManager);
      catalogManager->Release();
    }
    searchManager->Release();
  }
  return hr;
}

LSTATUS SetRegistryKeys() {
  LSTATUS rv = ERROR_SUCCESS;
  for (uint32_t i = 0; rv == ERROR_SUCCESS && i < _countof(sRegKeys); i++) {
    const RegKey* key = sRegKeys[i];
    HKEY subKey;
    // Since we're administrator, we should be able to do this just fine
    rv = RegCreateKeyExW(key->mRoot, key->mSubKey, 0, NULL,
                         REG_OPTION_NON_VOLATILE,
                         KEY_ALL_ACCESS | KEY_WOW64_64KEY, NULL, &subKey, NULL);
    if (rv == ERROR_SUCCESS)
      rv = RegSetValueExW(subKey, key->mName, 0, REG_SZ, (LPBYTE)key->mValue,
                          (lstrlenW(key->mValue) + 1) * sizeof(WCHAR));
    RegCloseKey(subKey);
  }

  return rv;
}

int APIENTRY wWinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance,
                      LPWSTR lpCmdLine, int nCmdShow) {
  UNREFERENCED_PARAMETER(lpCmdLine);

  HRESULT hr =
      CoInitializeEx(NULL, COINIT_APARTMENTTHREADED | COINIT_DISABLE_OLE1DDE);
  if (SUCCEEDED(hr)) {
    int argc;
    LPWSTR* argv = CommandLineToArgvW(lpCmdLine, &argc);
    if (argc != 2) hr = E_INVALIDARG;
    if (SUCCEEDED(hr)) {
      ISearchCrawlScopeManager* crawlScopeManager;
      hr = GetCrawlScopeManager(&crawlScopeManager);
      if (SUCCEEDED(hr)) {
        if (*argv[0] == L'1') {
          // We first add the required registry entries
          LSTATUS rv = SetRegistryKeys();
          if (rv != ERROR_SUCCESS) hr = E_FAIL;

          // Next, we add rules for each of the three folders
          for (uint32_t i = 0; SUCCEEDED(hr) && i < _countof(sFoldersToIndex);
               i++) {
            std::wstring path = L"file:///";
            path.append(argv[1]);
            path.append(sFoldersToIndex[i]);
            // Add only if the rule isn't already there
            BOOL isIncluded = FALSE;
            hr = crawlScopeManager->IncludedInCrawlScope(path.c_str(),
                                                         &isIncluded);
            if (SUCCEEDED(hr) && !isIncluded)
              hr = crawlScopeManager->AddUserScopeRule(path.c_str(), TRUE, TRUE,
                                                       TRUE);
          }
        } else if (*argv[0] == L'0') {
          // This is simple, we just exclude the profile dir and override
          // children
          std::wstring path = L"file:///";
          path.append(argv[1]);
          hr = crawlScopeManager->AddUserScopeRule(path.c_str(), FALSE, TRUE,
                                                   TRUE);
        } else
          hr = E_INVALIDARG;

        if (SUCCEEDED(hr)) {
          hr = crawlScopeManager->SaveAll();
        }
        crawlScopeManager->Release();
      }
    }
    LocalFree(argv);
  }

  return hr;
}