import Control.Monad (guard, when)
import Data.Array (Array, (!), (//), bounds, listArray)
import System.Console.ANSI
import System.IO (stdin, hReady)
import System.Random
import System.Timeout (timeout)
import Debug.Trace

data Pos
  = Pos
  { posX :: Int
  , posY :: Int
  }
  deriving (Show, Eq)

posPair :: Pos -> (Int, Int)
posPair (Pos x y) = (x, y)

data Delta
  = Delta
  { deltaX :: Int
  , deltaY :: Int
  }
  deriving (Show, Eq)

data Dir
  = DirN
  | DirE
  | DirS
  | DirW
  deriving (Show, Eq, Enum)

dirDelta :: Dir -> Delta
dirDelta DirN = Delta 0 (-1)
dirDelta DirE = Delta 1 0
dirDelta DirS = Delta 0 1
dirDelta DirW = Delta (-1) 0

(.+) :: Pos -> Delta -> Pos
Pos x y .+ Delta dx dy = Pos (x + dx) (y + dy)

type Snake = [Pos]

snakeHeadPair :: Snake -> (Int, Int)
snakeHeadPair = posPair . head

data Block
  = Empty
  | Wall
  | Snake
  | Food Int
  deriving (Show, Eq)

type StageSize = (Int, Int)

type Stage = Array StageSize Block

(.%) :: Pos -> StageSize -> Pos
Pos x y .% (w, h) = Pos (f x w) (f y h)
  where
    f i r = (i + r) `mod` r

moveSnake :: StageSize -> Dir -> Snake -> Snake
moveSnake sb dir snake = (head snake .+ dirDelta dir .% sb) : (init snake)

growSnake :: Int -> Snake -> Snake
growSnake n snake = snake ++ replicate n (last snake)

data Game
  = Game
  { gameStage :: Stage
  , gameScore :: Int
  , gameSnake :: Snake
  , gameDir :: Dir
  , gameOver :: Bool
  , gameRandom :: StdGen
  }
  deriving Show

instance RandomGen Game where
  next game =
    let (x, gen') = next $ gameRandom game
      in (x, game { gameRandom = gen' })

  split game =
    let (g1, g2) = split $ gameRandom game
      in (game { gameRandom = g1 }, game { gameRandom = g2 })

  genRange = genRange . gameRandom

gameBounds :: Game -> (StageSize, StageSize)
gameBounds = bounds . gameStage

gameSize :: Game -> StageSize
gameSize game =
  let (w, h) = snd (gameBounds game)
    in (w+1, h+1)

mapGameSnake :: (Snake -> Snake) -> Game -> Game
mapGameSnake f game = game { gameSnake = f (gameSnake game) }

mapGameStage :: (Stage -> Stage) -> Game -> Game
mapGameStage f game = game { gameStage = f (gameStage game) }

changeDir :: Dir -> Game -> Game
changeDir dir game = game { gameDir = dir }

endGame :: Game -> Game
endGame game = game { gameOver = True }

gameGrowSnake :: Int -> Game -> Game
gameGrowSnake n = mapGameSnake (growSnake n)

gameBlock :: Game -> Pos -> Block
gameBlock game pos =
  case gameStage game ! posPair pos of
    Empty | pos `elem` tail (gameSnake game) -> Snake
    block -> block

gameBlockDisplay :: Game -> Pos -> Block
gameBlockDisplay game pos
  | pos `elem` gameSnake game = Snake
  | otherwise = gameStage game ! posPair pos

gameSnakeHead :: Game -> Pos
gameSnakeHead = head . gameSnake

gameSnakeHeadBlock :: Game -> Block
gameSnakeHeadBlock game = gameBlock game (gameSnakeHead game)

gamePut :: Pos -> Block -> Game -> Game
gamePut pos block = mapGameStage (// [(posPair pos, block)])

gameAddFood :: Int -> Pos -> Game -> Game
gameAddFood n pos = gamePut pos (Food n)

gameRandomPos :: Game -> (Pos, Game)
gameRandomPos game =
  let (x, game') = randomR (0, w-1) game
      (y, game'') = randomR (0, h-1) game'
    in (Pos x y, game'')
  where
    (w, h) = gameSize game

-- should be Maybe
gameRandomEmptyPos :: Game -> (Pos, Game)
gameRandomEmptyPos game =
  let (pos, game') = gameRandomPos game
    in
      case gameBlockDisplay game pos of
        Empty     -> (pos, game')
        otherwise -> gameRandomEmptyPos game'

gameAddRandomFood :: Game -> Game
gameAddRandomFood game =
  let (pos, game') = gameRandomEmptyPos game
    in gameAddFood 3 (traceShowId pos) game'

gameClearAtHead :: Game -> Game
gameClearAtHead game = gamePut (gameSnakeHead game) Empty game

tickSnake :: Game -> Game
tickSnake game = mapGameSnake (moveSnake (gameSize game) (gameDir game)) game

tickCollision :: Game -> Game
tickCollision game =
  case gameSnakeHeadBlock game of
    Wall  -> endGame game
    Snake -> endGame game
    Empty  -> game
    Food n -> gameAddRandomFood . gameClearAtHead . gameGrowSnake n $ game

tick :: Game -> Game
tick game
  | gameOver game = game
  | otherwise = tickCollision . tickSnake $ game

emptyStage :: Int -> Int -> Stage
emptyStage w h = listArray ((0, 0), (w-1, h-1)) (repeat Empty)

newGame :: Int -> Int -> StdGen -> Game
newGame w h gen = gameAddFood 3 (Pos 4 0) game
  where
    game =
      Game
      { gameStage = emptyStage w h
      , gameScore = 0
      , gameSnake = [Pos 0 0]
      , gameDir = DirE
      , gameOver = False
      , gameRandom = gen
      }

displayGame :: Game -> String
displayGame game =
  unlines $ map row [0..h-1]
  where
    (w, h) = gameSize game
    row i = map (cell i) [0..w-1]
    cell i j =
      case gameBlockDisplay game (Pos j i) of
        Wall   -> '#'
        Food _ -> '%'
        Snake  -> '@'
        Empty  -> '.'

---

parseDir :: String -> Maybe Dir
parseDir key =
  case key of
    "\ESC[A" -> Just DirN
    "\ESC[B" -> Just DirS
    "\ESC[C" -> Just DirE
    "\ESC[D" -> Just DirW
    _        -> Nothing

getKey :: IO [Char]
getKey = reverse <$> getKey' ""
  where
    getKey' chars = do
      char <- getChar
      more <- hReady stdin
      (if more then getKey' else return) (char:chars)

readDir :: IO (Maybe Dir)
readDir = parseDir <$> getKey

play :: Game -> IO ()
play game
  | gameOver game = putStrLn "You lost ):"
  | otherwise = do
      clearScreen
      putStrLn ""
      putStrLn (displayGame game)
      action <- timeout 200000 readDir
      case action of
        Just (Just dir) -> play (tick . changeDir dir $ game)
        otherwise       -> play (tick game)