module Main where

import Control.Monad
import Data.Maybe
import Haskell.AST
import Haskell.Env
import Haskell.Rewrite
import Haskell.Parser
import Pretty
import System.Environment
import System.Exit
import System.IO


tryEither :: Either String a -> IO a
tryEither (Left err) = die err
tryEither (Right x) = return x

tryEither' :: Show e => Either e a -> IO a
tryEither' (Left err) = die (show err)
tryEither' (Right x) = return x

while :: Monad m => a -> (a -> m (Maybe a)) -> m ()
while val f = f val >>= \case
    Nothing   -> return ()
    Just val' -> while val' f

whenM :: Monad m => m Bool -> m a -> m ()
whenM b a = b >>= \b' -> if b' then void a else return ()

ifM :: Monad m => m Bool -> m a -> m a -> m a
ifM b t e = b >>= \b' -> if b' then t else e

splitOn :: Eq a => a -> [a] -> [[a]]
splitOn _   [] = [[]]
splitOn spl (x:xs)
    | x == spl  = [] : splitOn spl xs
    | otherwise = let (r : rs) = splitOn spl xs in (x:r) : rs


usageString :: String
usageString =
    "\x1B[1mMeta-commands\x1B[0m\n\
    \  env          -- Show the current environment\n\
    \  exit         -- Exit program immediately\n\
    \  undo         -- Undo last operation\n\
    \  focus \x1B[4mx\x1B[0m      -- Watch and operate on given variable\n\
    \  focus        -- Unfocus currently focused variable\n\
    \  log          -- Print list of commands that reproduces current environment\n\
    \  help         -- Show this help\n\
    \\n\
    \\x1B[1mExpression rewrite commands\x1B[0m\n\
    \These commands operate on the currently-focused expression.\n\
    \  rew \x1B[4mx\x1B[0m        -- Replaces occurrences of \x1B[4mx\x1B[0m with its value\n\
    \  beta         -- Applies some beta-reductions ((\\x -> M) y => M[y/x])\n\
    \  eta          -- Applies some eta-reductions ((\\x -> f x) => f)\n\
    \  case         -- Evaluates case statements\n\
    \  case!        -- Evaluates case statements, choosing first hit if multiple branches unify\n\
    \  etacase      -- If all case branches are the same as their patterns, remove case statement\n\
    \  casecase     -- case (case x of {A->a;B->b}) of {C->c;D->d} \n\
    \                    => case x of {A->case a of {C->c;D->d}; B->case b of {C->c;D->d}}\n\
    \  repeat \x1B[4mcmd\x1B[0m   -- Repeat \x1B[4mcmd\x1B[0m until no change occurs\n\
    \\n\
    \To evaluate a rewrite command on an unfocused variable, use:\n\
    \  action \x1B[4mx\x1B[0m \x1B[4mcmd\x1B[0m -- Evaluate \x1B[4mcmd\x1B[0m on variable \x1B[4mx\x1B[0m"


data AppState = AppState { asCtx :: Context, asFocus :: Maybe Name }
  deriving (Show)

data Context = Context { cStack :: [(Action, Env)], cBaseEnv :: Env }
  deriving (Show)

data Action = Action Cmd Name
  deriving (Show)

data Cmd = CRewrite Name
         | CBeta
         | CEta
         | CCase
         | CCaseForce
         | CEtaCase
         | CCaseCase
         | CRepeat Cmd
  deriving (Show)

data UserCmd = UCCmd Cmd
             | CAction Action
             | CExit
             | CUndo
             | CShowEnv
             | CUnFocus
             | CFocus Name
             | CLog
             | CHelp
  deriving (Show)

instance Pretty Action where
    pretty (Action cmd target) = Node ("action " ++ target) [pretty cmd]

instance Pretty Cmd where
    pretty (CRewrite name) = Leaf ("rew " ++ name)
    pretty CBeta = Leaf "beta"
    pretty CEta = Leaf "eta"
    pretty CCase = Leaf "case"
    pretty CCaseForce = Leaf "case!"
    pretty CEtaCase = Leaf "etacase"
    pretty CCaseCase = Leaf "casecase"
    pretty (CRepeat cmd) = Node "repeat" [pretty cmd]

topEnv :: Context -> Env
topEnv (Context [] env) = env
topEnv (Context ((_, env):_) _) = env

pushCtx :: (Action, Env) -> Context -> Context
pushCtx pair ctx = ctx { cStack = pair : cStack ctx }

exprTransformer :: Cmd -> Maybe (Expr -> Expr)
exprTransformer CBeta = Just betared
exprTransformer CEta = Just etared
exprTransformer CCase = Just (casered False)
exprTransformer CCaseForce = Just (casered True)
exprTransformer CEtaCase = Just etacase
exprTransformer CCaseCase = Just casecase
exprTransformer _ = Nothing

repeatTrans :: (Expr -> Expr) -> Expr -> Expr
repeatTrans f expr =
    let values = iterate f expr
        pairs = zip values (tail values)
    in fst . head $ dropWhile (uncurry (/=)) pairs

