/* -*- mode: C++; c-basic-offset: 2; indent-tabs-mode: nil -*- */

/*
 *  Main authors:
 *     Guido Tack <guido.tack@monash.edu>
 *     Gleb Belov <gleb.belov@monash.edu>
 */

/* 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/. */

#ifdef _MSC_VER
#define _CRT_SECURE_NO_WARNINGS
#endif

#include <minizinc/solns2out.hh>
#include <minizinc/solver.hh>

#include <fstream>
#include <utility>

using namespace std;
using namespace MiniZinc;

void Solns2Out::printHelp(ostream& os) {
  os << "Solution output options:" << std::endl
     << "  --ozn-file <file>\n    Read output specification from ozn file." << std::endl
     << "  -o <file>, --output-to-file <file>\n    Filename for generated output." << std::endl
     << "  -i <n>, --ignore-lines <n>, --ignore-leading-lines <n>\n    Ignore the first <n> lines "
        "in the FlatZinc solution stream."
     << std::endl
     << "  --soln-sep <s>, --soln-separator <s>, --solution-separator <s>\n    Specify the string "
        "printed after each solution (as a separate line).\n    The default is to use the same as "
        "FlatZinc, \"----------\"."
     << std::endl
     << "  --soln-comma <s>, --solution-comma <s>\n    Specify the string used to separate "
        "solutions.\n    The default is the empty string."
     << std::endl
     << "  --unsat-msg (--unsatisfiable-msg), --unbounded-msg, --unsatorunbnd-msg,\n"
        "        --unknown-msg, --error-msg, --search-complete-msg <msg>\n"
        "    Specify solution status messages. The defaults:\n"
        "    \"=====UNSATISFIABLE=====\", \"=====UNSATorUNBOUNDED=====\", "
        "\"=====UNBOUNDED=====\",\n"
        "    \"=====UNKNOWN=====\", \"=====ERROR=====\", \"==========\", respectively."
     << std::endl
     << "  --non-unique\n    Allow duplicate solutions.\n"
     << "  -c, --canonicalize\n    Canonicalize the output solution stream (i.e., buffer and "
        "sort).\n"
     << "  --output-non-canonical <file>\n    Non-buffered solution output file in case of "
        "canonicalization.\n"
     << "  --output-raw <file>\n    File to dump the solver's raw output (not for hard-linked "
        "solvers)\n"
     // Unclear how to exit then:
     //   << "  --number-output <n>\n    Maximal number of different solutions printed." <<
     //   std::endl
     << "  --no-output-comments\n    Do not print comments in the FlatZinc solution stream."
     << std::endl
     << "  --output-time\n    Print timing information in the FlatZinc solution stream."
     << std::endl
     << "  --no-flush-output\n    Don't flush output stream after every line." << std::endl;
}

bool Solns2Out::processOption(int& i, std::vector<std::string>& argv,
                              const std::string& workingDir) {
  CLOParser cop(i, argv);
  std::string buffer;
  if (cop.getOption("--ozn-file", &buffer)) {
    initFromOzn(FileUtils::file_path(buffer, workingDir));
  } else if (cop.getOption("-o --output-to-file", &buffer)) {
    opt.flagOutputFile = buffer;
  } else if (cop.getOption("--no-flush-output")) {
    opt.flagOutputFlush = false;
  } else if (cop.getOption("--no-output-comments")) {
    opt.flagOutputComments = false;
  } else if (cop.getOption("-i --ignore-lines --ignore-leading-lines",
                           &opt.flagIgnoreLines)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--output-time")) {
    opt.flagOutputTime = true;
  } else if (cop.getOption("--soln-sep --soln-separator --solution-separator",
                           &opt.solutionSeparator)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--soln-comma --solution-comma",
                           &opt.solutionComma)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--unsat-msg --unsatisfiable-msg",
                           &opt.unsatisfiableMsg)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--unbounded-msg",
                           &opt.unboundedMsg)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--unsatorunbnd-msg",
                           &opt.unsatorunbndMsg)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--unknown-msg", &opt.unknownMsg)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--error-msg", &opt.errorMsg)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--search-complete-msg",
                           &opt.searchCompleteMsg)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--unique")) {
    opt.flagUnique = true;
  } else if (cop.getOption("--non-unique")) {
    opt.flagUnique = false;
  } else if (cop.getOption("-c --canonicalize")) {
    opt.flagCanonicalize = true;
  } else if (cop.getOption("--output-non-canonical --output-non-canon",
                           &opt.flagOutputNoncanonical)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (cop.getOption("--output-raw",
                           &opt.flagOutputRaw)) {  // NOLINT: Allow repeated empty if
    // Parsed by reference
  } else if (opt.flagStandaloneSolns2Out) {
    std::string oznfile(argv[i]);
    if (oznfile.length() <= 4) {
      return false;
    }
    size_t last_dot = oznfile.find_last_of('.');
    if (last_dot == string::npos) {
      return false;
    }
    std::string extension = oznfile.substr(last_dot, string::npos);
    if (extension == ".ozn") {
      initFromOzn(oznfile);
      return true;
    }
    return false;
  } else {
    return false;
  }
  return true;
}

