/******************************************************************************
*
* Copyright (C) 1997-2019 by Dimitri van Heesch.
*
* Permission to use, copy, modify, and distribute this software and its
* documentation under the terms of the GNU General Public License is hereby
* granted. No representations are made about the suitability of this software
* for any purpose. It is provided "as is" without express or implied warranty.
* See the GNU General Public License for more details.
*
* Documents produced by Doxygen are derivative works derived from the
* input used in their production; they are not affected by this license.
*
*/

#include <cassert>
#include <cmath>
#include <map>
#include <set>
#include <string>
#include <algorithm>
#include <numeric>
#include <random>
#include "threadpool.h"

#ifdef _MSC_VER
#pragma warning( push )
#pragma warning( disable : 4242 )
#pragma warning( disable : 4244 )
#pragma warning( disable : 4996 )
#pragma warning( disable : 4456 )
#pragma warning( disable : 4805 )
#endif
#if defined(__clang__)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdeprecated-declarations"
#pragma clang diagnostic ignored "-Wshadow"
#endif
#if defined(__GNUC__)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wshadow"
#endif
#include <gunzip.hh>
#if defined(__GNUC__)
#pragma GCC diagnostic pop
#endif
#if defined(__clang__)
#pragma clang diagnostic pop
#endif
#ifdef _MSC_VER
#pragma warning( pop )
#endif

#include "dotrunner.h"
#include "util.h"
#include "portable.h"
#include "dot.h"
#include "message.h"
#include "config.h"
#include "dir.h"
#include "doxygen.h"

// the graphicx LaTeX has a limitation of maximum size of 16384
// To be on the save side we take it a little bit smaller i.e. 150 inch * 72 dpi
// It is anyway hard to view these size of images
#define MAX_LATEX_GRAPH_INCH  150
#define MAX_LATEX_GRAPH_SIZE  (MAX_LATEX_GRAPH_INCH * 72)

//#define DBG(x) printf x
#define DBG(x) do {} while(0)

//-----------------------------------------------------------------------------------------

// since dot silently reproduces the input file when it does not
// support the PNG format, we need to check the result.
static void checkPngResult(const DString &imgName)
{
  FILE *f = Portable::fopen(imgName,"rb");
  if (!f)
  {
    err("Could not read image '{}' generated by dot!\n",imgName);
    return;
  }

  char data[4];
  if (fread(data, 1, 4, f) != 4)
  {
    err("Could not read image '{}' generated by dot!\n",imgName);
    fclose(f);
    return;
  }

  if (!(data[1] == 'P' && data[2] == 'N' && data[3] == 'G'))
  {
    err("Image '{}' produced by dot is not a valid PNG!\n"
        "You should either select a different format "
        "(DOT_IMAGE_FORMAT in the config file) or install a more "
        "recent version of graphviz (1.7+)\n", imgName);
  }

  fclose(f);
}

static bool resetPDFSize(const int width,const int height, const DString &base)
{
  DString tmpName   = base+".tmp";
  DString patchFile = base+".dot";
  Dir thisDir;
  if (!thisDir.rename(patchFile.str(),tmpName.str()))
  {
    err("Failed to rename file {} to {}!\n",patchFile,tmpName);
    return false;
  }
  std::ifstream fi = Portable::openInputStream(tmpName);
  std::ofstream t  = Portable::openOutputStream(patchFile);
  if (!fi.is_open())
  {
    err("problem opening file {} for patching!\n",tmpName);
    thisDir.rename(tmpName.str(),patchFile.str());
    return false;
  }
  if (!t.is_open())
  {
    err("problem opening file {} for patching!\n",patchFile);
    thisDir.rename(tmpName.str(),patchFile.str());
    return false;
  }
  std::string line;
  while (getline(fi,line)) // foreach line
  {
    if (line.find("LATEX_PDF_SIZE") != std::string::npos)
    {
      double scale = (width > height ? width : height)/double(MAX_LATEX_GRAPH_INCH);
      t << "  size=\""<<width/scale << "," <<height/scale << "\";\n";
    }
    else
      t << line << "\n";
  }
  fi.close();
  t.close();
  // remove temporary file
  thisDir.remove(tmpName.str());
  return true;
}

