#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>


#define arrlen(x) (sizeof(x)/sizeof((x)[0]))


struct rock {
	unsigned char s[4];
};


struct rock const rocks[] = {
	{ { 0x1e, 0, 0, 0 } },		// Horizontal line
	{ { 0x08, 0x1c, 0x08, 0 } },	// Cross
	{ { 0x1c, 0x04, 0x04, 0 } },	// Elbow
	{ { 0x10, 0x10, 0x10, 0x10 } },	// Line piece
	{ { 0x18, 0x18, 0, 0 } },	// Block
};


struct cheat {
	unsigned long depth;
	struct rock rock;
	unsigned long num;
	unsigned long pos;
};


unsigned long
cheat_numdiff(struct cheat * cheats, int ncheats)
{
	return cheats[0].num - cheats[1].num;
}


unsigned long
cheat_posdiff(struct cheat * cheats, int ncheats)
{
	return cheats[0].pos - cheats[1].pos;
}


int
cheat_check(struct cheat * cheats, int ncheats,
		unsigned long depth, struct rock const * rock,
		unsigned long rocknum, unsigned long position)
{
	unsigned long numdiff, posdiff;
	int i;

	if (depth < cheats[0].depth)
		// Not deep enough.
		return 0;

	if (depth == cheats[0].depth
			&& memcmp(rock, &cheats[0].rock, sizeof(struct rock)) != 0)
		// Same depth as previously, but differnt rock, or in different
		// position in chamber.
		return 0;

	// Rock is either deeper than any previous, or is as deep and is the
	// same rock in the same position in the chamber

	fprintf(stderr, "Deepest rock %lu (depth %lu, height %lu)\n",
			rocknum, depth, position);

	memmove(cheats + 1, cheats, sizeof(struct cheat) * (ncheats - 1));
	cheats[0].depth = depth;
	cheats[0].rock = *rock;
	cheats[0].num = rocknum;
	cheats[0].pos = position;

	numdiff = cheat_numdiff(cheats, ncheats);
	posdiff = cheat_posdiff(cheats, ncheats);

	for (i = 2; i < ncheats; ++i)
		if (cheats[i - 1].num - cheats[i].num != numdiff
				|| cheats[i - 1].pos - cheats[i].pos != posdiff)
			// Not all the cheats are the same
			return 0;

	// All the cheats are the same!
	return 1;
}


#if 0
void
chamber_dump(unsigned char * chamber, int chamberheight, struct rock * rock, int rockheight)
{
	int i;

	for (i = chamberheight - 1; i >= 0; --i) {
		unsigned char c;

		fputc('|', stderr);

		for (c = 0x40; c != 0; c >>= 1) {
			if (rock && i >= rockheight
					&& i < rockheight + arrlen(rock->s)
					&& (rock->s[i - rockheight] & c))
				fputc('@', stderr);
			else if (chamber && (chamber[i] & c))
				fputc('#', stderr);
			else
				fputc('.', stderr);
		}


		fputc('|', stderr);
		fputc('\n', stderr);
	}
	fprintf(stderr, "+-------+\n");
}
#endif


