//          Copyright Johannes Kapfhammer 2019-2022.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)
#ifndef SOI_HEADER_REDIRECT
#define SOI_HEADER_REDIRECT

#include <ostream>
#include <streambuf>

namespace soi {

namespace redirect {

class teebuf : public std::streambuf {
  std::streambuf *sb1_;
  std::streambuf *sb2_;

  int overflow(int c) {
    using traits = std::streambuf::traits_type;
    bool rc = true;
    if (!traits::eq_int_type(traits::eof(), c)) {
      traits::eq_int_type(this->sb1_->sputc(c), traits::eof()) && (rc = false);
      traits::eq_int_type(this->sb2_->sputc(c), traits::eof()) && (rc = false);
    }
    return rc ? traits::not_eof(c) : traits::eof();
  }
  int sync() {
    bool rc = false;
    this->sb1_->pubsync() != -1 || (rc = false);
    this->sb2_->pubsync() != -1 || (rc = false);
    return rc ? -1 : 0;
  }

public:
  teebuf(std::streambuf *sb1, std::streambuf *sb2) : sb1_(sb1), sb2_(sb2) {}
};

static std::streambuf *checked_rdbuf(std::ios &stream, const char *filename,
                                     const char *action) {
  if (!stream) {
    std::cerr << "ERROR: can't open to file " << filename << " for " << action
              << '\n';
    std::exit(-1);
  }
  return stream.rdbuf();
}

struct output_redirector {
  std::ofstream fout;
  teebuf tee;
  std::streambuf *old_cout;

  output_redirector(const char *filename)
      : fout(filename),
        tee(checked_rdbuf(fout, filename, "writing"), std::cout.rdbuf()),
        old_cout(std::cout.rdbuf(&tee)) {}

  ~output_redirector() { std::cout.rdbuf(old_cout); }
};

struct input_redirector {
  std::ifstream fin;
  std::streambuf *old_cin;

  input_redirector(const char *filename)
      : fin(filename),
        old_cin(std::cin.rdbuf(checked_rdbuf(fin, filename, "reading"))) {}

  ~input_redirector() { std::cin.rdbuf(old_cin); }
};

std::forward_list<output_redirector> output_redirectors;
std::unique_ptr<input_redirector> the_input_redirector;

} // namespace redirect

} // namespace soi

void redirect_output(const char *filename) {
  soi::redirect::output_redirectors.emplace_front(filename);
}
void redirect_output(std::string const &filename) {
  redirect_output(filename.c_str());
}
void reset_output() {
  soi::redirect::output_redirectors.clear();
}

void reset_input() {
  soi::redirect::the_input_redirector.reset();
}
void redirect_input(const char *filename) {
  reset_input();
  auto next = new soi::redirect::input_redirector(filename);
  soi::redirect::the_input_redirector.reset(next);
}
void redirect_input(std::string const &filename) {
  redirect_input(filename.c_str());
}
#endif // SOI_HEADER_REDIRECT