//========================================================================
//
// pdf-fullrewrite.cc
//
// Copyright 2007 Julien Rebetez
// Copyright 2012 Fabio D'Urso
//
//========================================================================

#include "GlobalParams.h"
#include "Error.h"
#include "Object.h"
#include "PDFDoc.h"
#include "XRef.h"
#include "goo/GooString.h"
#include "utils/parseargs.h"

static bool compareDocuments(PDFDoc *origDoc, PDFDoc *newDoc);
static bool compareObjects(const Object *objA, const Object *objB);

static char ownerPassword[33] = "\001";
static char userPassword[33] = "\001";
static bool forceIncremental = false;
static bool checkOutput = false;
static bool printHelp = false;

static const ArgDesc argDesc[] = {
  {"-opw",    argString,   ownerPassword,    sizeof(ownerPassword),
   "owner password (for encrypted files)"},
  {"-upw",    argString,   userPassword,     sizeof(userPassword),
   "user password (for encrypted files)"},
  {"-i",      argFlag,     &forceIncremental,0,
   "incremental update mode"},
  {"-check",  argFlag,     &checkOutput,     0,
   "verify the generated document"},
  {"-h",      argFlag,     &printHelp,       0,
   "print usage information"},
  {"-help",   argFlag,     &printHelp,       0,
   "print usage information"},
  {"--help",  argFlag,     &printHelp,       0,
   "print usage information"},
  {"-?",      argFlag,     &printHelp,       0,
   "print usage information"},
  { }
};

int main (int argc, char *argv[])
{
  PDFDoc *doc = nullptr;
  PDFDoc *docOut = nullptr;
  GooString *inputName = nullptr;
  GooString *outputName = nullptr;
  GooString *ownerPW = nullptr;
  GooString *userPW = nullptr;
  int res = 0;

  // parse args
  bool ok = parseArgs(argDesc, &argc, argv);
  if (!ok || (argc < 3) || printHelp) {
    printUsage(argv[0], "INPUT-FILE OUTPUT-FILE", argDesc);
    if (!printHelp) {
      res = 1;
    }
    goto done;
  }

  inputName = new GooString(argv[1]);
  outputName = new GooString(argv[2]);

  if (ownerPassword[0] != '\001') {
    ownerPW = new GooString(ownerPassword);
  }
  if (userPassword[0] != '\001') {
    userPW = new GooString(userPassword);
  }

  // load input document
  globalParams = std::make_unique<GlobalParams>();
  doc = new PDFDoc(inputName, ownerPW, userPW);
  if (!doc->isOk()) {
    fprintf(stderr, "Error loading input document\n");
    res = 1;
    goto done;
  }

  // save it back (in rewrite or incremental update mode)
  if (doc->saveAs(outputName, forceIncremental ? writeForceIncremental : writeForceRewrite) != 0) {
    fprintf(stderr, "Error saving document\n");
    res = 1;
    goto done;
  }

  if (checkOutput) {
    // open the generated document to verify it
    docOut = new PDFDoc(outputName, ownerPW, userPW);
    if (!docOut->isOk()) {
      fprintf(stderr, "Error loading generated document\n");
      res = 1;
    } else if (!compareDocuments(doc, docOut)) {
      fprintf(stderr, "Verification failed\n");
      res = 1;
    }
  } else {
    delete outputName;
  }

done:
  delete docOut;
  delete doc;
  delete userPW;
  delete ownerPW;
  return res;
}

static bool compareDictionaries(Dict *dictA, Dict *dictB)
{
  const int length = dictA->getLength();
  if (dictB->getLength() != length)
    return false;

  /* Check that every key in dictA is contained in dictB.
   * Since keys are unique and we've already checked that dictA and dictB
   * contain the same number of entries, we don't need to check that every key
   * in dictB is also contained in dictA */
  for (int i = 0; i < length; ++i) {
    const char *key = dictA->getKey(i);
    const Object &valA = dictA->getValNF(i);
    const Object &valB = dictB->lookupNF(key);
    if (!compareObjects(&valA, &valB))
      return false;
  }

  return true;
}

