package main

import (
	"log"
	"math"
	"os"
	"regexp"
	"strings"

	"github.com/thoas/go-funk"
)

func main() {
	inputBytes, _ := os.ReadFile("input.txt")
	input := strings.Split(string(inputBytes), "\n")

	log.Println("Puzzle 1:", puzzle1(input))
	log.Println("Puzzle 2:", puzzle2(input))
}

const (
	len1 = 2
	len4 = 4
	len7 = 3
	len8 = 7
)

func puzzle1(input []string) int {
	inputRegex := regexp.MustCompile(`([a-g]+)`)
	count := 0
	for _, line := range input {
		matches := inputRegex.FindAllString(line, 14)
		if len(matches) != 14 {
			panic("Invalid input")
		}
		outputValues := matches[10:]
		for _, outputValue := range outputValues {
			if l := len(outputValue); l == len1 || l == len4 || l == len7 || l == len8 {
				count++
			}
		}
	}
	return count
}

func puzzle2(input []string) int {
	inputRegex := regexp.MustCompile(`([a-g]+)`)
	count := 0
	for _, line := range input {
		matches := inputRegex.FindAllString(line, 14)
		if len(matches) != 14 {
			panic("Invalid input")
		}
		lineDigits := matches[:10]
		digits := [10]string{}
		// Map by length
		lenGroups := make(map[int][]string)
		for _, digit := range lineDigits {
			l := len(digit)
			lenGroups[l] = append(lenGroups[l], digit)
		}
		// Get known digits
		digits[1] = lenGroups[len1][0]
		digits[4] = lenGroups[len4][0]
		digits[7] = lenGroups[len7][0]
		digits[8] = lenGroups[len8][0]
		// Deducing the rest
		// 9 (Length 6): Four chars are the same with 4
		for i, digit := range lenGroups[6] {
			if countIntersectChars(digits[4], digit) == 4 {
				digits[9] = digit
				lenGroups[6] = append(lenGroups[6][:i], lenGroups[6][i+1:]...)
				break
			}
		}
		// 0 (Length 6): Three chars are the same with 7
		for i, digit := range lenGroups[6] {
			if countIntersectChars(digits[7], digit) == 3 {
				digits[0] = digit
				lenGroups[6] = append(lenGroups[6][:i], lenGroups[6][i+1:]...)
				break
			}
		}
		// 6 (Length 6): Last remaining from lenGroup[6]
		digits[6] = lenGroups[6][0]
		// 3 (Length 5): Two chars are the same with 1
		for i, digit := range lenGroups[5] {
			if countIntersectChars(digits[1], digit) == 2 {
				digits[3] = digit
				lenGroups[5] = append(lenGroups[5][:i], lenGroups[5][i+1:]...)
				break
			}
		}
		// 2 (Length 5): Two chars are the same with 4
		for i, digit := range lenGroups[5] {
			if countIntersectChars(digits[4], digit) == 2 {
				digits[2] = digit
				lenGroups[5] = append(lenGroups[5][:i], lenGroups[5][i+1:]...)
				break
			}
		}
		// 5 (Length 5): Last remaining from lenGroup[5]
		digits[5] = lenGroups[5][0]
		// Add output values to count
		outputValues := matches[10:]
		for i, outputValue := range funk.ReverseStrings(outputValues) {
			l := len(outputValue)
			for j, digit := range digits {
				if l != len(digit) {
					continue
				}
				if countIntersectChars(outputValue, digit) == l {
					count += j * int(math.Pow(10, float64(i)))
				}
			}
		}
	}
	return count
}

func countIntersectChars(a, b string) int {
	count := 0
	for _, c := range []byte(a) {
		if funk.Contains([]byte(b), c) {
			count++
		}
	}
	return count
}