#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <err.h>
#include <math.h>
#include "input.h"

#define INPUT "input/08.txt"

#define INPUT_DIGIT_SIZE 10
#define OUTPUT_DIGIT_SIZE 4

#define EXPECTED1 504L
#define EXPECTED2 1073431L

void part1(struct input_str* input) {
	long sum = 0;
	for(size_t i = 0; i < input->line_count; i++) {
		// copy string as strtok will mess with it
		int len = strlen(input->lines[i]);
		char* line = malloc(len * sizeof(char));
		memset(line, '\0', len);
		strncpy(line, input->lines[i], len);
		char* output_start = strchr(line, '|');
		char* token = strtok(output_start + 1, " ");
		while(token != NULL) {
			int token_length = strlen(token);
			// respectively digists 1, 4, 7 and 8
			if(token_length == 2 || token_length == 4 || token_length == 3 || token_length == 7)
				sum++;

			token = strtok(NULL, " ");
		}
		//free(line); TODO causes memory corruption
	}

	CHECK(sum, EXPECTED1);
}

int count_segment(int digit) {
	int count = 0;
	for(int p = 0; p < 10; p++) {
		int segment_mask = 1 << p;
		if((digit & segment_mask) == segment_mask)
			count++;
	}
	return count;
}

int parse_digit(char* token) {
	// digits are parsed with segments mapped to bits as follows : 0abcdefg
	int result = 0;
	int p = 0;
	while(token[p] != '\0') {
		switch(token[p]) {
			case 'a':
				result |= (1 << 6);
				break;
			case 'b':
				result |= (1 << 5);
				break;
			case 'c':
				result |= (1 << 4);
				break;
			case 'd':
				result |= (1 << 3);
				break;
			case 'e':
				result |= (1 << 2);
				break;
			case 'f':
				result |= (1 << 1);
				break;
			case 'g':
				result |= 1;
				break;
			default:
				err(1, "unexpected token");
		}
		p++;
	}
	return result;
}

void part2(struct input_str* input) {
	long sum = 0;
	for(size_t i = 0; i < input->line_count; i++) {
		int digits[INPUT_DIGIT_SIZE];
		// copy string as strtok will mess with it
		int len = strlen(input->lines[i]);
		char* line = malloc(len * sizeof(char));
		memset(line, '\0', len);
		strncpy(line, input->lines[i], len);
		char* token = strtok(line, " ");
		for(int t = 0; t < INPUT_DIGIT_SIZE; t++) {
			digits[t] = parse_digit(token);
			token = strtok(NULL, " ");
		}

		int digits_code[10] = { 0 };
		// first pass to identify 1, 4, 7, 8
		for(int p = 0; p < 10; p++) {
			int digit = digits[p];
			int segment_count = count_segment(digit);
			if(segment_count == 3) {
				digits_code[7] = digit; // this is 7
				digits[p] = 0;
			}
			else if(segment_count == 4) {
				digits_code[4] = digit; // this is 4
				digits[p] = 0;
			}
			else if(segment_count == 2) {
				digits_code[1] = digit; // this is 0
				digits[p] = 0;
			}
			else if(segment_count == 7) {
				digits_code[8] = digit; // this is 8
				digits[p] = 0;
			}
		}
		// second pass identifies 2, 3 and 5
		for(int p = 0; p < 10; p++) {
			int digit = digits[p];
			if(digit == 0)
				continue; // this digit was identified by previous pass

			int digit_segment_count = count_segment(digit);
			if(digit_segment_count == 5) {
				if(count_segment(digits[p] & digits_code[7]) == 3) {
					digits_code[3] = digit; // this is 3 (5 segments and those of 7)
					digits[p] = 0;
				}
				else if(count_segment(digits[p] & digits_code[4]) == 2) {
					digits_code[2] = digit; // this is 2 (5 segments and 2 in common with 4)
					digits[p] = 0;
				}
				else if(count_segment(digits[p] & digits_code[4]) == 3) {
					digits_code[5] = digit; // this si 5 (5 segments and 3 in common with 4)
					digits[p] = 0;
				}
			}
		}
		// third pass identifies 0, 6 and 9
		for(int p = 0; p < 10; p++) {
			int digit = digits[p];
			if(digit == 0)
				continue;

#ifdef DEBUG
			// at this stage the remaining digits should have 6 segments
			if(count_segment(digit) != 6)
				printf("digit has %d segments while 6 expected\n", count_segment(digit));
#endif

			int segments_five_and_one = digits_code[5] | digits_code[1];
			if((segments_five_and_one & digit) == segments_five_and_one) {
				digits_code[9] = digit; // this is 9 (6 segments and has those of 5 and 1)
				digits[p] = 0;
			}
			else if((digits_code[5] & digit) == digits_code[5]) {
				digits_code[6] = digit; // this is 6 (6 segments, has those of 5 but not 1)
				digits[p] = 0;
			}
			else {
				digits_code[0] = digit; // this is 0
				digits[p] = 0;
			}
		}

#ifdef DEBUG
		// check that we found all digits
		for(int d = 0; d < 10; d++) {
			if(digits_code[d] == 0)
				printf("missing digit %d\n", d);
		}
#endif

		// use the table to parse output digits
		token = strtok(NULL, " "); // skip the '|' separator
		int current_number = 0;
		for(int n = 0; n < OUTPUT_DIGIT_SIZE && token != NULL; n++) {
			int parsed_digit = parse_digit(token);
			// search real digit
			for(int d = 0; d < 10; d++) {
				if(parsed_digit == digits_code[d]) {
					current_number += (int) pow(10, OUTPUT_DIGIT_SIZE - 1 - n) * d;
					break;
				}
			}
			// next digit
			token = strtok(NULL, " ");
		}
		sum += current_number;
		//free(line); TODO causes memory corruption
	}
	CHECK(sum, EXPECTED2);
}

int main() {
	// read input
	struct input_str input;
	input_str_read(&input, INPUT);

	// do stuff
	part1(&input);
	part2(&input);

	// cleanup & exit
	input_str_free(&input);
	return 0;
}