package main

import (
	"bufio"
	"fmt"
	"os"
	"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 Vector2 struct {
	x int
	y int
}

const (
	EMPTY     = iota
	ROBOT     = iota
	BOX       = iota
	WALL      = iota
	LEFT_BOX  = iota
	RIGHT_BOX = iota
)

const (
	UP    = '^'
	DOWN  = 'v'
	LEFT  = '<'
	RIGHT = '>'
)

var MOVEMENT = map[byte]Vector2{
	UP:    Vector2{0, -1},
	DOWN:  Vector2{0, 1},
	LEFT:  Vector2{-1, 0},
	RIGHT: Vector2{1, 0},
}

type Lab struct {
	field   [][]int
	command string
}

func cloneLab(lab Lab) Lab {
	var field = make([][]int, len(lab.field))

	for y := 0; y < len(lab.field); y++ {
		field[y] = make([]int, len(lab.field[y]))
		for x := 0; x < len(lab.field[y]); x++ {
			field[y][x] = lab.field[y][x]
		}
	}

	return Lab{field, lab.command}
}

func cloneWideLab(lab Lab) Lab {
	var field = make([][]int, len(lab.field))
	for y := 0; y < len(lab.field); y++ {
		field[y] = make([]int, len(lab.field[y])*2)
		for x := 0; x < len(lab.field[y]); x++ {
			if lab.field[y][x] == BOX {
				field[y][x*2] = LEFT_BOX
				field[y][x*2+1] = RIGHT_BOX
			} else if lab.field[y][x] == ROBOT {
				field[y][x*2] = ROBOT
				field[y][x*2+1] = EMPTY
			} else {
				field[y][x*2] = lab.field[y][x]
				field[y][x*2+1] = lab.field[y][x]
			}
		}
	}
	return Lab{field, lab.command}
}

type structuredInput = Lab

func TransformInput(lines []string) structuredInput {
	var field = make([][]int, 0)
	var row []int
	var i = 0
	var endField = false
	var command string = ""

	for !endField {
		if lines[i] == "" {
			endField = true
			break
		}
		row = make([]int, 0)

		for j := 0; j < len(lines[i]); j++ {
			if lines[i][j] == '@' {
				row = append(row, ROBOT)
			} else if lines[i][j] == 'O' {
				row = append(row, BOX)
			} else if lines[i][j] == '#' {
				row = append(row, WALL)
			} else {
				row = append(row, EMPTY)
			}
		}
		field = append(field, row)
		i += 1
	}

	for ; i < len(lines); i++ {
		command += lines[i]
	}

	return Lab{field, command}
}

func tryPush(lab *Lab, position Vector2, direction Vector2) {
	if lab.field[position.y+direction.y][position.x+direction.x] == WALL {
		return
	}

	if lab.field[position.y+direction.y][position.x+direction.x] == EMPTY {
		lab.field[position.y+direction.y][position.x+direction.x] = lab.field[position.y][position.x]
		lab.field[position.y][position.x] = EMPTY
		return
	}

	tryPush(lab, Vector2{position.x + direction.x, position.y + direction.y}, direction)

	if lab.field[position.y+direction.y][position.x+direction.x] == EMPTY {
		lab.field[position.y+direction.y][position.x+direction.x] = lab.field[position.y][position.x]
		lab.field[position.y][position.x] = EMPTY
	}
}

func getRobot(lab Lab) Vector2 {
	for y := 0; y < len(lab.field); y++ {
		for x := 0; x < len(lab.field[y]); x++ {
			if lab.field[y][x] == ROBOT {
				return Vector2{x, y}
			}
		}
	}
	return Vector2{-1, -1}
}

func printLab(lab Lab, command byte) {
	for y := 0; y < len(lab.field); y++ {
		for x := 0; x < len(lab.field[y]); x++ {
			if lab.field[y][x] == ROBOT {
				fmt.Printf("%c", command)
			}
			if lab.field[y][x] == EMPTY {
				fmt.Print(".")
			}
			if lab.field[y][x] == WALL {
				fmt.Print("#")
			}
			if lab.field[y][x] == LEFT_BOX {
				fmt.Print("[")
			}
			if lab.field[y][x] == RIGHT_BOX {
				fmt.Print("]")
			}
			if lab.field[y][x] == BOX {
				fmt.Print("O")
			}
		}
		fmt.Println()
	}
}

func calculateLabScore(lab Lab) int {
	var result = 0
	for y := 0; y < len(lab.field); y++ {
		for x := 0; x < len(lab.field[y]); x++ {
			if lab.field[y][x] == BOX {
				result += x + y*100
			}
		}
	}
	return result
}

func calculateWideLabScore(lab Lab) int {
	var result = 0
	for y := 0; y < len(lab.field); y++ {
		for x := 0; x < len(lab.field[y]); x++ {
			if lab.field[y][x] == LEFT_BOX {
				result += x + y*100
			}
		}
	}
	return result
}

func Part1(input structuredInput) int {
	var lab = cloneLab(input)

	for i := 0; i < len(lab.command); i++ {
		tryPush(&lab, getRobot(lab), MOVEMENT[lab.command[i]])
	}

	return calculateLabScore(lab)
}

func clearScreen() {
	fmt.Print("\033[H\033[2J") // Clear entire screen
}

