#include "stdafx.h"
#include "Ppm.h"
#include "Exception.h"
#include "Dither.h"
#include <math.h>

namespace graphics {

	static Bool CODECALL ppmApplicable(IStream *from) {
		const char *spec[] = {
			"P",
			"123456",
			" \n\r\t",
		};

		const Nat len = ARRAY_COUNT(spec);
		Buffer buffer = from->peek(storm::buffer(from->engine(), len));
		if (!buffer.full())
			return false;

		for (Nat i = 0; i < len; i++) {
			bool ok = false;
			for (const char *at = spec[i]; *at; at++) {
				ok |= byte(*at) == buffer[i];
			}

			if (!ok)
				return false;
		}

		return true;
	}

	static FormatOptions *CODECALL ppmCreate(ImageFormat *f) {
		return new (f) PPMOptions();
	}

	// Note: We go with ASCII by default, since the purpose of PPM is likely to write something that
	// is trivial to read somewhere else.
	PPMOptions::PPMOptions() : binary(false), mode(color) {}

	PPMOptions::PPMOptions(Bool binary, Mode mode) : binary(binary), mode(mode) {}

	void PPMOptions::toS(StrBuf *to) const {
		*to << S("PPM: { ") << (binary ? S("binary") : S("ascii")) << S(", ");
		switch (mode) {
		case mono1:
			*to << S("monochrome, 1-bit");
			break;
		case mono:
			*to << S("monochrome");
			break;
		case color:
			*to << S("color");
			break;
		}
		*to << S(" }");
	}

	ImageFormat *ppmFormat(Engine &e) {
		const wchar *exts[] = {
			S("ppm"),
			S("pbm"),
			null
		};
		return new (e) ImageFormat(S("Portable Pixel Map"), exts, &ppmApplicable, &ppmCreate);
	}

	// Is this a whitespace character?
	static bool whitespace(Byte ch) {
		switch (ch) {
		case ' ':
		case '\n':
		case '\r':
		case '\t':
			return true;
		default:
			return false;
		}
	}

	// Is this a numeric character?
	static bool numeric(Byte ch) {
		return ch >= '0' && ch <= '9';
	}

	struct ReadState {
		IStream *src;
		Buffer buffer;
		Nat pos;

		ReadState(IStream *src) :
			src(src),
			buffer(src->read(1024)),
			pos(0) {}

		// Get the next byte in the file.
		Byte next() {
			if (pos >= buffer.filled()) {
				pos = 0;
				buffer.filled(0);
				buffer = src->read(buffer);

				if (buffer.empty())
					throw new (src) ImageLoadError(S("Unexpected end of stream!"));
			}

			return buffer[pos++];
		}

		// Get the next character, ignoring comments.
		Byte nextText() {
			Byte r;
			for (r = next(); r == '#'; r = next()) {
				// In a comment. Skip until newline.
				while (next() != '\n')
					;
			}
			return r;
		}

		// Get the next number (in ASCII) from the file.
		Nat nextNum() {
			// Skip whitespace.
			Byte ch;
			while (whitespace(ch = nextText()))
				;

			// Read the number.
			if (!numeric(ch))
				throw new (src) ImageLoadError(S("Not a number!"));

			Nat result = 0;
			while (numeric(ch)) {
				result *= 10;
				result += ch - '0';
				ch = nextText();
			}

			return result;
		}
	};


	class WriteState {
	public:
		Image *image;
		OStream *out;
		Buffer buffer;

		WriteState(Image *image, OStream *out) :
			image(image),
			out(out),
			buffer(storm::buffer(out->engine(), 1024)) {}

		~WriteState() {
			flush();
		}

		void put(Byte b) {
			if (buffer.full())
				flush();
			Nat p = buffer.filled();
			buffer[p] = b;
			buffer.filled(p + 1);
		}

		void put(const Byte *buf, Nat count) {
			// We assume that 'count' is always smaller than the buffer.
			if (buffer.filled() + count >= buffer.count())
				flush();
			Nat p = buffer.filled();
			memcpy(buffer.dataPtr() + p, buf, count);
			buffer.filled(p + count);
		}

		void put(const char *buf, Nat count) {
			put((const Byte *)buf, count);
		}

		void putNum(Nat num, Nat minWidth) {
			Byte buffer[11] = "         0";
			Nat pos = 9;
			while (num > 0) {
				buffer[pos--] = '0' + num % 10;
				num = num / 10;
			}
			pos = min(pos, 9 - minWidth);
			put(buffer + pos + 1, 9 - pos);
		}

		void flush() {
			if (!buffer.empty())
				out->write(buffer);
			buffer.filled(0);
		}

	};


	struct Header {
		// Mode: We use 1, 2 and 3 regardless if binary or ascii.
		Byte mode;

		// Size.
		Nat width, height;
	};

	template <class T>
	Image *loadMonoPPM(ReadState &src, const Header &header, T read) {
		Image *out = new (src.src) Image(header.width, header.height);

		for (Nat y = 0; y < header.height; y++) {
			for (Nat x = 0; x < header.width; x++) {
				Float c = read.next(src);
				out->set(x, y, Color(c, c, c));
			}
			read.flush();
		}

		return out;
	}

	template <class T>
	Image *loadColorPPM(ReadState &src, const Header &header, T read) {
		Image *out = new (src.src) Image(header.width, header.height);

		for (Nat y = 0; y < header.height; y++) {
			for (Nat x = 0; x < header.width; x++) {
				Float r = read.next(src);
				Float g = read.next(src);
				Float b = read.next(src);
				out->set(x, y, Color(r, g, b));
			}
			read.flush();
		}

		return out;
	}

