import System.IO

data Player = X | O deriving (Eq, Show)
type Grid = [[Maybe Player]]

initialGrid :: Grid
initialGrid = [[Nothing, Nothing, Nothing],
               [Nothing, Nothing, Nothing],
               [Nothing, Nothing, Nothing]]

next :: Player -> Player
next X = O
next O = X

diagonal1 :: Grid -> [Maybe Player]
diagonal1 g = [head (head g), g !! 1 !! 1, g !! 2 !! 2]

diagonal2 :: Grid -> [Maybe Player]
diagonal2 g = [head g !! 2, g !! 1 !! 1, head (g !! 2)]

rows :: Grid -> [[Maybe Player]]
rows g = g

columns :: Grid -> [[Maybe Player]]
columns = transpose

transpose :: Grid -> Grid
transpose g = [[g !! j !! i | j <- [0..2]] | i <- [0..2]]

wonBy :: Grid -> Maybe Player
wonBy g = case filter (all (== Just X)) (rows g ++ columns g ++ [diagonal1 g, diagonal2 g]) of
            [_] -> Just X
            _ -> case filter (all (== Just O)) (rows g ++ columns g ++ [diagonal1 g, diagonal2 g]) of
                   [_] -> Just O
                   _ -> Nothing

play :: Grid -> Player -> IO ()
play g p = do putStrLn (render g)
              if wonBy g == Just X then putStrLn "X won!"
              else if wonBy g == Just O then putStrLn "O won!"
              else if all (all isJust) g then putStrLn "It's a draw!"
              else do putStr "Enter row: "
                      hFlush stdout
                      row <- readLn
                      putStr "Enter column: "
                      hFlush stdout
                      col <- readLn
                      let g' = update g row col p
                      if g' == g then do
                                        putStrLn "Invalid move, try again."
                                        play g p
                      else play g' (next p)

update :: Grid -> Int -> Int -> Player -> Grid
update g row col p =
  let (before, this:after) = splitAt row g
      (before', _:after') = splitAt col this
  in before ++ [before' ++ [Just p] ++ after'] ++ after

isJust :: Maybe a -> Bool
isJust (Just _) = True
isJust _ = False

render :: Grid -> String
render g = unlines [[case g !! i !! j of
                       Just X -> 'X'
                       Just O -> 'O'
                       Nothing -> ' ' | j <- [0..2]] | i <- [0..2]]

main :: IO ()
main = do putStrLn "Tic Tac Toe!"
          play initialGrid X