#ifndef CHUNKSTREAM_H
#define CHUNKSTREAM_H

#include <QDataStream>
#include <QBuffer>
#include <QByteArray>
#include <QMap>


class ChunkStream
{
public:
	enum Version {
		Version_1 /*! First, old version */,
		Version_2 /*! Second, more optimized version */ = 2,
	};
	ChunkStream(const QByteArray & data): version_(Version_2) {setSource(data);}
	ChunkStream(QDataStream & str): version_(Version_2) {setSource(str);}
	ChunkStream(QByteArray * data = 0, Version v = Version_2): version_(v) {setSource(data);}
	ChunkStream(Version v): version_(v) {setSource(0);}
	~ChunkStream();

	template <typename T>
	struct Chunk {
		Chunk(int i, const T & d): id(i), data(d) {}
		int id;
		T data;
	};
	template <typename T> static Chunk<T> chunk(int id, const T & data) {return Chunk<T>(id, data);}

	template <typename T> ChunkStream & add(int id, const T & data) {*this << Chunk<T>(id, data); return *this;}

	void setSource(const QByteArray & data);
	void setSource(QDataStream & str);
	void setSource(QByteArray * data);
	QDataStream & dataStream() {return stream_;}
	QByteArray data() const {return tmp_data;}
	bool atEnd() const {return stream_.atEnd();}
	Version version() const {return (Version)version_;}

	int read();
	void readAll();
	int getID() {return last_id;}
	template <typename T>
	T getData() const {T ret; QDataStream s(last_data); s.setVersion(QDataStream::Qt_4_8); s >> ret; return ret;}
	template <typename T>
	void get(T & v) const {v = getData<T>();}
	template <typename T>
	const ChunkStream & get(int id, T & v) const {
		QByteArray ba = data_map.value(id);
		if (!ba.isEmpty()) {
			QDataStream s(ba); s.setVersion(QDataStream::Qt_4_8);
			s >> v;
		}
		return *this;
	}
private:
	void _init();

	static uint readVInt(QDataStream & s);
	static void writeVInt(QDataStream & s, uint val);

	int last_id;
	uchar version_;
	QByteArray * data_, last_data, tmp_data;
	QBuffer buffer;
	QDataStream stream_;
	QMap<int, QByteArray> data_map;

	template <typename T> friend ChunkStream & operator <<(ChunkStream & s, const ChunkStream::Chunk<T> & c);
};

template <typename T>
ChunkStream & operator <<(ChunkStream & s, const ChunkStream::Chunk<T> & c) {
	QByteArray ba;
	QDataStream bas(&ba, QIODevice::WriteOnly);
	bas.setVersion(QDataStream::Qt_4_8);
	bas << c.data;
	switch (s.version_) {
	case ChunkStream::Version_1:
		s.stream_ << c.id << ba;
		break;
	case ChunkStream::Version_2:
		if (s.data_->isEmpty())
			s.stream_ << uchar(uchar(s.version_) | 0x80);
		ChunkStream::writeVInt(s.stream_, c.id);
		ChunkStream::writeVInt(s.stream_, ba.size());
		s.stream_.writeRawData(ba.data(), ba.size());
		break;
	default: break;
	}
	return s;
}

#endif // CHUNKSTREAM_H