import fs from "fs";

type Point = {x: number, y: number};

enum Direction {
    UP = '^',
    DOWN = 'v',
    LEFT = '<',
    RIGHT = '>'
}

const input = fs.readFileSync("./15/input.txt", "utf-8").split(/[\r\n]{4,}/);
const warehouse: string[][] = input[0]
    .split(/[\r\n]+/)
    .map(row => row.split(""));
const movements: Direction[] = input[1]
    .split("")
    .map(char => char.trim())
    .filter(Boolean)
    .map(char => char as Direction);

// Part 1
console.info("Part 1: " + solve(warehouse, movements));

// Part 2
const secondWarehouse = warehouse.map(row => {
    const newRow: string[] = [];
    for (const char of row) {
        if (char === '#') { newRow.push('#', '#'); }
        else if (char === 'O') { newRow.push('[', ']'); }
        else if (char === '.') { newRow.push('.', '.'); }
        else { newRow.push('@', '.'); }
    }
    return newRow;
});
console.info("Part 2: " + solve(secondWarehouse, movements));

function solve(warehouse: string[][], movements: Direction[]): number {
    let _warehouse = warehouse.map(row => [...row]); // Take a copy to avoid modifying the original
    const robotLocation: Point = findStartLocation(_warehouse);

    for (const move of movements) {
        // Under some very specific circumstances in part 2, tryMove returns false, but the grid has already been modified
        // "Fix" the issue ba taking a copy so we can easily revert all changes made
        // Slow AF of course but rest of this code isn't optimized either, so...
        const copy = _warehouse.map(row => [...row]);
    
        if (tryMove(robotLocation, move, _warehouse)) {
            if (move === Direction.UP) { robotLocation.y--; }
            else if (move === Direction.DOWN) { robotLocation.y++; }
            else if (move === Direction.LEFT) { robotLocation.x--; }
            else { robotLocation.x++; }
        } else {
            _warehouse = copy; // Revert changes
        }
    }

    // GPS
    let result = 0;
    for (let y = 0; y < _warehouse.length; y++) {
        for (let x = 0; x < _warehouse[y].length; x++) {
            if (_warehouse[y][x] === "O" || _warehouse[y][x] === "[") {
                result += 100 * y + x;
            }
        }
    }
    return result;
}

function tryMove(from: Point, direction: Direction, warehouse: string[][], movingPair = false): boolean {
    const moveWhat = warehouse[from.y][from.x];
    if (moveWhat === "#") {
        return false;
    }

    let to: Point;
    switch (direction) {
        case Direction.UP: to = {x: from.x, y: from.y - 1}; break;
        case Direction.DOWN: to = {x: from.x, y: from.y + 1}; break;
        case Direction.LEFT: to = {x: from.x - 1, y: from.y}; break;
        case Direction.RIGHT: to = {x: from.x + 1, y: from.y}; break;
    }

    const allowMove = warehouse[to.y][to.x] === "."
        || (direction === Direction.UP && tryMove({x: from.x, y: from.y - 1}, direction, warehouse))
        || (direction === Direction.DOWN && tryMove({x: from.x, y: from.y + 1}, direction, warehouse))
        || (direction === Direction.LEFT && tryMove({x: from.x - 1, y: from.y}, direction, warehouse))
        || (direction === Direction.RIGHT && tryMove({x: from.x + 1, y: from.y}, direction, warehouse));

    if (allowMove) {
        // Part 1 logic handles horizontal movement of larger boxes just fine. Needs special handling for vertical movement
        if (!movingPair && (direction === Direction.UP || direction === Direction.DOWN)) {
            if (moveWhat === "[" && !tryMove({x: from.x + 1, y: from.y}, direction, warehouse, true)) {
                return false;
            }
            if (moveWhat === "]" && !tryMove({x: from.x - 1, y: from.y}, direction, warehouse, true)) {
                return false;
            }
        }

        // Make the move
        warehouse[to.y][to.x] = moveWhat;
        warehouse[from.y][from.x] = ".";
        return true;
    }
    return false;
}

