package main

import (
	"bufio"
	"fmt"
	"os"
	"slices"
	"strconv"
	"strings"
	"time"
)

const FILE_PATH = "./input.txt"

//const FILE_PATH = "./sample.txt"

func LoadInput(filename string) ([]string, error) {
	file, err := os.Open(filename)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var lines []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		lines = append(lines, scanner.Text())
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return lines, nil
}

type structuredInput = []int

type Stone struct {
	number int
	count  int
}

func TransformInput(lines []string) structuredInput {
	var result = make([]int, 0)
	var numbs = strings.Split(lines[0], " ")
	var temp int

	for _, numb := range numbs {
		temp, _ = strconv.Atoi(numb)
		result = append(result, temp)
	}

	return result
}

func CreateStones(input []int) []Stone {
	var result = make([]Stone, 0)
	var temp Stone

	for _, numb := range input {
		temp = Stone{numb, 1}
		result = append(result, temp)
	}

	return result
}

func CombineStones(stones []Stone) []Stone {
	for i := 0; i < len(stones)-1; i++ {
		for j := len(stones) - 1; j > i; j-- {
			if stones[i].number == stones[j].number {
				stones[i].count += stones[j].count
				stones = slices.Delete(stones, j, j+1)
			}
		}
	}
	return stones
}

func GenerateNextStep(stones []Stone) []Stone {
	var result = make([]Stone, 0)
	var temp string
	var numb int
	for _, stone := range stones {
		if stone.number == 0 {
			result = append(result, Stone{1, stone.count})
			continue
		}
		temp = strconv.Itoa(stone.number)
		if len(temp)%2 == 0 {
			numb, _ = strconv.Atoi(temp[:len(temp)/2])
			result = append(result, Stone{numb, stone.count})
			numb, _ = strconv.Atoi(temp[len(temp)/2:])
			result = append(result, Stone{numb, stone.count})
			continue
		}

		result = append(result, Stone{stone.number * 2024, stone.count})
	}

	return result
}

func Part1(input structuredInput) int {
	var result = 0

	var stones = CreateStones(input)

	for i := 0; i < 25; i++ {
		stones = GenerateNextStep(stones)
		stones = CombineStones(stones)
	}
	for i := 0; i < len(stones); i++ {
		result += stones[i].count
	}

	return result
}

func Part2(input structuredInput) int {
	var result = 0

	var stones = CreateStones(input)

	for i := 0; i < 75; i++ {
		stones = GenerateNextStep(stones)
		stones = CombineStones(stones)
	}
	for i := 0; i < len(stones); i++ {
		result += stones[i].count
	}

	return result
}

func main() {
	var start time.Time
	var elapsed time.Duration = 0
	input, err := LoadInput(FILE_PATH)
	if err != nil {
		fmt.Println("Error loading input:", err)
		return
	}

	structuredInput := TransformInput(input)
	fmt.Println("Structured Input:", structuredInput)

	start = time.Now()
	fmt.Println("Solution Part 1: ", Part1(structuredInput))
	elapsed = time.Since(start)
	fmt.Printf("Part 1 took %s\n", elapsed)

	start = time.Now()
	fmt.Println("Solution Part 2: ", Part2(structuredInput))
	elapsed = time.Since(start)
	fmt.Printf("Part 2 took %s\n", elapsed)
}