#pragma once

#include <iostream>
#include <stdlib.h>
#include <string.h>

namespace bufio {

template<size_t bufsiz = 1024>
class IBuf {
public:
	IBuf(std::istream &is): is_(is) {}

	char get();
	int peek();
	int peek2();

private:
	std::istream &is_;
	char buf_[bufsiz];
	size_t idx_ = 0;
	size_t len_ = 0;
};

template<size_t bufsiz = 1024>
class OBuf {
public:
	OBuf(std::ostream &os): os_(os) {}
	~OBuf();

	void put(char ch);
	void put(const char *str, size_t len);
	void put(const char *str) { put(str, strlen(str)); }
	void put(const std::string &str) { put(str.c_str(), str.size()); }

private:
	std::ostream &os_;
	char buf_[bufsiz];
	size_t idx_ = 0;
};

/*
 * IBuf
 */

template<size_t bufsiz>
inline char IBuf<bufsiz>::get() {
	if (idx_ < len_) {
		return buf_[idx_++];
	}

	idx_ = 0;
	is_.read(buf_, sizeof(buf_));
	len_ = is_.gcount();
	if (len_ == 0) {
		return EOF;
	}

	return buf_[idx_++];
}

template<size_t bufsiz>
inline int IBuf<bufsiz>::peek() {
	if (idx_ < len_) {
		return buf_[idx_];
	} else {
		return is_.peek();
	}
}

template<size_t bufsiz>
inline int IBuf<bufsiz>::peek2() {
	if (idx_ + 1 < len_) { 
		return buf_[idx_ + 1];
	} else if (idx_ < len_) {
		return is_.peek();
	} else {
		is_.get();
		int ch = is_.peek();
		is_.unget();
		return ch;
	}
}

/*
 * OBuf
 */

template<size_t bufsiz>
inline OBuf<bufsiz>::~OBuf() {
	if (idx_ > 0) {
		os_.write(buf_, idx_);
	}
}

template<size_t bufsiz>
inline void OBuf<bufsiz>::put(char ch) {
	buf_[idx_++] = ch;

	if (idx_ == sizeof(buf_)) {
		os_.write(buf_, sizeof(buf_));
		idx_ = 0;
	}
}

template<size_t bufsiz>
inline void OBuf<bufsiz>::put(const char *str, size_t len) {
	size_t w = sizeof(buf_) - idx_ - 1;
	if (w > len) {
		w = len;
	}

	memcpy(buf_ + idx_, str, w);
	if (len - w > 0) {
		os_.write(buf_, idx_ + w);
		os_.write(str + w, len - w);
		idx_ = 0;
	} else {
		idx_ += w;
	}
}

}