bool Solns2Out::initFromEnv(Env* pE) {
  assert(pE);
  _env = pE;
  _includePaths.push_back(_stdlibDir + "/std/");
  init();
  return true;
}

void Solns2Out::initFromOzn(const std::string& filename) {
  std::vector<string> filenames(1, filename);

  _includePaths.push_back(_stdlibDir + "/std/");

  for (auto& includePath : _includePaths) {
    if (!FileUtils::directory_exists(includePath)) {
      std::cerr << "solns2out: cannot access include directory " << includePath << "\n";
      std::exit(EXIT_FAILURE);
    }
  }

  {
    _env = new Env();
    std::stringstream errstream;
    _outputModel = parse(*_env, filenames, std::vector<std::string>(), "", "", _includePaths, {},
                         false, false, false, false, errstream);
    if (_outputModel != nullptr) {
      std::vector<TypeError> typeErrors;
      _env->model(_outputModel);
      MZN_ASSERT_HARD_MSG(_env, "solns2out: could not allocate Env");
      _envGuard.reset(_env);
      MiniZinc::typecheck(*_env, _outputModel, typeErrors, false, false);
      MiniZinc::register_builtins(*_env);
      _env->model()->checkFnValid(_env->envi(), typeErrors);
      _env->envi().swapOutput();
      init();
    } else {
      throw Error(errstream.str());
    }
  }
}

Solns2Out::DE& Solns2Out::findOutputVar(const ASTString& name) {
  declNewOutput();
  auto it = _declmap.find(name);
  MZN_ASSERT_HARD_MSG(_declmap.end() != it, "solns2out_base: unexpected id in output: " << name);
  return it->second;
}

void Solns2Out::restoreDefaults() {
  for (auto& i : *getModel()) {
    if (auto* vdi = i->dynamicCast<VarDeclI>()) {
      if (vdi->e()->id()->idn() != -1 || (vdi->e()->id()->v() != "_mzn_solution_checker" &&
                                          vdi->e()->id()->v() != "_mzn_stats_checker" &&
                                          vdi->e()->id()->v() != "_mzn_assumption_map")) {
        GCLock lock;
        auto& de = findOutputVar(vdi->e()->id()->str());
        vdi->e()->e(de.second());
        vdi->e()->evaluated(false);
      }
    }
  }
  _fNewSol2Print = false;
}

void Solns2Out::parseAssignments(string& solution) {
  unique_ptr<Model> sm(parse_from_string(*_env, solution, "solution received from solver",
                                         _includePaths, false, true, false, false, _log));
  if (sm == nullptr) {
    throw Error("solns2out_base: could not parse solution");
  }
  solution = "";
  for (unsigned int i = 0; i < sm->size(); i++) {
    if (auto* ai = (*sm)[i]->dynamicCast<AssignI>()) {
      auto& de = findOutputVar(ai->id());
      if (!Expression::isa<BoolLit>(ai->e()) && !Expression::isa<IntLit>(ai->e()) &&
          !Expression::isa<FloatLit>(ai->e())) {
        Type de_t = de.first->type();
        de_t.cv(false);
        Expression::type(ai->e(), de_t);
      }
      ai->decl(de.first);
      typecheck(*_env, getModel(), ai);
      if (Call* c = Expression::dynamicCast<Call>(ai->e())) {
        // This is an arrayXd call, make sure we get the right builtin
        assert(Expression::isa<ArrayLit>(c->arg(c->argCount() - 1)));
        for (unsigned int i = 0; i < c->argCount(); i++) {
          Expression::type(c->arg(i), Type::parsetint());
        }
        Expression::type(c->arg(c->argCount() - 1), de.first->type());
        c->decl(getModel()->matchFn(_env->envi(), c, false));
      }
      de.first->e(ai->e());
    }
  }
  declNewOutput();
}