bool DotRunner::readBoundingBox(const DString &fileName,int *width,int *height,bool isEps)
{
  std::ifstream f = Portable::openInputStream(fileName);
  if (!f.is_open())
  {
    err("Failed to open file {} for extracting bounding box\n",fileName);
    return false;
  }

  // read file contents into string 'contents'
  std::stringstream buffer;
  buffer << f.rdbuf();
  std::string contents = buffer.str();

  // start of bounding box marker we are looking for
  const std::string boundingBox = isEps ? "%%PageBoundingBox:" : "/MediaBox [";

  // helper routine to extract the bounding boxes width and height
  auto extractBoundingBox = [&fileName,&boundingBox,&width,&height](const char *s) -> bool
  {
    int x=0, y=0;
    double w=0, h=0;
    if (sscanf(s+boundingBox.length(),"%d %d %lf %lf",&x,&y,&w,&h)==4)
    {
      *width  = static_cast<int>(std::ceil(w));
      *height = static_cast<int>(std::ceil(h));
      return true;
    }
    err("Failed to extract bounding box from generated diagram file {}\n",fileName);
    return false;
  };

  // compressed segment start and end markers
  const std::string streamStart = "stream\n";
  const std::string streamEnd = "\nendstream";

  auto detectDeflateStreamStart = [&streamStart](const char *s)
  {
    size_t len = streamStart.length();
    bool streamOK = strncmp(s,streamStart.c_str(),len)==0;
    if (streamOK) // ASCII marker matches, check stream header bytes as well
    {
      unsigned short header1 = static_cast<unsigned char>(s[len])<<8; // CMF byte
      if (header1) // not end of string
      {
        unsigned short header = (static_cast<unsigned char>(s[len+1])) | header1; // FLG byte
        // check for correct header (see https://www.rfc-editor.org/rfc/rfc1950)
        return ((header&0x8F20)==0x0800) && (header%31)==0;
      }
    }
    return false;
  };

  const size_t l = contents.length();
  size_t i=0;
  while (i<l)
  {
    if (!isEps && contents[i]=='s' && detectDeflateStreamStart(&contents[i]))
    { // compressed stream start
      int col=17;
      i+=streamStart.length();
      const size_t start=i;
      DBG(("---- start stream at offset %08x\n",(int)i));
      while (i<l)
      {
        if (contents[i]=='\n' && strncmp(&contents[i],streamEnd.c_str(),streamEnd.length())==0)
        { // compressed block found in range [start..i]
          DBG(("\n---- end stream at offset %08x\n",(int)i));
          // decompress it into decompressBuf
          std::vector<char> decompressBuf;
          const char *source = &contents[start];
          const size_t sourceLen = i-start;
          size_t sourcePos = 0;
          decompressBuf.reserve(sourceLen*2);
          auto getter = [source,&sourcePos,sourceLen]() -> int {
            return sourcePos<sourceLen ? static_cast<unsigned char>(source[sourcePos++]) : EOF;
          };
          auto putter = [&decompressBuf](const char c) -> int {
            decompressBuf.push_back(c); return c;
          };
          Deflate(getter,putter);
          // convert decompression buffer to string
          std::string s(decompressBuf.begin(), decompressBuf.end());
          DBG(("decompressed_data=[[[\n%s\n]]]\n",s.c_str()));
          // search for bounding box marker
          const size_t idx = s.find(boundingBox);
          if (idx!=std::string::npos) // found bounding box in uncompressed data
          {
            return extractBoundingBox(s.c_str()+idx);
          }
          // continue searching after end stream marker
          i+=streamEnd.length();
          break;
        }
        else // compressed stream character
        {
          if (col>16) { col=0; DBG(("\n%08x: ",static_cast<int>(i))); }
          DBG(("%02x ",static_cast<unsigned char>(contents[i])));
          col++;
          i++;
        }
      }
    }
    else if (((isEps && contents[i]=='%') || (!isEps && contents[i]=='/')) &&
             strncmp(&contents[i],boundingBox.c_str(),boundingBox.length())==0)
    { // uncompressed bounding box
      return extractBoundingBox(&contents[i]);
    }
    else // uncompressed stream character
    {
      i++;
    }
  }
  err("Failed to find bounding box in generated diagram file {}\n",fileName);
  // nothing found
  return false;
}

//---------------------------------------------------------------------------------

static DString getBaseNameOfOutput(const DString &output)
{
  size_t index = output.rfind('.');
  if (index==DString::npos) return output;
  return output.left(index);
}

