package day05

import (
	"strconv"
	"strings"
)

type Rule struct {
	X, Y int
}

func ParseInput(input string) ([]Rule, [][]int) {
	sections := strings.Split(input, "\n\n")

	ruleLines := strings.Split(strings.TrimSpace(sections[0]), "\n")
	var rules []Rule
	for _, ruleLine := range ruleLines {
		parts := strings.Split(ruleLine, "|")
		x, _ := strconv.Atoi(parts[0])
		y, _ := strconv.Atoi(parts[1])
		rules = append(rules, Rule{X: x, Y: y})
	}

	updateLines := strings.Split(strings.TrimSpace(sections[1]), "\n")
	var updates [][]int
	for _, updateLine := range updateLines {
		parts := strings.Split(updateLine, ",")
		var update []int
		for _, part := range parts {
			page, _ := strconv.Atoi(part)
			update = append(update, page)
		}
		updates = append(updates, update)
	}

	return rules, updates
}

func ValidateUpdate(update []int, rules []Rule) bool {
	indexMap := make(map[int]int)
	for i, page := range update {
		indexMap[page] = i
	}

	for _, rule := range rules {
		pageX, indexX := indexMap[rule.X]
		pageY, indexY := indexMap[rule.Y]
		if indexX && indexY && pageX > pageY {
			return false
		}
	}

	return true
}

func FindMiddle(update []int) int {
	n := len(update)
	return update[n/2]
}

func SortUpdate(update []int, rules []Rule) []int {
	depen := make(map[int][]int)
	for _, rule := range rules {
		depen[rule.Y] = append(depen[rule.Y], rule.X)
	}

	visited := make(map[int]bool)
	var sorted []int

	for _, page := range update {
		if !visited[page] {
			sorted = Visit(page, depen, visited, update, sorted)
		}
	}

	return sorted
}

func Visit(page int, depen map[int][]int, visited map[int]bool, update []int, sorted []int) []int {
	if visited[page] {
		return sorted
	}
	visited[page] = true

	for _, d := range depen[page] {
		if contains(update, d) {
			sorted = Visit(d, depen, visited, update, sorted)
		}
	}

	sorted = append(sorted, page)
	return sorted
}

func contains(slice []int, element int) bool {
	for _, item := range slice {
		if item == element {
			return true
		}
	}
	return false
}

func Part1(input string) int {
	rules, updates := ParseInput(input)
	total := 0

	for _, update := range updates {
		if ValidateUpdate(update, rules) {
			total += FindMiddle(update)
		}
	}

	return total
}

func Part2(input string) int {
	rules, updates := ParseInput(input)
	var invalidUpdates [][]int
	total := 0

	for _, update := range updates {
		if !ValidateUpdate(update, rules) {
			invalidUpdates = append(invalidUpdates, update)
		}
	}

	for _, update := range invalidUpdates {
		sortedUpdate := SortUpdate(update, rules)
		total += FindMiddle(sortedUpdate)
	}

	return total
}