void Solns2Out::declNewOutput() {
  _fNewSol2Print = true;
  status = SolverInstance::SAT;
}

void Solns2Out::printSolution(std::istream& sol, std::ostream& os, bool outputTime) {
  if (opt.flagEncapsulateJSON) {
    os << "{\"type\": \"solution\", ";
    std::string line;
    while (std::getline(sol, line)) {
      // Remove line breaks from JSON object
      os << line;
    }
    if (outputTime) {
      os << ", \"time\": " << _starttime.ms().count();
    }
    os << "}\n";
  } else {
    os << sol.rdbuf();
    os.clear();  // Writing sol.rdbuf() sets the fail bit if the solution was empty, so clear it
    if (outputTime) {
      os << "% time elapsed: " << _starttime.stoptime() << "\n";
    }
    if (!opt.solutionSeparator.empty()) {
      os << opt.solutionSeparator << '\n';
    }
  }
  if (opt.flagOutputFlush) {
    os.flush();
  }
}

void Solns2Out::parseStatistics(const std::string& stats, std::ostream& os) {
  unique_ptr<Model> sm(parse_from_string(*_env, stats, "statistics received from solver",
                                         _includePaths, false, true, false, false, _log));
  if (sm == nullptr) {
    throw Error("solns2out_base: could not parse statistics");
  }
  StatisticsStream ss(os, opt.flagEncapsulateJSON);
  for (unsigned int i = 0; i < sm->size(); i++) {
    if (auto* ai = (*sm)[i]->dynamicCast<AssignI>()) {
      ss.add(ai->id().c_str(), ai->e());
    }
  }
}

bool Solns2Out::evalOutput() {
  if (!_fNewSol2Print) {
    return true;
  }

  stringstream oss;
  if (!_checkerModel.empty()) {
    auto& checkerStream = _env->envi().checkerOutput;
    checkerStream.clear();
    checkerStream.str("");
    checkSolution(checkerStream);
  }
  if (!evalOutputInternal(oss)) {
    return false;
  }
  bool fNew = true;
  if (opt.flagUnique || opt.flagCanonicalize) {
    auto res = _sSolsCanon.insert(oss.str());
    if (!res.second) {  // repeated solution
      fNew = false;
    }
  }
  if (fNew) {
    {
      auto& checkerStream = _env->envi().checkerOutput;
      checkerStream.flush();
      std::string line;
      if (std::getline(checkerStream, line)) {
        if (opt.flagEncapsulateJSON) {
          std::string last = line;
          _os << "{\"type\": \"checker\", \"messages\": [" << line;
          while (std::getline(checkerStream, line)) {
            last = line;
            _os << ", " << line;
          }
          _os << "]";
          if (last.size() > 19 && last.substr(0, 19) == "{\"type\": \"solution\"") {
            // For backwards compatibility, output solution message as part of this message
            _os << last.substr(19) << "\n";
          } else {
            _os << "}\n";
          }
        } else {
          _os << "% Solution checker report:\n";
          _os << "% " << line << "\n";
          while (std::getline(checkerStream, line)) {
            _os << "% " << line << "\n";
          }
        }
      }
    }
    ++stats.nSolns;
    if (opt.flagCanonicalize) {
      if (_outStreamNonCanon != nullptr && _outStreamNonCanon->good()) {
        printSolution(oss, *_outStreamNonCanon, opt.flagOutputTime);
      }
      if (opt.flagOutputTime) {
        // Print time as we get solutions
        if (opt.flagEncapsulateJSON) {
          getOutput() << "{\"type\": \"time\", \"time\": " << _starttime.ms().count() << "}\n";
        } else {
          getOutput() << "% time elapsed: " << _starttime.stoptime() << "\n";
        }
      }
    } else {
      if ((!opt.solutionComma.empty()) && stats.nSolns > 1) {
        getOutput() << opt.solutionComma << '\n';
      }

      printSolution(oss, getOutput(), opt.flagOutputTime);
    }
  }
  restoreDefaults();  // cleans data. evalOutput() should not be called again w/o assigning new
                      // data.
  return true;
}