apply :: Context -> Action -> Either String Context
apply ctx act@(Action cmd target) = case cmd of
    CRewrite name -> genTransform (\env -> liftM2 (,) (get env name) (get env target))
                                  (\(repl, expr) -> normalise $ rewrite name repl expr)
    CRepeat subcmd -> case exprTransformer subcmd of
                          Nothing -> Left $ "Cannot repeat '" ++ pprint 80 subcmd ++ "'"
                          Just trans -> transform (repeatTrans (normalise . trans))
    _ -> transform (normalise . fromJust (exprTransformer cmd))
  where
    transform :: (Expr -> Expr) -> Either String Context
    transform = genTransform (\env -> get env target)

    genTransform :: (Env -> Either String a) -> (a -> Expr) -> Either String Context
    genTransform getter f = do
        expr <- f <$> getter (topEnv ctx)
        env' <- reAdd (topEnv ctx) (Def target expr)
        return $ pushCtx (act, env') ctx

applyUserCmd :: AppState -> UserCmd -> IO (Either String (Maybe AppState))
applyUserCmd appstate = \case
    UCCmd cmd ->
        case asFocus appstate of
            Nothing -> return (Left "No focus set")
            Just focus ->
                case apply (asCtx appstate) (Action cmd focus) of
                    Left err -> return (Left err)
                    Right ctx -> return (Right (Just appstate { asCtx = ctx }))
    CAction act -> case apply (asCtx appstate) act of
                       Left err -> return (Left err)
                       Right ctx -> return (Right (Just appstate { asCtx = ctx }))
    CExit -> return (Right Nothing)
    CUndo ->
        let ctx = asCtx appstate
        in case cStack ctx of
                 [] -> return (Left "Empty environment stack")
                 (_:stk) -> return (Right (Just appstate { asCtx = ctx { cStack = stk } }))
    CShowEnv -> putStrLn (pprint 80 (topEnv (asCtx appstate))) >> return (Right (Just appstate))
    CUnFocus -> return (Right (Just appstate { asFocus = Nothing }))
    CFocus name -> if envContains (topEnv (asCtx appstate)) name
                       then return (Right (Just appstate { asFocus = Just name }))
                       else return (Left "Name doesn't exist")
    CLog -> do
        mapM_ (putStrLn . pprint 80) $ map fst (reverse $ cStack (asCtx appstate))
        return (Right (Just appstate))
    CHelp -> putStrLn usageString >> return (Right (Just appstate))

applyUserCmds :: AppState -> [UserCmd] -> IO (Either String (Maybe AppState))
applyUserCmds appstate [] = return (Right (Just appstate))
applyUserCmds appstate (ucmd : ucmds) =
    applyUserCmd appstate ucmd >>= \case
        Left err -> return (Left err)
        Right Nothing -> return (Right Nothing)
        Right (Just as) -> applyUserCmds as ucmds

interface :: AppState -> IO (Maybe AppState)
interface appstate = do
    putStrLn ""

    let appstate' = case asFocus appstate of
                        Nothing -> appstate
                        Just focus ->
                            if envContains (topEnv (asCtx appstate)) focus
                                then appstate
                                else appstate { asFocus = Nothing }

    case asFocus appstate' of
        Nothing -> return ()
        Just n  ->
            case get (topEnv (asCtx appstate')) n of
                Left _ -> return ()
                Right expr -> putStrLn $ pprint 80 (Def n expr)

    putStr "> " >> hFlush stdout

    ifM isEOF (return Nothing)
              (getLine >>= handleInputLine appstate')

handleInputLine :: AppState -> String -> IO (Maybe AppState)
handleInputLine appstate line = do
    let phrases = map words $ splitOn ';' line

    ucmds <- case mapM parseCmd $ filter (not . null) phrases of
                Left err -> putStrLn err >> return []
                Right ucmds -> return ucmds

    applyUserCmds appstate ucmds >>= \case
        Left err -> putStrLn err >> return (Just appstate)
        Right Nothing -> return Nothing
        Right (Just as) -> return (Just as)

parseCmd :: [String] -> Either String UserCmd
parseCmd ("action" : target : rest) = parseCmd rest >>= \case
    UCCmd cmd -> Right (CAction (Action cmd target))
    _ -> Left "Invalid command within 'action'"
parseCmd ["exit"] = Right CExit
parseCmd ["undo"] = Right CUndo
parseCmd ["env"] = Right CShowEnv
parseCmd ["focus"] = Right CUnFocus
parseCmd ["focus", name] = Right (CFocus name)
parseCmd ["log"] = Right CLog
parseCmd ["help"] = Right CHelp
parseCmd ["rew", func] = Right (UCCmd (CRewrite func))
parseCmd ["beta"] = Right (UCCmd CBeta)
parseCmd ["eta"] = Right (UCCmd CEta)
parseCmd ["case"] = Right (UCCmd CCase)
parseCmd ["case!"] = Right (UCCmd CCaseForce)
parseCmd ["etacase"] = Right (UCCmd CEtaCase)
parseCmd ["casecase"] = Right (UCCmd CCaseCase)
parseCmd ("repeat" : rest) = parseCmd rest >>= \case
    UCCmd (CRepeat _) -> Left "Cannot repeat 'repeat'"
    UCCmd cmd -> Right (UCCmd (CRepeat cmd))
    _ -> Left "Invalid command within 'repeat'"
parseCmd cmd = Left $ "Unrecognised command: " ++ show cmd

main :: IO ()
main = do
    fname <- getArgs >>= \case
        [fname] -> return fname
        _       -> die "Usage: verify-hs <file.hs>"

    source <- readFile fname
    ast <- tryEither' (parseAST fname source)
    print ast
    env0 <- tryEither (envFromAST ast)
    print env0

    while (AppState (Context [] env0) Nothing) interface