package main

import (
	"bufio"
	"fmt"
	"os"
	"unicode"

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

type EnginePart struct {
	number     uint64
	line       uint64
	startIndex uint64
	endIndex   uint64
}

func readFileLines(filename string) ([]string, error) {
	var lines []string

	file, err := os.Open(filename)
	if err != nil {
		return lines, err
	}
	fileScanner := bufio.NewScanner(file)
	fileScanner.Split(bufio.ScanLines)

	for fileScanner.Scan() {
		lines = append(lines, fileScanner.Text())
	}

	return lines, nil
}

func getEngineParts(lines *[]string) []EnginePart {
	var result []EnginePart
	var tmpPart EnginePart
	for lineIdx := 0; lineIdx < len(*lines); lineIdx++ {
		line := &(*lines)[lineIdx]
		tmpPart.number = 0
		tmpPart.line = uint64(lineIdx)
		numStarted := false
		for charIdx, char := range *line {
			if unicode.IsDigit(char) {
				if !numStarted {
					tmpPart.startIndex = uint64(charIdx)
					numStarted = true
				}
				tmpPart.number *= 10
				tmpPart.number += uint64(char) - uint64('0')
			} else if numStarted {
				numStarted = false
				tmpPart.endIndex = uint64(charIdx) - 1
				result = append(result, tmpPart)
				tmpPart.number = 0
			}
		}
		if numStarted {
			tmpPart.endIndex = uint64(len(*line)) - 1
			result = append(result, tmpPart)
			tmpPart.number = 0
		}
	}
	return result
}

func hasSymbol(line *string, startIndex uint64, endIndex uint64) bool {
	for i := startIndex; i <= endIndex; i++ {
		if (*line)[i] != '.' && !unicode.IsDigit(rune((*line)[i])) {
			return true
		}
	}
	return false
}

func getValidEngineParts(parts *[]EnginePart, lines *[]string) []uint64 {
	var validParts []uint64
	for i := 0; i < len(*parts); i++ {
		part := &(*parts)[i]
		startIndex := part.startIndex
		endIndex := part.endIndex
		if startIndex > 0 {
			startIndex--
		}
		if endIndex < uint64(len((*lines)[0]))-1 {
			endIndex++
		}
		for i := part.line - 1; i != part.line+2; i++ {
			if i >= uint64(len(*lines)) {
				continue
			}
			if hasSymbol(&(*lines)[i], startIndex, endIndex) {
				validParts = append(validParts, part.number)
				break
			}
		}
	}
	return validParts
}

func hasTwoAdjecent(parts *[]EnginePart, lineIdx uint64, charIdx uint64) (bool, []EnginePart) {
	var result []EnginePart
	for i := 0; i < len(*parts); i++ {
		part := &(*parts)[i]
		if part.line >= lineIdx-1 && part.line <= lineIdx+1 {
			if part.startIndex <= charIdx+1 && part.endIndex >= charIdx-1 {
				result = append(result, *part)
			}
		}
	}
	return len(result) == 2, result
}

func getGears(parts *[]EnginePart, lines *[]string) []uint64 {
	var result []uint64
	for lineIdx := 0; lineIdx < len(*lines); lineIdx++ {
		line := &((*lines)[lineIdx])
		for charIdx, char := range *line {
			if char == '*' {
				valid, engineParts := hasTwoAdjecent(parts, uint64(lineIdx), uint64(charIdx))
				if valid {
					result = append(result, engineParts[0].number*engineParts[1].number)
				}
			}
		}
	}
	return result
}

func main() {
	redPrint := color.New(color.FgRed)
	yellowPrint := color.New(color.FgYellow)
	if len(os.Args) < 2 {
		redPrint.Fprintln(os.Stderr, "You must provide the input file")
		os.Exit(1)
	}
	lines, err := readFileLines(os.Args[1])
	if err != nil {
		redPrint.Fprintln(os.Stderr, "Could not read input file")
		os.Exit(1)
	}
	engineParts := getEngineParts(&lines)

	validParts := getValidEngineParts(&engineParts, &lines)
	validPartSum := funk.Reduce(validParts, func(acc, elem uint64) uint64 { return acc + elem }, uint64(0))

	fmt.Print("Part 1: ")
	yellowPrint.Print(validPartSum)
	fmt.Println()

	gearRatios := getGears(&engineParts, &lines)
	gearRatiosSum := funk.Reduce(gearRatios, func(acc, elem uint64) uint64 { return acc + elem }, uint64(0))

	fmt.Print("Part 2: ")
	yellowPrint.Print(gearRatiosSum)
	fmt.Println()
}