{-# LANGUAGE Rank2Types
           , ExistentialQuantification
           , GeneralizedNewtypeDeriving
           , PatternGuards
           , ViewPatterns
  #-}

module Main (main) where

import Control.Monad.State
import Control.Monad.Error

import Data.Char
import Data.List

import System.Console.Haskeline hiding (display)

import Language.PTS.Names
import Language.PTS.Parser
import Language.PTS.Systems
import Language.PTS.TypeCheck
import Language.PTS.SyntaxTree

import Text.PrettyPrint.Leijen

import REPL.CMD
import REPL.Error
import REPL.Environment

main :: IO ()
main = repl

repl :: IO ()
repl = do runInputT defaultSettings $
            evalStateT (runErrorT (unREPL go)
                          >>= either (lift . outputStrLn . unPE) return)
                       (towerPTI, "tower> ")

data PureTypeInterp =
  forall c. (Pretty c, Show c, Eq c) =>
              PTI (PureTypeSystem c) -- type system
                  (Parser c)         -- parser
                  (Environment c)    -- environment


newtype REPL a = REPL { 
    unREPL :: ErrorT PrintedError 
                (StateT (PureTypeInterp, String) (InputT IO)) a
  } deriving ( Functor, Monad, MonadIO
             , MonadState (PureTypeInterp, String)
             , MonadError PrintedError
             )

-- Lift a CMD computation into a REPL computation.
liftCMD :: (forall c. (Pretty c, Show c, Eq c) => CMD c a) -> REPL a
liftCMD cmd = REPL . ErrorT . StateT $ \((PTI pts pars env), prm) ->
  unCMD cmd pts pars (process env) >>= \(e', res) ->
  return (res, (PTI pts pars (plainEnv e'), prm))

outputStrLnR :: String -> REPL ()
outputStrLnR = REPL . lift . lift . outputStrLn

-- Type systems
towerPTI :: PureTypeInterp
towerPTI = PTI tower towerParse []

itowerPTI :: PureTypeInterp
itowerPTI = PTI itower towerParse []

f2PTI :: PureTypeInterp
f2PTI = PTI f2 cocParse []

pf2PTI :: PureTypeInterp
pf2PTI = PTI pf2 cocParse []

fwPTI :: PureTypeInterp
fwPTI = PTI fw cocParse []

pfwPTI :: PureTypeInterp
pfwPTI = PTI pfw cocParse []

lfPTI :: PureTypeInterp
lfPTI = PTI lf cocParse []

lf2PTI :: PureTypeInterp
lf2PTI = PTI lf2 cocParse []

cocPTI :: PureTypeInterp
cocPTI = PTI coc cocParse []

icocPTI :: PureTypeInterp
icocPTI = PTI icoc cocParse []

simpPTI :: PureTypeInterp
simpPTI = PTI simp cocParse []

hmPTI :: PureTypeInterp
hmPTI = PTI hm hmParse []

ihmPTI :: PureTypeInterp
ihmPTI = PTI hm ihmParse []

hhmPTI :: PureTypeInterp
hhmPTI = PTI hhm hmParse []

ihhmPTI :: PureTypeInterp
ihhmPTI = PTI hhm ihmParse []

sysUPTI :: PureTypeInterp
sysUPTI = PTI sysU uParse []

-- The command for switching to a new type system.
newInterp :: String -> PureTypeInterp -> REPL ()
newInterp s pti = put (pti, s)

-- Interpreter commands.
showType :: (Pretty c, Eq c) => String -> CMD c ()
showType = checkEmpty $ prettyCMD . unnotate <=< inferType 
                                             <=< annotate <=< parseTerm

showEvaluatedForm :: (Pretty c, Eq c) => EvalStrategy c -> String -> CMD c ()
showEvaluatedForm strat = checkEmpty $ \s -> do
  tm <- annotate =<< parseTerm s
  _ <- inferType tm
  prettyCMD . unnotate =<< evaluate strat tm

doLet :: (Pretty c, Eq c) => String -> CMD c ()
doLet = checkEmpty $ \s -> do
  (name, mty, stx) <- parseDefn s
  ensureFree name
  tyInf <- inferType =<< annotate stx
  ty <- case mty of
    Nothing  -> return tyInf
    Just ty' -> do tyGiven <- annotate ty'
                   _ <- inferType tyGiven
                   equate "Inferred type does not match expected type"
                          tyInf tyGiven
                   return tyGiven
  pushEnv name stx (unnotate ty)
  writeCMD "Defined."

assume :: (Pretty c, Eq c) => String -> CMD c ()
assume = checkEmpty $ \s -> do
  (n, ty) <- parseAssume s
  ensureFree n
  _ <- inferType =<< annotate ty
  pushEnv n (axiom n) ty
  writeCMD "Assumed."

echo :: (Pretty c, Eq c) => String -> CMD c ()
echo = checkEmpty $ prettyCMD <=< parseTerm

dump :: (Show c) => String -> CMD c ()
dump = checkEmpty $ printCMD <=< parseTerm

dispEnv :: (Pretty c, Eq c) => Bool -> String -> CMD c ()
dispEnv dType _ = do
  env <- plainEnv `fmap` getEnv
  forM_ env $ \(name, (_, ty)) ->
    let str | dType     = show $ pretty (display name)
                                    <+> string ":" <+> pretty ty
            | otherwise = show $ pretty (display name)
    in writeCMD str

-- This is the only command (aside from the trivial system-changing
-- ones, for which that isn't possible) that hasn't been ported to the
-- CMD monad. It could be, but that would disallow files declaring their
-- own type system, if that is desirable. Perhaps at a later date I'll
-- decide to restrict it to CMD, but for now, it stays REPL.
load :: String -> REPL ()
load = mapM_ (uncurry dispatch . slice) . merge . lines <=< liftIO . readFile

merge :: [String] -> [String]
merge [              ] = []
merge ((':':'{':s):ss) = case break (==":}") ss of
  (ls, rest) -> unwords (s:ls) : merge (drop 1 rest)
merge (l:ls)           = l : merge ls

newtype RCMD = RC (forall c. (Show c, Pretty c, Eq c) => String -> CMD c ())

commands :: [(String, RCMD)]
commands = [ (":assume" , RC assume)
           , (":type"   , RC showType)
           , (":tenv"   , RC (dispEnv True))
           , (":env"    , RC (dispEnv False))
           , (":echo"   , RC echo)
           , (":let"    , RC doLet)
           , (":nf"     , RC (showEvaluatedForm nfWithEnv))
           , (":normal" , RC (showEvaluatedForm nfWithEnv))
           , (":whnf"   , RC (showEvaluatedForm whnfWithEnv))
           , (":dump"   , RC dump)
           , (":default", RC echo)
           ]

lookupCommand :: String -> Maybe RCMD
lookupCommand s = aux commands
 where
 aux [] = Nothing
 aux ((n, cmd):cmds)
   | s `isPrefixOf` n = Just cmd
   | otherwise        = aux cmds

systems :: [(String, String, PureTypeInterp)]
systems = [ (":tower"  , "tower> "  , towerPTI)
          , (":itower" , "itower> " , itowerPTI)
          , (":f2"     , "F2> "     , f2PTI)
          , (":pf2"    , "pred F2> ", pf2PTI)
          , (":fw"     , "Fω> "     , fwPTI)
          , (":pfw"    , "pred Fω> ", pfwPTI)
          , (":lf"     , "LF> "     , lfPTI)
          , (":lf2"    , "LF2> "    , lf2PTI)
          , (":coc"    , "coc> "    , cocPTI)
          , (":icoc"   , "icoc> "   , icocPTI)
          , (":simp"   , "simp> "   , simpPTI)
          , (":hhm"    , "hhm> "    , hhmPTI)
          , (":ihhm"   , "ihhm> "   , ihhmPTI)
          , (":hm"     , "hm> "     , hmPTI)
          , (":ihm"    , "ihm> "    , ihmPTI)
          , (":sysu"   , "sysu> "   , sysUPTI)
          ] 

lookupSystem :: String -> Maybe (String, PureTypeInterp)
lookupSystem s = aux systems
 where
 aux [] = Nothing
 aux ((n, prm, sys):syss)
   | s `isPrefixOf` n = Just (prm, sys)
   | otherwise        = aux syss

slice :: String -> (String, String)
slice str@(':':_) = break isSpace str
slice str         = (":default", str)

dispatch :: String -> String -> REPL ()
dispatch []  _     = return ()
dispatch cmd (dropWhile isSpace -> input)
  | cmd `isPrefixOf` ":load" = load input
  | Just (RC exec) <- lookupCommand cmd = liftCMD (exec input) 
                                       `catchError` \(PE s) -> outputStrLnR s
  | Just (prm, sys) <- lookupSystem cmd = newInterp prm sys
  | otherwise = outputStrLnR "Unrecognized command."

getCommand :: String -> REPL (Maybe String)
getCommand prompt = REPL . lift . lift $ do
  line <- getInputLine prompt
  case line of
    Just (':':'{':s) -> aux s
    _                -> return line
 where
 aux s = do line <- getInputLine prompt
            case line of
              Just ":}" -> return $ Just s
              Just i    -> aux (s ++ ' ' : i)
              Nothing   -> aux s

go :: REPL ()
go = do prompt <- gets snd
        line <- getCommand prompt
        case fmap slice line of
          Nothing          -> go
          Just (cmd, rest) ->
            when (not $ length cmd > 1 && cmd `isPrefixOf` ":quit") $
            dispatch cmd rest >> go