DotRunner::DotRunner()
  : m_dotExe(Doxygen::verifiedDotPath)
{
}

bool DotRunner::run(const DotJobs &dotJobs)
{
  if (dotJobs.empty()) return true;

  // Group jobs by format, then by directory so we can cd once per group
  std::map<std::string, std::map<std::string, std::vector<const DotJob*>>> byFormatAndDir;
  for (const auto &job : dotJobs)
  {
    byFormatAndDir[job.format.str()][job.absPath.str()].push_back(&job);
  }

  std::mt19937 rng(std::random_device{}());
  bool ok = true;
  size_t prev=0;
  for (const auto &[fmtStr, byDir] : byFormatAndDir)
  {
    DString format = DString(fmtStr);

    for (const auto &[dirStr, jobs] : byDir)
    {
      std::string oldDir = Dir::currentDirPath();
      Dir::setCurrent(dirStr);

      // settings controlling how to distribute the graphs over threads and batches
      const size_t numThreads = static_cast<size_t>(Config_getInt(DOT_NUM_THREADS));
      const size_t batchSize  = static_cast<size_t>(Config_getInt(DOT_BATCH_SIZE));
      const size_t exeLen     = m_dotExe.length() + 1; // "exe " prefix
      const size_t maxArgLen  = 32000-exeLen; // Windows CreateProcess limit is 32767; keep safe margin

      // create a pseudo random ordering in which to process the dot files
      std::vector<size_t> indices(jobs.size());
      std::iota(indices.begin(), indices.end(), 0);
      std::shuffle(indices.begin(), indices.end(), rng);

      // helper to keep track of dot command to run later
      struct CommandArgument
      {
        CommandArgument(const DString &args) : arguments(args) {}
        DString arguments;
        size_t numDotFiles = 0;
        const DotJob *firstJob = nullptr;
      };

      std::vector<CommandArgument> partialCommands;
      std::vector<CommandArgument> finalCommands;

      bool hasImageMap = std::any_of(jobs.begin(),jobs.end(),[](const auto &j) { return j->generateImageMap; });

      // each dot command has a command arguments of the form: -Tformat -O basename1.dot basename2.dot ...
      DString baseArgs = DString("-T") + format;
      if (hasImageMap) // if any image needs a map we generate one for all images
      {
        baseArgs += " -Tcmapx";
      }
      baseArgs += " -O";

      // prepare partial commands for each thread (command is later skipped if numDotFiles==0).
      for (size_t i=0; i<numThreads; i++)
      {
        partialCommands.emplace_back(baseArgs);
      }

      // split the jobs into batches per thread iterating in pseudo random order to fill each batch with a random selection of graphs
      size_t index=0;
      for (size_t i : indices)
      {
        const auto &job = jobs[i];
        DString fileArg = DString(" ") + job->relDotName;
        auto &cmd = partialCommands[index];
        if (cmd.numDotFiles<batchSize && cmd.arguments.length()+fileArg.length()<maxArgLen) // still room in this batch
        {
          cmd.arguments+=fileArg;
          cmd.numDotFiles++;
        }
        else // this batch is full, move to finished commands and start a new one
        {
          finalCommands.push_back(cmd);
          cmd.arguments=baseArgs+fileArg;
          cmd.numDotFiles=1;
        }
        if (cmd.firstJob==nullptr) cmd.firstJob=job;
        index = (index+1)%numThreads;
      }

      // append partial commands to the final commands
      finalCommands.insert(finalCommands.end(),partialCommands.begin(),partialCommands.end());

      // now run the finalCommands.
      if (Config_getInt(DOT_NUM_THREADS)<=1) // no threads to work with
      {
        for (const auto &cmd : finalCommands)
        {
          if (cmd.numDotFiles>0) // check if there are graphs to generate first
          {
            if (cmd.numDotFiles>1) // batch mode
            {
              msg("Running dot for graphs {}-{}/{}\n",prev+1,prev+cmd.numDotFiles,dotJobs.size());
            }
            else // single graph mode
            {
              msg("Running dot for graph {}/{}\n",prev+1,dotJobs.size());
            }
            prev+=cmd.numDotFiles;
            int exitCode;
            if ((exitCode = Portable::system(m_dotExe, cmd.arguments, false)) != 0)
            {
              err_full(cmd.firstJob->srcFile, 1,
                  "Problems running dot: exit code={}, command='{}', dir='{}', arguments='{}'",
                  exitCode, m_dotExe, dirStr, cmd.arguments);
              ok = false;
            }
          }
        }
      }
      else // use multiple threads to run instances of dot in parallel
      {
        ThreadPool workers(numThreads);
        std::vector< std::future<size_t> > results;
        for (auto & cmd: finalCommands)
        {
          if (cmd.numDotFiles>0)
          {
            auto locDirStr = dirStr;
            auto process = [this,cmd,locDirStr]() -> size_t
            {
              int exitCode;
              if ((exitCode = Portable::system(m_dotExe, cmd.arguments, false)) != 0)
              {
                err_full(cmd.firstJob->srcFile, 1,
                    "Problems running dot: exit code={}, command='{}', dir='{}', arguments='{}'",
                    exitCode, m_dotExe, locDirStr, cmd.arguments);
              }
              return cmd.numDotFiles;
            };
            results.emplace_back(workers.queue(process));
          }
        }
        for (auto &f : results)
        {
          size_t numDotFiles = f.get();
          if (numDotFiles>1) // batch mode
          {
            msg("Finished running dot for graphs {}-{}/{}\n",prev+1,prev+numDotFiles,dotJobs.size());
          }
          else // single graph mode
          {
            msg("Finished running dot for graph {}/{}\n",prev+1,dotJobs.size());
          }
          prev+=numDotFiles;
        }
      }

      // Post-process each output file. dot -O appends the format suffix to the
      // full input filename, so the output is absPath + relDotName + "." + format.
      // Rename to remove the .dot infix, producing absPath + baseName + "." + format.
      for (const auto *job : jobs)
      {
        DString base   = job->absPath + getBaseNameOfOutput(job->relDotName);
        DString dotOutput = job->absPath + job->relDotName + "." + format;
        DString output = base + "." + format;
        Dir d;
        if (!d.rename(dotOutput.str(), output.str()))
        {
          err("Failed to rename {} to {}!\n", dotOutput, output);
          ok = false;
          continue;
        }
        if (job->generateImageMap)
        {
          DString dotMapOutput = job->absPath + job->relDotName + ".cmapx";
          DString mapOutput = base + ".map";
          if (!d.rename(dotMapOutput.str(), mapOutput.str()))
          {
            err("Failed to rename {} to {}!\n", dotMapOutput, mapOutput);
            ok = false;
            continue;
          }
        }

        if (format.startsWith("pdf"))
        {
          int width=0, height=0;
          if (!readBoundingBox(output, &width, &height, false))
          {
            ok = false;
            continue;
          }
          if ((width > MAX_LATEX_GRAPH_SIZE) || (height > MAX_LATEX_GRAPH_SIZE))
          {
            if (!resetPDFSize(width, height, base))
            {
              ok = false;
              continue;
            }
            // Re-run dot for just this one file
            DString rerunArgs = DString("-T") + format + " -O \"" + job->relDotName + "\"";
            int exitCode;
            if ((exitCode = Portable::system(m_dotExe, rerunArgs, false)) != 0)
            {
              err_full(job->srcFile, 1,
                       "Problems running dot: exit code={}, command='{}', dir='{}', arguments='{}'",
                       exitCode, m_dotExe, dirStr, rerunArgs);
              ok = false;
            }
            else
            {
              Dir d2;
              if (!d2.rename(dotOutput.str(), output.str()))
              {
                err("Failed to rename {} to {}!\n", dotOutput, output);
                ok = false;
              }
            }
          }
        }
        else if (format.startsWith("png"))
        {
          checkPngResult(output);
        }
      }
      Dir::setCurrent(oldDir);
    }
  }

  // Write .md5 files and clean up .dot files (once per unique dotFile)
  std::set<std::string> processed;
  for (const auto &job : dotJobs)
  {
    if (!processed.insert((job.absPath + job.relDotName).str()).second) continue;

    if (!job.md5Hash.empty())
    {
      DString md5Name = job.absPath + getBaseNameOfOutput(job.relDotName) + ".md5";
      FILE *f = Portable::fopen(md5Name, "w");
      if (f)
      {
        fwrite(job.md5Hash.data(), 1, 32, f);
        fclose(f);
      }
    }

    if (Config_getBool(DOT_CLEANUP))
    {
      Portable::unlink(job.absPath + job.relDotName);
    }
  }

  return ok;
}
