module World
  ( World,
    Direction (..),
    Vector,
    worldWithDimensions,
    getBeepersAt,
    setBeepersAt,
    hasWallAt,
    addWallAt,
    removeWallAt,
    positionWithinWorld,
  )
where

import qualified Data.Map.Strict as M
import qualified Data.Set as S

type Vector = (Int, Int)

data Direction = North | East | South | West deriving (Show, Eq, Ord)

data World = World
  { dimensions :: Vector,
    walls :: M.Map Vector (S.Set Direction),
    beepers :: M.Map Vector Int
  }
  deriving (Show, Eq)

worldWithDimensions :: Vector -> World
worldWithDimensions (width, height)
  | width <= 0 || height <= 0 = error "Cannot have world with zero or negative dimension"
  | otherwise =
    World
      { dimensions = (width, height),
        walls = M.empty,
        beepers = M.empty
      }

positionWithinWorld :: World -> Vector -> Bool
positionWithinWorld world (x, y) = valid_x && valid_y
  where
    (world_width, world_height) = dimensions world
    valid_x = x > 0 && x <= world_width
    valid_y = y > 0 && y <= world_height

setBeepersAt :: Vector -> Int -> World -> World
setBeepersAt position count world
  | not $ positionWithinWorld world position = error "Invalid position"
  | count < 0 = error "Invalid count of beepers"
  | count == 0 = world {beepers = M.delete position old_beepers}
  | otherwise = world {beepers = M.insert position count old_beepers}
  where
    old_beepers = beepers world

getBeepersAt :: Vector -> World -> Int
getBeepersAt position world
  | not $ positionWithinWorld world position = error "Invalid position"
  | otherwise = M.findWithDefault 0 position (beepers world)

switchDirectionToInternal :: Vector -> Direction -> (Vector, Direction)
switchDirectionToInternal position direction
  | direction == West = ((x - 1, y), East)
  | direction == South = ((x, y - 1), North)
  | otherwise = (position, direction)
  where
    (x, y) = position

isOnEdge :: Vector -> Vector -> Direction -> Bool
isOnEdge world_dimensions position direction = horizontally || vertically
  where
    (width, height) = world_dimensions
    (x, y) = position
    horizontally = x == width && direction == East
    vertically = y == height && direction == North

addWallAt :: Vector -> Direction -> World -> World
addWallAt position direction world
  | not $ positionWithinWorld world newPosition = error "Invalid position"
  | isOnEdge (dimensions world) newPosition newDirection = world
  | otherwise = world {walls = M.insert position newWalls (walls world)}
  where
    (newPosition, newDirection) = switchDirectionToInternal position direction
    oldWalls = M.findWithDefault S.empty position (walls world)
    newWalls = S.union oldWalls $ S.singleton direction

removeWallAt :: Vector -> Direction -> World -> World
removeWallAt position direction world
  | not $ positionWithinWorld world newPosition = error "Invalid position"
  | isOnEdge (dimensions world) newPosition newDirection = world
  | S.null newWalls = world {walls = M.delete position (walls world)}
  | otherwise = world {walls = M.insert position newWalls (walls world)}
  where
    (newPosition, newDirection) = switchDirectionToInternal position direction
    oldWalls = M.findWithDefault S.empty position (walls world)
    newWalls = oldWalls S.\\ S.singleton direction

hasWallAt :: Vector -> Direction -> World -> Bool
hasWallAt position direction world
  | not $ positionWithinWorld world newPosition = False
  | isOnEdge (dimensions world) newPosition newDirection = True
  | otherwise = case walls world M.!? position of
    Just wallsInPlace -> S.member direction wallsInPlace
    Nothing -> False
  where
    (newPosition, newDirection) = switchDirectionToInternal position direction