int
main(int argc, char ** argv)
{
	int docheat = 0;
	size_t jetsize = 0;
	char * jets = NULL;
	unsigned char * chamber = NULL;
	unsigned long chamberheight = 0, top = 0, cheatby = 0;
	unsigned long movement, rocknum, rockheight;
	unsigned long rocklimit = 2022;
	struct cheat cheats[10] = {0};
	struct rock rock;
	int i;

	while ((i = getopt(argc, argv, "p:r:c")) != -1) {
		switch (i) {
		case 'p':
			switch (atoi(optarg)) {
			case 1: rocklimit = 2022; break;
			case 2: rocklimit = 1000000000000; break;
			default:
				fprintf(stderr, "Bad part: %s\n", optarg);
				return -1;
			}
			break;

		case 'r':
			rocklimit = atol(optarg);
			break;

		case 'c':
			docheat = 1;
			break;

		default:
			return -1;
		}
	}

	// Read jets.
	while (1) {
		void * p;
		size_t s;

		if ((p = realloc(jets, jetsize + BUFSIZ)) == NULL) {
			fprintf(stderr, "jets: bad realloc (%zu)\n", jetsize + BUFSIZ);
			free(jets);
			return -1;
		}
		jets = p;
		if ((s = fread(jets + jetsize, 1, BUFSIZ, stdin)) == 0)
			break;
		jetsize += s;
	}

	// Remove any characters that are not '<' or '>' (e.g. newlines)
	for (i = 0; i < jetsize; ++i) {
		while (i < jetsize && jets[i] != '<' && jets[i] != '>') {
			memmove(jets + i, jets + i + 1, jetsize - (i + 1));
			--jetsize;
		}
	}

	// Place first rock, and start movement
	rocknum = 0;
	rock = rocks[rocknum];
	rockheight = 3;

	for (movement = 0; rocknum < rocklimit; ++movement) {
		int jet;

		if (!chamberheight || top > chamberheight - 10) {
			void * p;

			if ((p = realloc(chamber, chamberheight + BUFSIZ)) == NULL) {
				fprintf(stderr, "Bad realloc(%lu)\n", chamberheight + BUFSIZ);
				free(chamber);
				free(jets);
				return -1;
			}
			chamber = p;
			memset(chamber + chamberheight, 0, BUFSIZ);
			chamberheight += BUFSIZ;
		}

		jet = jets[movement % jetsize];

		// Check if jet can push rock.
		for (i = 0; i < arrlen(rock.s); ++i) {
			// Check left wall collision
			if (jet == '<' && (rock.s[i] & 0x40) != 0)
				jet = 0;

			// Check right wall collision
			if (jet == '>' && (rock.s[i] & 0x01) != 0)
				jet = 0;

			// Check moving left into existing rocks.
			if (jet == '<' && ((rock.s[i] << 1) & chamber[rockheight + i]) != 0)
				jet = 0;

			// Check moving right into existing rocks
			if (jet == '>' && ((rock.s[i] >> 1) & chamber[rockheight + i]) != 0)
				jet = 0;
		}

		// If jet can push rock, do so.
		if (jet == '<')
			for (i = 0; jet && i < arrlen(rock.s); ++i)
				rock.s[i] <<= 1;
		if (jet == '>')
			for (i = 0; jet && i < arrlen(rock.s); ++i)
				rock.s[i] >>= 1;

		// Check if rock can fall.
		for (i = 0; rockheight > 0 && i < arrlen(rock.s); ++i)
			if ((rock.s[i] & chamber[rockheight - 1 + i]) != 0)
				break;
		if (i < arrlen(rock.s)) {
			// Rock cannot fall. Place rock.
			for (i = 0; i < arrlen(rock.s) && rock.s[i] != 0; ++i) {
				chamber[rockheight + i] |= rock.s[i];
				if (rockheight + i > top)
					top = rockheight + i;
			}

			// Check to see if we can cheat
			if (docheat
					&& rocknum > arrlen(rocks) * jetsize
					&& rockheight < top
					&& cheat_check(cheats, arrlen(cheats),
						top - rockheight, &rock,
						rocknum, rockheight))
			{
				unsigned long numdiff, posdiff, steps;

				numdiff = cheat_numdiff(cheats, arrlen(cheats));
				posdiff = cheat_posdiff(cheats, arrlen(cheats));

				steps = (rocklimit - rocknum) / numdiff;

				fprintf(stderr, "Cheating from rock %lu to %lu in %lu steps of %lu rocks/%lu height\n",
						rocknum, rocklimit, steps, numdiff, posdiff);

				rocknum += steps * numdiff;
				cheatby += steps * posdiff;
			}

			// Pick next rock.
			rock = rocks[++rocknum % arrlen(rocks)];
			rockheight = top + 4;

#if 0
			chamber_dump(chamber, top + 7, &rock, rockheight);
#endif
		}
		else {
			// Rock can fall, so reduce its height by 1
			--rockheight;
		}
	}

	printf("Height of rocks in chamber: %lu\n", top + cheatby + 1);

	free(chamber);
	free(jets);

	return 0;
}