// NOLINTNEXTLINE(readability-convert-member-functions-to-static): Appears static without Gecode
void Solns2Out::checkSolution(std::ostream& oss) const {
#ifdef HAS_GECODE

  std::ostringstream checker;
  checker << _checkerModel;
  Printer p(checker, 0, false, &getEnv()->envi());
  {
    GCLock lock;
    for (auto& i : *getModel()) {
      if (auto* vdi = i->dynamicCast<VarDeclI>()) {
        if (Expression::ann(vdi->e()).contains(Constants::constants().ann.mzn_check_var)) {
          checker << vdi->e()->id()->str() << " = ";
          Expression* e = eval_par(getEnv()->envi(), vdi->e()->e());
          auto* al = Expression::dynamicCast<ArrayLit>(e);
          std::vector<Id*> enumids;
          if (Call* cev = Expression::ann(vdi->e()).getCall(
                  Constants::constants().ann.mzn_check_enum_var)) {
            auto* enumIdsAl = eval_array_lit(getEnv()->envi(), cev->arg(0));
            for (unsigned int j = 0; j < enumIdsAl->size(); j++) {
              enumids.push_back(Expression::dynamicCast<Id>((*enumIdsAl)[j]));
            }
          }
          bool isArray = al != nullptr && al->type().dim() > 0;
          if (isArray) {
            checker << "array" << al->dims() << "d(";
            for (unsigned int i = 0; i < al->dims(); i++) {
              if (!enumids.empty() && enumids[i] != nullptr) {
                checker << "to_enum(" << *enumids[i] << ",";
              }
              checker << al->min(i) << ".." << al->max(i);
              if (!enumids.empty() && enumids[i] != nullptr) {
                checker << ")";
              }
              checker << ",";
            }
          }
          if (!enumids.empty() && enumids.back() != nullptr) {
            checker << "to_enum(" << *enumids.back() << "," << *e << ")";
          } else {
            p.print(e);
          }
          if (isArray) {
            checker << ")";
          }
          checker << ";\n";
        }
      }
    }
  }

  MznSolver slv(oss, _log, _starttime);
  slv.s2out.opt.solutionSeparator = "";
  slv.s2out.opt.searchCompleteMsg = "";
  try {
    slv.run(opt.checkerArgs, checker.str(), "minizinc", "checker.mzc");
  } catch (const Exception& e) {
    if (opt.flagEncapsulateJSON) {
      e.json(std::cout);
    } else {
      e.print(oss);
    }
  } catch (const exception& e) {
    oss << e.what() << std::endl;
  } catch (...) {
    oss << "  UNKNOWN EXCEPTION." << std::endl;
  }

#else
  Warning w("solution checking not supported (need built-in Gecode)");
  if (opt.flagEncapsulateJSON) {
    w.json(oss, false);
  } else {
    w.print(oss, false);
  }
#endif
}

// NOLINTNEXTLINE(readability-convert-member-functions-to-static): Appears static without Gecode
void Solns2Out::checkStatistics(std::ostream& oss) {
#ifdef HAS_GECODE

  std::ostringstream checker;
  checker << _statisticsCheckerModel;
  checker << "mzn_stats_failures = " << stats.nFails << ";\n";
  checker << "mzn_stats_solutions = " << stats.nSolns << ";\n";
  checker << "mzn_stats_nodes = " << stats.nNodes << ";\n";
  checker << "mzn_stats_time = " << _starttime.ms().count() << ";\n";

  MznSolver slv(oss, oss, _starttime);
  slv.s2out.opt.solutionSeparator = "";
  try {
    slv.run(opt.checkerArgs, checker.str(), "minizinc", "checker.mzc");
  } catch (const LocationException& e) {
    oss << e.loc() << ":" << std::endl;
    oss << e.what() << ": " << e.msg() << std::endl;
  } catch (const Exception& e) {
    std::string what = e.what();
    oss << what << (what.empty() ? "" : ": ") << e.msg() << std::endl;
  } catch (const exception& e) {
    oss << e.what() << std::endl;
  } catch (...) {
    oss << "  UNKNOWN EXCEPTION." << std::endl;
  }

#else
  oss << "% statistics checking not supported (need built-in Gecode)" << std::endl;
#endif
}

bool Solns2Out::evalOutputInternal(ostream& fout) {
  if (nullptr != _outputExpr) {
    _env->envi().evalOutput(fout, _log);
  }
  return true;
}