function findStartLocation(warehouse: string[][]): Point {
    for (let y = 0; y < warehouse.length; y++) {
        for (let x = 0; x < warehouse[y].length; x++) {
            if (warehouse[y][x] === "@") {
                return {x,y};
            }
        }
    }

    throw new Error("Could not find start location!");
}

type
  Vec2 = tuple[x,y: int]
  Box = array[2, Vec2]
  Dir = enum
    U = "^"
    R = ">"
    D = "v"
    L = "<"

proc convertPart2(grid: seq[string]): seq[string] =
  for y in 0..grid.high:
    result.add ""
    for x in 0..grid[0].high:
      result[^1] &= (
        if grid[y][x] == 'O': "[]"
        elif grid[y][x] == '#': "##"
        else: "..")

proc shiftLeft(grid: var seq[string], col: int, range: HSlice[int,int]) =
  for i in range.a ..< range.b:
    grid[col][i] = grid[col][i+1]
  grid[col][range.b] = '.'

proc shiftRight(grid: var seq[string], col: int, range: HSlice[int,int]) =
  for i in countDown(range.b, range.a+1):
    grid[col][i] = grid[col][i-1]
  grid[col][range.a] = '.'

proc box(pos: Vec2, grid: seq[string]): array[2, Vec2] =
  if grid[pos.y][pos.x] == '[':
    [pos, (pos.x+1, pos.y)]
  else:
    [(pos.x-1, pos.y), pos]

proc step(grid: var seq[string], bot: var Vec2, dir: Dir) =
  var (x, y) = bot
  case dir
  of U:
    while (dec y; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y-1][bot.x] == 'O': swap(grid[bot.y-1][bot.x], grid[y][x])
    dec bot.y
  of R:
    while (inc x; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y][bot.x+1] == 'O': swap(grid[bot.y][bot.x+1], grid[y][x])
    inc bot.x
  of L:
    while (dec x; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y][bot.x-1] == 'O': swap(grid[bot.y][bot.x-1], grid[y][x])
    dec bot.x
  of D:
    while (inc y; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y+1][bot.x] == 'O': swap(grid[bot.y+1][bot.x], grid[y][x])
    inc bot.y

proc canMoveVert(box: Box, grid: seq[string], boxes: var HashSet[Box], dy: int): bool =
  boxes.incl box
  var left, right = false
  let (lbox, rbox) = (box[0], box[1])
  let lbigBox = box((lbox.x, lbox.y+dy), grid)
  let rbigBox = box((rbox.x, lbox.y+dy), grid)

  if grid[lbox.y+dy][lbox.x] == '#' or
     grid[rbox.y+dy][rbox.x] == '#': return false
  elif grid[lbox.y+dy][lbox.x] == '.': left = true
  else:
    left = canMoveVert(box((lbox.x,lbox.y+dy), grid), grid, boxes, dy)

  if grid[rbox.y+dy][rbox.x] == '.': right = true
  elif lbigBox == rbigBox: right = left
  else:
    right = canMoveVert(box((rbox.x, rbox.y+dy), grid), grid, boxes, dy)

  left and right

proc moveBoxes(grid: var seq[string], boxes: var HashSet[Box], d: Vec2) =
  for box in boxes:
    grid[box[0].y][box[0].x] = '.'
    grid[box[1].y][box[1].x] = '.'
  for box in boxes:
    grid[box[0].y+d.y][box[0].x+d.x] = '['
    grid[box[1].y+d.y][box[1].x+d.x] = ']'
  boxes.clear()

