#pragma once

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

namespace bufio {

class IStream {
public:
	virtual size_t read(char *buf, size_t maxlen) = 0;
};

class IFStream: public IStream {
public:
	IFStream(const char *path);
	IFStream(const std::string &path): IFStream(path.c_str()) {}

	size_t read(char *buf, size_t maxlen) override;

private:
	std::ifstream is_;
};

class ISStream: public IStream {
public:
	ISStream(std::string str): str_(std::move(str)) {}

	size_t read(char *buf, size_t maxlen) override;

private:
	size_t idx_ = 0;
	const std::string str_;
};

class OStream {
public:
	virtual void write(const char *buf, size_t len) = 0;
};

class OFStream: public OStream {
public:
	OFStream(const char *path);
	OFStream(const std::string &path): OFStream(path.c_str()) {}

	void write(const char *buf, size_t len) override;

private:
	std::ofstream os_;
};

template<typename IS = IStream, size_t bufsiz = 1024>
class IBuf {
public:
	IBuf(IS &is): is_(is) {}

	char get();
	int peek(size_t count = 1);

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

template<typename OS = OStream, size_t bufsiz = 1024>
class OBuf {
public:
	OBuf(OS &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:
	OS &os_;
	char buf_[bufsiz];
	size_t idx_ = 0;
};

/*
 * IFStream
 */

inline IFStream::IFStream(const char *path) {
	is_.exceptions(std::ifstream::badbit);
	is_.open(path);
	if (!is_.good()) {
		throw std::system_error(errno, std::generic_category(), path);
	}
}

inline size_t IFStream::read(char *buf, size_t maxlen) {
	is_.read(buf, maxlen);
	return is_.gcount();
}

/*
 * ISStream
 */

inline size_t ISStream::read(char *buf, size_t maxlen) {
	size_t left = str_.size() - idx_;
	size_t len = maxlen < left ? maxlen : left;
	idx_ += len;
	memcpy(buf, str_.c_str(), len);
	return len;
}

/*
 * OFStream
 */

inline OFStream::OFStream(const char *path) {
	os_.exceptions(std::ofstream::badbit);
	os_.open(path);
	if (!os_.good()) {
		throw std::system_error(errno, std::generic_category(), path);
	}
}

inline void OFStream::write(const char *buf, size_t len) {
	os_.write(buf, len);
}

/*
 * IBuf
 */

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

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

	return buf_[idx_++];
}

template<typename IS, size_t bufsiz>
inline int IBuf<IS, bufsiz>::peek(size_t count) {
	size_t offset = count - 1;
	if (idx_ + offset < len_) {
		return buf_[idx_ + offset];
	} else {
		len_ -= idx_;
		memmove(buf_, buf_ + idx_, len_);
		idx_ = 0;
		len_ += is_.read(buf_ + len_, sizeof(buf_) - len_);
		if (len_ <= offset) {
			return EOF;
		}

		return buf_[idx_ + offset];
	}
}

/*
 * OBuf
 */

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

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

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

template<typename OS, size_t bufsiz>
inline void OBuf<OS, 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;
	}
}

}