bool Solns2Out::evalStatus(SolverInstance::Status status) {
  if (opt.flagCanonicalize) {
    evalOutputFinalInternal(opt.flagOutputFlush);
  }
  evalStatusMsg(status);
  fStatusPrinted = true;
  return true;
}

bool Solns2Out::evalOutputFinalInternal(bool /*b*/) {
  /// Print the canonical list
  for (const auto& sol : _sSolsCanon) {
    if ((!opt.solutionComma.empty()) && &sol != &*_sSolsCanon.begin()) {
      getOutput() << opt.solutionComma << '\n';
    }
    std::stringstream oss;
    oss << sol;
    printSolution(oss, getOutput(), false);
  }
  return true;
}

bool Solns2Out::evalStatusMsg(SolverInstance::Status status) {
  getOutput() << comments;
  comments = "";

  Solns2Out::status = status;
  std::string label;
  switch (status) {
    case SolverInstance::SAT:
      if (opt.flagOutputFlush) {
        getOutput().flush();
      }
      return true;
    case SolverInstance::OPT: {
      if (opt.flagEncapsulateJSON && !opt.searchCompleteMsg.empty()) {
        bool sat = getEnv()->flat()->solveItem()->st() == SolveI::ST_SAT;
        label = sat ? "ALL_SOLUTIONS" : "OPTIMAL_SOLUTION";
      } else {
        label = opt.searchCompleteMsg;
      }
      break;
    }
    case SolverInstance::UNSAT:
      label = opt.flagEncapsulateJSON ? "UNSATISFIABLE" : opt.unsatisfiableMsg;
      break;
    case SolverInstance::UNBND:
      label = opt.flagEncapsulateJSON ? "UNBOUNDED" : opt.unboundedMsg;
      break;
    case SolverInstance::UNSATorUNBND:
      label = opt.flagEncapsulateJSON ? "UNSAT_OR_UNBOUNDED" : opt.unsatorunbndMsg;
      break;
    case SolverInstance::UNKNOWN:
      label = opt.flagEncapsulateJSON ? "UNKNOWN" : opt.unknownMsg;
      break;
    case SolverInstance::ERROR:
      label = opt.flagEncapsulateJSON ? "ERROR" : opt.errorMsg;
      break;
    case SolverInstance::NONE:
      label = "";
      break;
    default:
      assert(false);
  }
  if (opt.flagEncapsulateJSON) {
    if (label.empty()) {
      if (opt.flagOutputTime) {
        getOutput() << "{\"type\": \"time\", \"time\": " << _starttime.ms().count() << "}\n";
      }
    } else {
      getOutput() << "{\"type\": \"status\", \"status\": \"" << label << "\"";
      if (opt.flagOutputTime) {
        getOutput() << ", \"time\": " << _starttime.ms().count();
      }
      getOutput() << "}\n";
    }
  } else {
    if (!label.empty()) {
      getOutput() << label << '\n';
    }
    if (opt.flagOutputTime) {
      getOutput() << "% time elapsed: " << _starttime.stoptime() << "\n";
    }
  }

  if (opt.flagOutputFlush) {
    getOutput().flush();
  }
  return true;
}