proc step2(grid: var seq[string], bot: var Vec2, dir: Dir) =
  case dir
  of U:
    if grid[bot.y-1][bot.x] == '#': return
    if grid[bot.y-1][bot.x] == '.': dec bot.y
    else:
      var boxes: HashSet[Box]
      if canMoveVert(box((x:bot.x, y:bot.y-1), grid), grid, boxes, -1):
        grid.moveBoxes(boxes, (0, -1))
        dec bot.y
  of R:
    var (x, y) = bot
    while (inc x; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y][bot.x+1] == '[': grid.shiftRight(bot.y, bot.x+1..x)
    inc bot.x
  of L:
    var (x, y) = bot
    while (dec x; grid[y][x] != '#' and grid[y][x] != '.'): discard
    if grid[y][x] == '#': return
    if grid[bot.y][bot.x-1] == ']': grid.shiftLeft(bot.y, x..bot.x-1)
    dec bot.x
  of D:
    if grid[bot.y+1][bot.x] == '#': return
    if grid[bot.y+1][bot.x] == '.': inc bot.y
    else:
      var boxes: HashSet[Box]
      if canMoveVert(box((x:bot.x, y:bot.y+1), grid), grid, boxes, 1):
        grid.moveBoxes(boxes, (0, 1))
        inc bot.y


proc solve(input: string): AOCSolution[int, int] =
  let chunks = input.split("\n\n")
  var grid = chunks[0].splitLines()
  let movements = chunks[1].splitLines().join().join()

  var robot: Vec2
  for y in 0..grid.high:
    for x in 0..grid[0].high:
      if grid[y][x] == '@':
        grid[y][x] = '.'
        robot = (x,y)

  block p1:
    var grid = grid
    var robot = robot
    for m in movements:
      let dir = parseEnum[Dir]($m)
      step(grid, robot, dir)
    for y in 0..grid.high:
      for x in 0..grid[0].high:
        if grid[y][x] == 'O':
          result.part1 += 100 * y + x

  block p2:
    var grid = grid.convertPart2()
    var robot = (robot.x*2, robot.y)
    for m in movements:
      let dir = parseEnum[Dir]($m)
      step2(grid, robot, dir)
      #grid.inspect(robot)

    for y in 0..grid.high:
      for x in 0..grid[0].high:
        if grid[y][x] == '[':
          result.part2 += 100 * y + x

import Control.Monad
import Data.Bifunctor
import Data.List
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set

type C = (Int, Int)

readInput :: String -> (Map C Char, [C])
readInput s =
  let (room, _ : moves) = break null $ lines s
   in ( Map.fromList [((i, j), c) | (i, l) <- zip [0 ..] room, (j, c) <- zip [0 ..] l],
        map dir $ concat moves
      )
  where
    dir '^' = (-1, 0)
    dir 'v' = (1, 0)
    dir '<' = (0, -1)
    dir '>' = (0, 1)

moveInto :: Int -> Set C -> C -> C -> Set C -> Maybe (Set C)
moveInto boxWidth walls (di, dj) = go
  where
    go (i, j) boxes
      | (i, j) `Set.member` walls = Nothing
      | Just j' <- find (\j' -> (i, j') `Set.member` boxes) $ map (j -) [0 .. boxWidth - 1] =
          Set.insert (i + di, j' + dj)
            <$> foldM
              (flip go)
              (Set.delete (i, j') boxes)
              [(i + di, j' + z + dj) | z <- [0 .. boxWidth - 1]]
      | otherwise = Just boxes

runMoves :: (Map C Char, [C]) -> Int -> Int
runMoves (room, moves) scale = score $ snd $ foldl' move (start, boxes) moves
  where
    room' = Map.mapKeysMonotonic (second (* scale)) room
    Just start = fst <$> find ((== '@') . snd) (Map.assocs room')
    walls =
      let ps = Map.keysSet $ Map.filter (== '#') room'
       in Set.unions [Set.mapMonotonic (second (+ z)) ps | z <- [0 .. scale - 1]]
    boxes = Map.keysSet $ Map.filter (== 'O') room'
    move (pos@(i, j), boxes) dir@(di, dj) =
      let pos' = (i + di, j + dj)
       in maybe (pos, boxes) (pos',) $ moveInto scale walls dir pos' boxes
    score = sum . map (\(i, j) -> i * 100 + j) . Set.elems

main = do
  input <- readInput <$> readFile "input15"
  mapM_ (print . runMoves input) [1, 2]