static bool compareObjects(const Object *objA, const Object *objB)
{
  switch (objA->getType()) {
    case objBool:
    {
      if (objB->getType() != objBool) {
        return false;
      } else {
        return (objA->getBool() == objB->getBool());
      }
    }
    case objInt:
    case objInt64:
    case objReal:
    {
      if (!objB->isNum()) {
        return false;
      } else {
        // Fuzzy comparison
        const double diff = objA->getNum() - objB->getNum();
        return (-0.01 < diff) && (diff < 0.01);
      }
    }
    case objString:
    {
      if (objB->getType() != objString) {
        return false;
      } else {
        const GooString *strA = objA->getString();
        const GooString *strB = objB->getString();
        return (strA->cmp(strB) == 0);
      }
    }
    case objName:
    {
      if (objB->getType() != objName) {
        return false;
      } else {
        GooString nameA(objA->getName());
        GooString nameB(objB->getName());
        return (nameA.cmp(&nameB) == 0);
      }
    }
    case objNull:
    {
      if (objB->getType() != objNull) {
        return false;
      } else {
        return true;
      }
    }
    case objArray:
    {
      if (objB->getType() != objArray) {
        return false;
      } else {
        Array *arrayA = objA->getArray();
        Array *arrayB = objB->getArray();
        const int length = arrayA->getLength();
        if (arrayB->getLength() != length) {
          return false;
        } else {
          for (int i = 0; i < length; ++i) {
            const Object &elemA = arrayA->getNF(i);
            const Object &elemB = arrayB->getNF(i);
            if (!compareObjects(&elemA, &elemB)) {
              return false;
            }
          }
          return true;
        }
      }
    }
    case objDict:
    {
      if (objB->getType() != objDict) {
        return false;
      } else {
        Dict *dictA = objA->getDict();
        Dict *dictB = objB->getDict();
        return compareDictionaries(dictA, dictB);
      }
    }
    case objStream:
    {
      if (objB->getType() != objStream) {
        return false;
      } else {
        Stream *streamA = objA->getStream();
        Stream *streamB = objB->getStream();
        if (!compareDictionaries(streamA->getDict(), streamB->getDict())) {
          return false;
        } else {
          int c;
          streamA->reset();
          streamB->reset();
          do
          {
            c = streamA->getChar();
            if (c != streamB->getChar()) {
              return false;
            }
          } while (c != EOF);
          return true;
        }
      }
      return true;
    }
    case objRef:
    {
      if (objB->getType() != objRef) {
        return false;
      } else {
        const Ref refA = objA->getRef();
        const Ref refB = objB->getRef();
        return refA == refB;
      }
    }
    default:
    {
      fprintf(stderr, "compareObjects failed: unexpected object type %u\n", objA->getType());
      return false;
    }
  }
}

static bool compareDocuments(PDFDoc *origDoc, PDFDoc *newDoc)
{
  bool result = true;
  XRef *origXRef = origDoc->getXRef();
  XRef *newXRef = newDoc->getXRef();

  // Make sure that special flags are set in both documents
  origXRef->scanSpecialFlags();
  newXRef->scanSpecialFlags();

  // Compare XRef tables' size
  const int origNumObjects = origXRef->getNumObjects();
  const int newNumObjects = newXRef->getNumObjects();
  if (forceIncremental && origXRef->isXRefStream()) {
    // In case of incremental update, expect a new entry to be appended to store the new XRef stream
    if (origNumObjects+1 != newNumObjects) {
      fprintf(stderr, "XRef table: Unexpected number of entries (%d+1 != %d)\n", origNumObjects, newNumObjects);
      result = false;
    }
  } else {
    // In all other cases the number of entries must be the same
    if (origNumObjects != newNumObjects) {
      fprintf(stderr, "XRef table: Different number of entries (%d != %d)\n", origNumObjects, newNumObjects);
      result = false;
    }
  }

  // Compare each XRef entry
  const int numObjects = (origNumObjects < newNumObjects) ? origNumObjects : newNumObjects;
  for (int i = 0; i < numObjects; ++i) {
    XRefEntryType origType = origXRef->getEntry(i)->type;
    XRefEntryType newType = newXRef->getEntry(i)->type;
    const int origGenNum = (origType != xrefEntryCompressed) ? origXRef->getEntry(i)->gen : 0;
    const int newGenNum = (newType != xrefEntryCompressed) ? newXRef->getEntry(i)->gen : 0;

    // Check that DontRewrite entries are freed in full rewrite mode
    if (!forceIncremental && origXRef->getEntry(i)->getFlag(XRefEntry::DontRewrite)) {
      if (newType != xrefEntryFree || origGenNum+1 != newGenNum) {
        fprintf(stderr, "XRef entry %u: DontRewrite entry was not freed correctly\n", i);
        result = false;
      }
      continue; // There's nothing left to check for this entry
    }

    // Compare generation numbers
    // Object num 0 should always have gen 65535 according to specs, but some
    // documents have it set to 0. We always write 65535 in output
    if (i != 0) {
      if (origGenNum != newGenNum) {
        fprintf(stderr, "XRef entry %u: generation numbers differ (%d != %d)\n", i, origGenNum, newGenNum);
        result = false;
        continue;
      }
    } else {
      if (newGenNum != 65535) {
        fprintf(stderr, "XRef entry %u: generation number was expected to be 65535 (%d != 65535)\n", i, newGenNum);
        result = false;
        continue;
      }
    }

    // Compare object flags. A failure shows that there's some error in XRef::scanSpecialFlags()
    if (origXRef->getEntry(i)->flags != newXRef->getEntry(i)->flags) {
      fprintf(stderr, "XRef entry %u: flags detected by scanSpecialFlags differ (%d != %d)\n", i, origXRef->getEntry(i)->flags, newXRef->getEntry(i)->flags);
      result = false;
    }

    // Check that either both are free or both are in use
    if ((origType == xrefEntryFree) != (newType == xrefEntryFree)) {
      const char *origStatus = (origType == xrefEntryFree) ? "free" : "in use";
      const char *newStatus = (newType == xrefEntryFree) ? "free" : "in use";
      fprintf(stderr, "XRef entry %u: usage status differs (%s != %s)\n", i, origStatus, newStatus);
      result = false;
      continue;
    }

    // Skip free entries
    if (origType == xrefEntryFree) {
      continue;
    }

    // Compare contents
    Object origObj = origXRef->fetch(i, origGenNum);
    Object newObj = newXRef->fetch(i, newGenNum);
    if (!compareObjects(&origObj, &newObj)) {
      fprintf(stderr, "XRef entry %u: contents differ\n", i);
      result = false;
    }
  }

  return result;
}