void Solns2Out::init() {
  _declmap.clear();
  _assumptionMap.clear();
  for (auto& i : *getModel()) {
    if (auto* oi = i->dynamicCast<OutputI>()) {
      _outputExpr = oi->e();
    } else if (auto* vdi = i->dynamicCast<VarDeclI>()) {
      if (vdi->e()->id()->idn() == -1 && vdi->e()->id()->v() == "_mzn_solution_checker") {
        _checkerModel = eval_string(getEnv()->envi(), vdi->e()->e());
        if (!_checkerModel.empty() && _checkerModel[0] == '@') {
          _checkerModel = FileUtils::decode_base64(_checkerModel);
          FileUtils::inflate_string(_checkerModel);
        }
      } else if (vdi->e()->id()->idn() == -1 && vdi->e()->id()->v() == "_mzn_stats_checker") {
        _statisticsCheckerModel = eval_string(getEnv()->envi(), vdi->e()->e());
        if (!_statisticsCheckerModel.empty() && _statisticsCheckerModel[0] == '@') {
          _statisticsCheckerModel = FileUtils::decode_base64(_statisticsCheckerModel);
          FileUtils::inflate_string(_statisticsCheckerModel);
        }
      } else if (vdi->e()->id()->idn() == -1 && vdi->e()->id()->v() == "_mzn_assumption_map") {
        // A `list of tuple(string, string)` of `(variable name, expression)` pairs.
        GCLock lock;
        ArrayLit* al = eval_array_lit(getEnv()->envi(), vdi->e()->e());
        for (unsigned int k = 0; k < al->size(); k++) {
          ArrayLit* tuple = eval_array_lit(getEnv()->envi(), (*al)[k]);
          _assumptionMap[eval_string(getEnv()->envi(), (*tuple)[0])] =
              eval_string(getEnv()->envi(), (*tuple)[1]);
        }
      } else {
        _declmap.insert(make_pair(vdi->e()->id()->str(), DE(vdi->e(), vdi->e()->e())));
      }
    }
  }

  /// Main output file
  if (nullptr == _outStream) {
    if (!opt.flagOutputFile.empty()) {
      _outStream.reset(new ofstream(FILE_PATH(opt.flagOutputFile)));
      MZN_ASSERT_HARD_MSG(_outStream.get(),
                          "solns2out_base: could not allocate stream object for file output into "
                              << opt.flagOutputFile);
      check_io_status(_outStream->good(), opt.flagOutputFile);
    }
  }
  /// Non-canonical output
  if (opt.flagCanonicalize && (!opt.flagOutputNoncanonical.empty())) {
    _outStreamNonCanon.reset(new ofstream(FILE_PATH(opt.flagOutputNoncanonical)));
    MZN_ASSERT_HARD_MSG(_outStreamNonCanon.get(),
                        "solns2out_base: could not allocate stream object for non-canon output");
    check_io_status(_outStreamNonCanon->good(), opt.flagOutputNoncanonical, false);
  }
  /// Raw output
  if (!opt.flagOutputRaw.empty()) {
    _outStreamRaw.reset(new ofstream(FILE_PATH(opt.flagOutputRaw)));
    MZN_ASSERT_HARD_MSG(_outStreamRaw.get(),
                        "solns2out_base: could not allocate stream object for raw output");
    check_io_status(_outStreamRaw->good(), opt.flagOutputRaw, false);
  }
  /// Assume all options are set before
  nLinesIgnore = opt.flagIgnoreLines;
}

Solns2Out::Solns2Out(std::ostream& os0, std::ostream& log0, std::string stdlibDir0)
    : _os(os0), _log(log0), _stdlibDir(std::move(stdlibDir0)) {}

Solns2Out::~Solns2Out() {
  getOutput() << comments;
  if (opt.flagOutputFlush) {
    getOutput() << flush;
  }
}

ostream& Solns2Out::getOutput() {
  return (((_outStream != nullptr) && _outStream->good()) ? *_outStream : _os);
}

ostream& Solns2Out::getLog() { return _log; }

