module Main where

import Control.Monad.State
import Prelude hiding (Word, null)
import ListT (fromFoldable, ListT, toList, fold)
import Text.Read (readMaybe)
import qualified Data.Map as M
import Text.Parsec.String (Parser)
import Text.Parsec hiding (State)
import Data.Maybe (catMaybes)
import Data.Either (partitionEithers)
import qualified Text.Parsec.Token as P
import Text.Parsec.Language (haskellDef)
import Data.List (delete, null)
import System.Environment

data Direction = U | L | R | D deriving (Show, Eq)
type Word = Either String Integer
type Position = (Int, Int)
type Opcode = (Char, Char)
data Car = Car Position Direction Word deriving Show
type Maze = [String]

data Condition = Signal | Cmp Word deriving Show
data Function = Eq Word
                | Dec Integer 
                | Inc 
                | Dir Direction
                | If Condition Function Function deriving Show
type Env = M.Map Opcode Function 

advance' :: Maze -> Car -> Car
-- so this is a little weird
-- this is how the "canonical" Maze implementation turns cars around on walls
-- its weird, i would have preferred fixed clockwise/counterclockwise
advance' m (Car pos dir v) = head $ filter
    ((('#','#')/=) . tile m)
    (step . flip (Car pos) v <$> moveDesired [U, D, L, R]) where
  -- prefer to move forward, turn around only if you have to
  moveDesired = (dir:) . delete dir
                . (++ [opposite dir]) . delete (opposite dir)
  opposite U = D
  opposite D = U
  opposite R = L
  opposite L = R

  step (Car (x,y) U v) = Car (x,y-1) U v
  step (Car (x,y) D v) = Car (x,y+1) D v
  step (Car (x,y) L v) = Car (x-3,y) L v
  step (Car (x,y) R v) = Car (x+3,y) R v

tile :: Maze -> Car -> Opcode
tile m (Car (x,y) _ _) = (row !! x, row !! (x+1)) where
  row = m !! y

act :: Function -> Car -> Car
act (Eq v) (Car pos dir _) = Car pos dir v 
act (If (Cmp v') lhs rhs) car@(Car _ _ v) = if v' == v then act lhs car else act rhs car 
act (Dir dir) (Car pos _ v) = Car pos dir v
act (Dec x) (Car pos dir (Right y)) = Car pos dir (Right $ y - x)
  
maze :: Env -> Maze -> [Car] -> ListT IO Car
maze env m cars = do
  let signalled' = any (\car -> tile m car == ('*', '*')) cars
  car@(Car pos dir v) <- fromFoldable cars
  let advance = advance' m;
  case tile m car of
    ('.', '.') -> return $ advance car
    ('^', '^') -> return $ advance car
    ('*', '*') -> return $ advance car
    ('(', ')') -> mempty 
    ('>', '>') -> do 
      lift $ case v of
        Right v' -> putStr $ show v' ++ " "
        Left v' -> putStrLn v'
      return $ advance car 
    ('<', '<') -> do
      lift $ putStr "> "
      line <- lift getLine
      case readMaybe line :: Maybe Integer of
        Just v' -> return $ advance (Car pos dir (Right v'))
        Nothing -> return $ advance (Car pos dir (Left line))
    ('<', '>') -> fromFoldable [advance $ Car pos L v, advance $ Car pos R v]
    ('%', 'D') -> return $ advance (Car pos D v)
    ('%', 'U') -> return $ advance (Car pos U v)
    ('%', 'L') -> return $ advance (Car pos L v)
    ('%', 'R') -> return $ advance (Car pos R v)
    unk -> case (M.!?) env unk of
      Just fn -> case fn of
        If Signal lhs rhs -> fromFoldable $ if signalled' then [advance $ act lhs car] else [advance $ act rhs car]
        fn' -> return $ advance (act fn' car)
      Nothing -> error $ show unk

assignment :: Parser Function
assignment = string "=" >> Eq <$> (Left <$> stringLiteral' <|> Right <$> decimal') where
  def = P.makeTokenParser haskellDef
  stringLiteral' = P.stringLiteral def
  decimal' = P.decimal def

condition :: Parser Condition
condition = (string "**" *> return Signal) <|> (string "== " >> Cmp <$> (Right . read <$> many1 digit))

ifStatement :: Parser Function
ifStatement = string "IF " >> If <$> condition <* string " THEN " <*> (assignment <|> direction) <* string "ELSE " <*> (assignment <|> direction)

direction :: Parser Function
direction = char '%' >> Dir <$> choice [char 'D' >> return D, char 'U' >> return U, char 'L' >> return L, char 'R' >> return R] <* spaces

decrement = string "-= " >> Dec <$> P.decimal (P.makeTokenParser haskellDef) <* spaces 
  
parseMaze = catMaybes <$> many1 maybeFunction where
  maybeFunction = (Just . Left <$> function) <|> (Just . Right <$> mazeStart)
  mazeStart = catMaybes <$> manyTill (Just <$> (try (string "^^") *> getPosition) <|> (anyChar >> return Nothing)) newline
  function :: Parser (Opcode, Function)
  function = (,) <$> ((,) <$> letter <*> letter) <* string "-> " <*> (assignment <|> ifStatement <|> decrement)

maze' :: Env -> Maze -> StateT [Car] IO ()
maze' env m = do
  cars <- get
  cars' <- lift $ toList $ maze env m cars
  put cars'
  if null cars'
    then return ()
    else maze' env m

runMaze m = do
  let (Right (fns, starts)) = partitionEithers <$> parse (parseMaze <* eof) "" m
  let starts' = [Car (sourceColumn x - 3, sourceLine x - 1) U (Right 0) | x <- concat starts]
  let env = M.fromList fns
  let m' = lines m
  execStateT (maze' env m') starts'
  return ()

main :: IO ()
main = getArgs >>= readFile . head >>= runMaze