	template <class T>
	Image *loadPPM(ReadState &src, const Header &header) {
		Nat maxval = 0;

		switch (header.mode) {
		case 1:
			return loadMonoPPM(src, header, typename T::Mono());
		case 2:
			maxval = src.nextNum();
			return loadMonoPPM(src, header, typename T::Multi(maxval));
		case 3:
			maxval = src.nextNum();
			return loadColorPPM(src, header, typename T::Multi(maxval));
		default:
			return null;
		}
	}

	struct Raw {
		struct Multi {
			Float maxval;
			Bool multibyte;

			Multi(Nat maxval) : maxval(Float(maxval)), multibyte(maxval > 255) {}

			Float next(ReadState &src) const {
				Nat result = src.next();
				if (multibyte)
					result = (result << 8) | src.next();
				return Float(result) / maxval;
			}

			void put(WriteState &to, Byte px) {
				to.put(px);
			}
			void put(WriteState &to, Float pxData) {
				Int value = Int(pxData * 255.0f);
				put(to, Byte(max(255, min(0, value))));
			}

			void nextLine(WriteState &to) {
				// Nothing to do...
				(void)to;
			}

			void flush() const {}
		};

		struct Mono {
			Byte data;
			Byte fill;
			DitherState dither;

			Mono() : data(0), fill(0) {}

			Float next(ReadState &src) {
				if (fill == 0) {
					data = src.next();
					fill = 8;
				}

				fill--;
				return Float(~(data >> fill) & 0x1);
			}

			void flush() {
				fill = 0;
			}

			void put(WriteState &to, Float pxData) {
				if (fill >= 8)
					nextLine(to);
				data <<= 1;
				data |= dither.pixelValue(pxData) ? 0 : 1;
				fill++;
			}

			void nextLine(WriteState &to) {
				if (fill) {
					data <<= 8 - fill;
					to.put(data);
				}
				data = fill = 0;
			}
		};
	};

	struct Ascii {
		struct Multi {
			Float maxval;
			Bool outFirst;

			Multi(Nat maxval) : maxval(Float(maxval)), outFirst(true) {}

			Float next(ReadState &src) const {
				return Float(src.nextNum()) / maxval;
			}

			void flush() const {}

			void put(WriteState &to, Byte pxData) {
				if (!outFirst)
					to.put(' ');
				to.putNum(pxData, 3);
				outFirst = false;
			}
			void put(WriteState &to, Float pxData) {
				Int value = Int(floorf((pxData * 255.0f) + 0.5f));
				put(to, Byte(min(255, max(0, value))));
			}

			void nextLine(WriteState &to) {
				to.put('\n');
				outFirst = true;
			}
		};

		struct Mono {
			Bool outFirst;
			DitherState dither;

			Mono() : outFirst(true) {}

			Float next(ReadState &src) const {
				return Float(src.nextNum());
			}

			void flush() const {}

			void put(WriteState &to, Float pxData) {
				if (!outFirst)
					to.put(' ');
				to.put(dither.pixelValue(pxData) ? '0' : '1');
				outFirst = false;
			}

			void nextLine(WriteState &to) {
				to.put('\n');
				outFirst = true;
			}
		};
	};

	Image *PPMOptions::load(IStream *src) {
		ReadState s(src);

		if (s.next() != 'P')
			throw new (this) ImageLoadError(S("Not a supported PPM file."));

		Header header;

		header.mode = s.next();
		if (header.mode >= '1' && header.mode <= '6')
			header.mode -= '0';
		else
			return null;

		if (!whitespace(s.next()))
			return null;

		header.width = s.nextNum();
		header.height = s.nextNum();

		if (header.mode > 3) {
			header.mode -= 3;
			this->mode = Mode(header.mode);
			this->binary = true;
			return loadPPM<Raw>(s, header);
		} else {
			this->mode = Mode(header.mode);
			this->binary = false;
			return loadPPM<Ascii>(s, header);
		}
	}

	template <class T>
	void saveMonoPPM(WriteState &s, T write) {
		// Now we can just put numbers:
		for (Nat y = 0; y < s.image->height(); y++) {
			for (Nat x = 0; x < s.image->width(); x++) {
				Color color = s.image->get(x, y);
				Float average = fromLinear(color.toLinear().brightness());
				write.put(s, average);
			}
			write.nextLine(s);
		}
	}

	template <class T>
	void saveColorPPM(WriteState &s, T write) {
		// Now we can just put numbers for RGB:
		for (Nat y = 0; y < s.image->height(); y++) {
			for (Nat x = 0; x < s.image->width(); x++) {
				Byte *buf = s.image->buffer(x, y);
				write.put(s, buf[0]);
				write.put(s, buf[1]);
				write.put(s, buf[2]);
			}
			write.nextLine(s);
		}
	}

	template <class T>
	void savePPM(WriteState &s, PPMOptions::Mode mode) {
		switch (mode) {
		case PPMOptions::mono1:
			saveMonoPPM(s, typename T::Mono());
			break;
		case PPMOptions::mono:
			// Second part of the header, max value.
			s.put("255\n", 4);
			saveMonoPPM(s, typename T::Multi(255));
			break;
		case PPMOptions::color:
			// Second part of the header, max value.
			s.put("255\n", 4);
			saveColorPPM(s, typename T::Multi(255));
			break;
		}
	}

	void PPMOptions::save(Image *image, OStream *to) {
		WriteState s(image, to);

		// Write the header.
		Byte header[4] = "P0\n";
		header[1] += Byte(mode) + (binary ? 3 : 0);
		s.put(header, 3);

		s.putNum(image->width(), 1);
		s.put(' ');
		s.putNum(image->height(), 1);
		s.put('\n');

		if (binary) {
			savePPM<Raw>(s, mode);
		} else {
			savePPM<Ascii>(s, mode);
		}
	}

}