bool Solns2Out::feedRawDataChunk(const char* data) {
  istringstream solstream(data);
  while (solstream.good()) {
    string line;
    getline(solstream, line);
    if (!_linePart.empty()) {
      line = _linePart + line;
      _linePart.clear();
    }
    if (solstream.eof()) {  // wait next chunk
      _linePart = line;
      break;  // to get to raw output
    }
    if (!line.empty()) {
      if ('\r' == line.back()) {
        line.pop_back();  // For WIN files
      }
    }
    if (nLinesIgnore > 0) {
      --nLinesIgnore;
      continue;
    }
    if (_mapInputStatus.empty()) {
      createInputMap();
    }
    auto it = _mapInputStatus.find(line);
    if (_mapInputStatus.end() != it) {
      flushStatistics(getOutput());
      if (SolverInstance::SAT == it->second) {
        parseAssignments(solution);
        evalOutput();
      } else {
        evalStatus(it->second);
      }
    } else {
      solution += line + '\n';
      if (opt.flagOutputComments) {
        std::istringstream iss(line);
        char c = '_';
        iss >> skipws >> c;
        if (iss.good() && '%' == c) {
          bool is_statistic = line.substr(0, 13) == "%%%mzn-stat: ";
          bool is_core = line.substr(0, 13) == "%%%mzn-core: ";
          std::ostringstream message;
          if (is_core) {
            // Resolve the FlatZinc variable names in the unsatisfiable core back to the
            // original `assume` expressions and report them.
            std::vector<std::string> resolved = resolveAssumptionCore(line.substr(13));
            if (opt.flagEncapsulateJSON) {
              message << "{\"type\": \"core\", \"core\": [";
              for (size_t i = 0; i < resolved.size(); ++i) {
                message << (i == 0 ? "" : ", ") << "\"" << Printer::escapeStringLit(resolved[i])
                        << "\"";
              }
              message << "]}\n";
            } else {
              message << "%%%mzn-core: [";
              for (size_t i = 0; i < resolved.size(); ++i) {
                if (i > 0) {
                  message << ", ";
                }
                message << resolved[i];
              }
              message << "]\n";
            }
          } else if (opt.flagEncapsulateJSON) {
            if (is_statistic) {
              _stats += line.substr(13) + ";";
            } else if (line == "%%%mzn-stat-end") {
              parseStatistics(_stats, message);
              _stats.clear();
            } else {
              message << "{\"type\": \"comment\", \"comment\": \"" << Printer::escapeStringLit(line)
                      << "\\n\"}\n";
            }
          } else {
            message << line << '\n';
          }
          getOutput() << message.str();
          if (opt.flagOutputFlush) {
            getOutput().flush();
          }
          if (_outStreamNonCanon != nullptr && _outStreamNonCanon->good()) {
            (*_outStreamNonCanon) << message.str();
            if (opt.flagOutputFlush) {
              _outStreamNonCanon->flush();
            }
          }

          if (is_statistic && line.size() > 13) {
            if (line.substr(13, 6) == "nodes=") {
              std::istringstream iss(line.substr(19));
              int n_nodes;
              iss >> n_nodes;
              stats.nNodes = n_nodes;
            } else if (line.substr(13, 9) == "failures=") {
              std::istringstream iss(line.substr(22));
              int n_failures;
              iss >> n_failures;
              stats.nFails = n_failures;
            }
          }
        }
      }
    }
  }
  if (_outStreamRaw != nullptr) {
    *_outStreamRaw << data;
    if (opt.flagOutputFlush) {
      _outStreamRaw->flush();
    }
  }
  return true;
}

namespace {
/// True if \a s could still grow into \a marker (i.e. it is a prefix of, or
/// equal to, the marker). Used to decide whether an incomplete line must be
/// held back because it might turn out to be a marker.
bool is_marker_prefix(const std::string& s, const std::string& marker) {
  return s.size() <= marker.size() && marker.compare(0, s.size(), s) == 0;
}
}  // namespace

bool Solns2OutInteractive::processCompleteLine(const std::string& line) {
  // Anything before _echoed is a prompt prefix already shown to the user (a
  // prompt has no trailing newline, so the marker that follows it ends up glued
  // onto the same line in the solver's output stream). Only the not-yet-echoed
  // remainder is eligible to be a marker, or fresh verbatim/solution text.
  std::string rest = line.substr(_echoed);
  std::string marker = rest;
  if (!marker.empty() && '\r' == marker.back()) {
    marker.pop_back();  // tolerate CRLF when matching markers
  }
  if (!_inSolution) {
    if (marker == _beginMarker) {
      _inSolution = true;  // consume the marker (any prompt prefix stays shown)
      return true;
    }
    // Verbatim solver chatter: echo whatever has not been echoed yet.
    _inner.getOutput() << rest << '\n';
    _inner.getOutput().flush();
    return true;
  }
  if (marker == _endMarker) {
    _inSolution = false;  // consume the marker, back to verbatim mode
    return true;
  }
  // FlatZinc solution text: hand it to the real output processor.
  std::string forwarded = rest + '\n';
  return _inner.feedRawDataChunk(forwarded.c_str());
}

bool Solns2OutInteractive::feedRawDataChunk(const char* data) {
  // Split the (arbitrarily chunked) raw stream into whole lines, mirroring the
  // partial-line handling in Solns2Out::feedRawDataChunk. Marker detection is
  // line-based, but to keep an interactive prompt (text with no trailing
  // newline) in sync we echo the incomplete trailing line eagerly, holding it
  // back only while it could still become the begin marker.
  std::istringstream solstream(data);
  bool ret = true;
  while (solstream.good()) {
    std::string line;
    std::getline(solstream, line);
    line = _linePart + line;  // prepend any carried-over partial line
    if (solstream.eof()) {    // incomplete line: carry it over to the next chunk
      _linePart = line;
      if (!_inSolution) {
        std::string remainder = _linePart.substr(_echoed);
        if (!remainder.empty() && !is_marker_prefix(remainder, _beginMarker)) {
          // The remainder cannot (yet) be the begin marker: flush it now so a
          // prompt without a trailing newline appears immediately.
          _inner.getOutput() << remainder << std::flush;
          _echoed = _linePart.size();
        }
      }
      break;
    }
    // `line` is now a complete line (without its terminating '\n').
    if (!processCompleteLine(line)) {
      ret = false;
    }
    _linePart.clear();
    _echoed = 0;
  }
  return ret;
}