func moveCursorUp(lines int) {
	fmt.Printf("\033[%dA", lines) // Move cursor up n lines
}

func canPushRobot(lab *Lab, position Vector2, direction Vector2) bool {
	if lab.field[position.y+direction.y][position.x+direction.x] == WALL {
		return false
	}
	if lab.field[position.y+direction.y][position.x+direction.x] == EMPTY {
		return true
	}
	return canPushBox(lab, Vector2{position.x + direction.x, position.y + direction.y}, direction)
}

func canPushBox(lab *Lab, position Vector2, direction Vector2) bool {
	var left_position = position
	if lab.field[position.y][position.x] == RIGHT_BOX {
		left_position = Vector2{position.x - 1, position.y}
	}

	if direction == MOVEMENT[UP] || direction == MOVEMENT[DOWN] {
		if lab.field[left_position.y+direction.y][left_position.x] == WALL ||
			lab.field[left_position.y+direction.y][left_position.x+1] == WALL {
			return false
		}
		if lab.field[left_position.y+direction.y][left_position.x] == EMPTY &&
			lab.field[left_position.y+direction.y][left_position.x+1] == EMPTY {
			return true
		}
		if (lab.field[left_position.y+direction.y][left_position.x] == RIGHT_BOX || lab.field[left_position.y+direction.y][left_position.x] == LEFT_BOX) && !canPushBox(lab, Vector2{left_position.x, left_position.y + direction.y}, direction) {
			return false
		}

		if lab.field[left_position.y+direction.y][left_position.x+1] == LEFT_BOX && !canPushBox(lab, Vector2{left_position.x + 1, left_position.y + direction.y}, direction) {
			return false
		}

		return true
	}

	if direction == MOVEMENT[LEFT] {
		if lab.field[left_position.y][left_position.x-1] == WALL {
			return false
		}
		if lab.field[left_position.y][left_position.x-1] == EMPTY {
			return true
		}

		return canPushBox(lab, Vector2{left_position.x - 1, left_position.y}, direction)
	}

	if lab.field[left_position.y][left_position.x+2] == WALL {
		return false
	}

	if lab.field[left_position.y][left_position.x+2] == EMPTY {
		return true
	}

	return canPushBox(lab, Vector2{left_position.x + 2, left_position.y}, direction)
}

func pushRobot(lab *Lab, position Vector2, direction Vector2) {
	if lab.field[position.y+direction.y][position.x+direction.x] != EMPTY {
		pushBox(lab, Vector2{position.x + direction.x, position.y + direction.y}, direction)
	}
	lab.field[position.y+direction.y][position.x+direction.x] = ROBOT
	lab.field[position.y][position.x] = EMPTY
}

func pushBox(lab *Lab, position Vector2, direction Vector2) {
	var left_position = position
	if lab.field[position.y][position.x] == RIGHT_BOX {
		left_position = Vector2{position.x - 1, position.y}
	}

	if direction == MOVEMENT[UP] || direction == MOVEMENT[DOWN] {
		if lab.field[left_position.y+direction.y][left_position.x] == LEFT_BOX || lab.field[left_position.y+direction.y][left_position.x] == RIGHT_BOX {
			pushBox(lab, Vector2{left_position.x, left_position.y + direction.y}, direction)
		}

		if lab.field[left_position.y+direction.y][left_position.x+1] == LEFT_BOX {
			pushBox(lab, Vector2{left_position.x + 1, left_position.y + direction.y}, direction)
		}

		lab.field[left_position.y+direction.y][left_position.x] = LEFT_BOX
		lab.field[left_position.y+direction.y][left_position.x+1] = RIGHT_BOX
		lab.field[left_position.y][left_position.x] = EMPTY
		lab.field[left_position.y][left_position.x+1] = EMPTY
	}

	if direction == MOVEMENT[LEFT] {
		if lab.field[left_position.y][left_position.x-1] == RIGHT_BOX {
			pushBox(lab, Vector2{position.x - 1, position.y}, direction)
		}

		lab.field[left_position.y][left_position.x-1] = LEFT_BOX
		lab.field[left_position.y][left_position.x] = RIGHT_BOX
		lab.field[left_position.y][left_position.x+1] = EMPTY
	}

	if direction == MOVEMENT[RIGHT] {
		if lab.field[left_position.y][left_position.x+2] == LEFT_BOX {
			pushBox(lab, Vector2{position.x + 2, position.y}, direction)
		}

		lab.field[left_position.y][left_position.x+1] = LEFT_BOX
		lab.field[left_position.y][left_position.x+2] = RIGHT_BOX
		lab.field[left_position.y][left_position.x] = EMPTY
	}
}

func Part2(input structuredInput) int {
	var lab = cloneWideLab(input)
	//clearScreen()
	//printLab(lab, lab.command[0])
	//time.Sleep(1000 * time.Millisecond)
	for i := 0; i < len(lab.command); i++ {
		//moveCursorUp(len(lab.field))
		if canPushRobot(&lab, getRobot(lab), MOVEMENT[lab.command[i]]) {
			pushRobot(&lab, getRobot(lab), MOVEMENT[lab.command[i]])
		}
		//printLab(lab, lab.command[i+1])
		//time.Sleep(100 * time.Millisecond)
	}

	return calculateWideLabScore(lab)
}

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)
}