std::vector<std::string> Solns2Out::resolveAssumptionCore(const std::string& payload) const {
  std::vector<std::string> result;
  // Trim and strip a surrounding pair of square brackets.
  size_t begin = payload.find_first_not_of(" \t\r\n");
  size_t end = payload.find_last_not_of(" \t\r\n");
  if (begin == std::string::npos) {
    return result;
  }
  std::string body = payload.substr(begin, end - begin + 1);
  if (!body.empty() && body.front() == '[' && body.back() == ']') {
    body = body.substr(1, body.size() - 2);
  }
  auto is_ident_start = [](char ch) {
    return (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || ch == '_';
  };
  auto is_ident_char = [&](char ch) { return is_ident_start(ch) || (ch >= '0' && ch <= '9'); };
  // Split on commas (the solver emits FlatZinc names/literals, which contain no commas) and
  // resolve each. A resolvable entry is either a known variable name (a Boolean assumption) or a
  // comparison `<objective> <op> <const>`; in both cases the variable is the leading identifier.
  // Only that leading identifier is substituted: any entry whose leading identifier is unknown is
  // kept verbatim.
  size_t pos = 0;
  while (pos <= body.size()) {
    size_t comma = body.find(',', pos);
    std::string token =
        body.substr(pos, comma == std::string::npos ? std::string::npos : comma - pos);
    size_t nb = token.find_first_not_of(" \t\r\n");
    size_t ne = token.find_last_not_of(" \t\r\n");
    if (nb != std::string::npos) {
      token = token.substr(nb, ne - nb + 1);
      // Extract the leading identifier (if any).
      size_t idEnd = 0;
      if (is_ident_start(token[0])) {
        idEnd = 1;
        while (idEnd < token.size() && is_ident_char(token[idEnd])) {
          idEnd++;
        }
      }
      auto it = idEnd > 0 ? _assumptionMap.find(token.substr(0, idEnd)) : _assumptionMap.end();
      if (it == _assumptionMap.end()) {
        // Unknown entry: keep it verbatim.
        result.push_back(token);
      } else if (idEnd == token.size()) {
        // The whole entry is a known variable: report its expression as-is.
        result.push_back(it->second);
      } else {
        // A known variable followed by a comparison (e.g. an objective bound): substitute the
        // variable, parenthesising it so the surrounding operator still reads correctly.
        result.push_back("(" + it->second + ")" + token.substr(idEnd));
      }
    }
    if (comma == std::string::npos) {
      break;
    }
    pos = comma + 1;
  }
  return result;
}

void Solns2Out::createInputMap() {
  _mapInputStatus[opt.searchCompleteMsgDef] = SolverInstance::OPT;
  _mapInputStatus[opt.solutionSeparatorDef] = SolverInstance::SAT;
  _mapInputStatus[opt.unsatisfiableMsgDef] = SolverInstance::UNSAT;
  _mapInputStatus[opt.unboundedMsgDef] = SolverInstance::UNBND;
  _mapInputStatus[opt.unsatorunbndMsgDef] = SolverInstance::UNSATorUNBND;
  _mapInputStatus[opt.unknownMsgDef] = SolverInstance::UNKNOWN;
  _mapInputStatus[opt.errorMsgDef] = SolverInstance::ERROR;
}

void Solns2Out::flushStatistics(ostream& os) {
  if (!_stats.empty()) {
    // Print any unprocessed statistics
    parseStatistics(_stats, getOutput());
    _stats.clear();
  }
}

void Solns2Out::printStatistics(ostream& os) {
  StatisticsStream ss(os, opt.flagEncapsulateJSON);
  ss.add("nSolutions", stats.nSolns);
  if (!_statisticsCheckerModel.empty()) {
    std::ostringstream oss;
    checkStatistics(oss);
    ss.add("statisticsCheck", oss.str());
  }
}
