{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-tabs #-}
module Hcompta.Calculus.Explicit.REPL where

import Control.Applicative (Applicative(..), (<$>))
import Control.Exception
import Control.Monad
import Control.Monad.State
import Data.Bool
import qualified Data.Char as Char
import Data.Either (Either(..))
import Data.Eq (Eq(..))
import Data.Foldable (Foldable(..))
import Data.Function (($), (.), const, flip)
import Data.Functor.Identity
import Data.List ((++), break, concat, dropWhile, last, reverse)
import Data.List (isPrefixOf, find)
import qualified Data.List as List
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe (Maybe(..))
import Data.Monoid ((<>))
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..))
import Data.String (String)
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Text.Buildable (Buildable(..))
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Builder as Builder
import Data.Tuple (fst)
import Data.Typeable as Typeable
import Prelude (Integer, Num(..), div, error)
import System.Console.Haskeline
import System.Directory
import System.FilePath
import System.IO (IO, readFile)
import qualified Text.Parsec as R
import Text.Show (Show(..))

-- import Debug.Trace

import Calculus.Abstraction.DeBruijn.Generalized
import Calculus.Lambda.Omega.Explicit
import Calculus.Lambda.Omega.Explicit.Read

-- * Type 'REPL'

-- | /Read Eval Print Loop/ monad.
newtype REPL a
 =      REPL
 {    unREPL :: StateT REPL_State (InputT IO) a }
 deriving ( Functor
          , Applicative
          , Monad
          , MonadIO
          , MonadState REPL_State )

data REPL_State
 =   REPL_State
 {   repl_state_env       :: Env
 ,   repl_state_load_dir  :: FilePath
 ,   repl_state_load_done :: Map FilePath ()
 }

main :: IO ()
main = do
	cwd <- getCurrentDirectory
	runInputT defaultSettings $
		evalStateT (unREPL main_loop) $
		REPL_State
		 { repl_state_env       = prelude axioms
		 , repl_state_load_dir  = cwd
		 , repl_state_load_done = mempty
		 }
	where
	main_loop :: REPL ()
	main_loop = do
		let prompt = "> "
		line <- repl_read_command prompt
		case slice <$> line of
		 Just (cmd, input) | (not (length cmd > 1 && cmd `isPrefixOf` "quit")) ->
			dispatch cmd input
		 _ -> return ()
		main_loop

slice :: String -> (String, String)
slice (':':str) = break Char.isSpace str
slice str       = ("", str)


-- ** I/O
print :: Buildable a => a -> TL.Text
print = Builder.toLazyText . build

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

repl_write_string :: String -> REPL ()
repl_write_string = REPL . lift . outputStr

repl_show :: Show a => a -> REPL ()
repl_show = repl_write_string_ln . show

repl_print :: Buildable a => a -> REPL ()
repl_print = repl_write_string . TL.unpack . print

repl_print_ln :: Buildable a => a -> REPL ()
repl_print_ln = repl_write_string_ln . TL.unpack . print

repl_read :: Parser Identity x -> String -> (x -> REPL ()) ->  REPL ()
repl_read _p [] _k = return ()
repl_read p s k =
	case runIdentity $ read p s of
	 Right x -> k x
	 Left err -> do
		repl_write_string_ln "Parsing error:"
		repl_show err
		repl_write_string " Command_Input: "
		repl_write_string_ln s

-- ** Commands
repl_read_command :: String -> REPL (Maybe String)
repl_read_command prompt = do
	line <- REPL . lift $ getInputLine prompt
	case line of
	 Just s@(':':_) -> return $ Just s
	 Just l@(_:_) | last l == ' ' -> parse_block l
	 Just s -> return $ Just s
	 Nothing -> return Nothing
	where
	parse_block :: String -> REPL (Maybe String)
	parse_block s = do
		line <- REPL . lift $ getInputLine ""
		case line of
		 Just l@(_:_) | last l == ' ' -> parse_block (s ++ '\n' : l)
		 _ -> return $ Just s

type Command_Name = String
type Command = Command_Input -> IO Command_Output

command :: Command_Name -> Command
command "" = command "let"
command cmd =
	case find (\p -> cmd `isPrefixOf` fst p) commands of
	 Nothing -> (\_ -> return $ Left $ Error_Command cmd)
	 Just (_, c) -> c
	where
	commands :: [(String, Command)]
	commands =
	 [ ("assume", command_assume) -- Type -> IO ()
	 , ("code" , command_code) -- Var -> Term
	 , ("dump" , command_dump) -- Term -> Text
	 , ("equiv", command_equiv) -- Term -> Term -> Bool
	 -- , ("echo" , command_echo) -- Term -> Term
	 , ("let"  , command_let) -- Var -> Term -> IO ()
	 , ("load" , command_load) -- FilePath -> IO ()
	 , ("nf"   , command_normalize normalize) -- Term -> Term
	 , ("nf_dump", command_normalize_dump normalize) -- Term -> Term
	 , ("reset", command_reset) -- () -> IO ()
	 , ("type" , command_type) -- Term -> Type
	 , ("type_dump", command_type_dump) -- Term -> Text
	 , ("whnf" , command_normalize whnf) -- Term -> Term
	 ]
command_run :: Command_Name -> Command
command_run cmd (Command_Input i st) = do
	command cmd $
		Command_Input (strip_spaces i) st
	where
	strip_spaces =
		reverse . dropWhile Char.isSpace .
		reverse . dropWhile Char.isSpace

dispatch :: Command_Name -> String -> REPL ()
dispatch cmd i = do
	st <- get
	output <- liftIO $ command_run cmd (Command_Input i st)
	case output of
	 Left (err::Error) ->
		repl_print err
	 Right (Command_Result msg new_st) -> do
		case TL.unpack $ print msg of
		 [] -> return ()
		 o -> repl_write_string_ln o
		put new_st

data Command_Input
 =   Command_Input String REPL_State
type Command_Output
 = Either Error Command_Result
data Error
 =   Error_Parse String R.ParseError
 |   Error_Type Type_Error
 |   Error_Let  Var_Name Type_Error
 |   Error_Code Var_Name
 |   Error_Command Command_Name
 |   Error_IO IOException
 |   Error_Load FilePath Error
 deriving (Show)
instance Buildable Error where
	build err =
		case err of
		 Error_Parse s e -> "Error: parsing: " <> build s <> "\n" <> build (show e) <> "\n"
		 Error_Type e -> build e
		 Error_Let var e -> "Error: in let: " <> build var <> "\n " <> build e
		 Error_Code var -> "Error: no such variable in environment: " <> build var
		 Error_Command cmd -> "Error: unrecognized command: " <> build cmd <> "\n"
		 Error_IO e -> "Error: " <> build (show e) <> "\n"
		 Error_Load file e ->
			"Error: loading: " <> build file <> "\n"
			<> build e
data Command_Result
 =   forall msg.
     ( Buildable msg
     ) => Command_Result msg REPL_State

command_assume :: Command
command_assume (Command_Input str st) =
	return $
	let ctx = context_from_env $ repl_state_env st in
	case runIdentity $ read parse_assume str of
	 Left err -> Left $ Error_Parse str err
	 Right (v, ty) ->
		case type_of ctx ty of
			Left err -> Left $ Error_Type err
			Right _ty_ty -> Right $
				Command_Result (""::Text) $
					st{repl_state_env = env_insert v
						 (TeTy_Axiom (Axiom_Type_Assume ty)) ty $
						repl_state_env st
					 }

command_code :: Command
command_code (Command_Input str st) =
	return $
	let env = repl_state_env st in
	case runIdentity $ read parse_var_name str of
	 Left err -> Left $ Error_Parse str err
	 Right var ->
		case Map.lookup var env of
		 Nothing -> Left $ Error_Code var
		 Just item -> Right $ Command_Result
			 (form_given $ env_item_term item)
			 st

command_dump :: Command
command_dump (Command_Input str st) =
	return $
		case runIdentity $ read parse_term str of
		 Left err -> Left $ Error_Parse str err
		 Right te -> Right $ Command_Result (show te) st

command_let :: Command
command_let (Command_Input [] st) =
	return $ Right $ Command_Result (""::Text) st
command_let (Command_Input str st) =
	let toks_or_err = runIdentity $ lex lex_all str in
	case toks_or_err of
	 Left err -> return $ Left $ Error_Parse str err
	 Right [] -> return $ Right $ Command_Result (""::Text) st
	 Right toks ->
		case runIdentity $ parse parse_let_or_term toks of
		 Left err -> return $ Left $ Error_Parse str err
		 Right let_or_term -> do
			let ctx = context_from_env $ repl_state_env st
			case let_or_term of
			 Left (v, mty, te) ->
				let ety = case mty of
					 Nothing -> type_of ctx te
					 Just ty -> do
						_ty_ty <- type_of ctx ty
						const ty <$> check ctx (context_apply ctx ty) te
				 in
				return $
				case ety of
				 Left err -> Left $ Error_Let v err
				 Right ty -> do
					Right $ Command_Result (build v <> " : " <> build ty) $
						st{repl_state_env = env_insert v te ty $ repl_state_env st}
			 Right let_te ->
				case type_of ctx let_te of
				 Left err -> return $ Left $ Error_Type err
				 Right _ty ->
					let norm_te = normalize ctx let_te in
					case norm_te of
					 TeTy_Axiom (axiom_cast -> Just (Axiom_Term (io::IO (Term Var_Name)) _o_ty)) -> do
						io_te <- io
						return $ Right $ Command_Result io_te st
					 _ ->
						return $ Right $ Command_Result norm_te st

command_load :: Command
command_load (Command_Input file input_st) = do
	err_or_io <- try $ do
		path <- canonicalizePath (repl_state_load_dir input_st </> file)
		content <- readFile path
		return (path, content)
	case err_or_io of
	 Left (err::IOException) -> return $ Left $ Error_IO err
	 Right (modul, content) -> do
		case Map.lookup modul $ repl_state_load_done input_st of
		 Just _ -> return $ Right $
			Command_Result ("Module already loaded: " ++ modul) input_st
		 _ ->
			case runIdentity $ R.runParserT (parse_commands <* R.eof) () modul content of
			 Left err -> return $ Left $ Error_Load modul $ Error_Parse content err
			 Right cmds -> do
				let old_dir = repl_state_load_dir input_st
				err_or_st <-
					foldM (\err_or_st (cmd, i) ->
						case err_or_st of
						 Left _ -> return $ err_or_st
						 Right (Command_Result last_msg last_st) -> do
							o <- command_run cmd (Command_Input i last_st)
							case o of
							 Left _ -> return o
							 Right (Command_Result msg running_st) ->
								return $ Right $
									Command_Result (build last_msg <> "\n" <> build msg) running_st
					 )
					 (Right $ Command_Result ("Loading: " <> modul)
						 input_st{repl_state_load_dir = takeDirectory modul})
					 ({-trace ("cmds: " ++ show cmds)-} cmds)
				return $
					case err_or_st of
					 Left err -> Left $ Error_Load modul $ err
					 Right (Command_Result msg result_st) ->
						Right $ Command_Result msg $ result_st
						 { repl_state_load_dir  = old_dir
						 , repl_state_load_done =
							Map.insert modul () $
							repl_state_load_done result_st
						 }

command_normalize
 :: (Context Var_Name -> Term Var_Name -> Term Var_Name)
 -> Command
command_normalize norm (Command_Input str st) = do
	let ctx = context_from_env $ repl_state_env st
	case runIdentity $ read parse_term str of
	 Left err -> return $ Left $ Error_Parse str err
	 Right (te::Term Var_Name) ->
		let n_te = norm ctx te in
		case n_te of
		 TeTy_Axiom (axiom_cast -> Just (Axiom_Term (o::IO (Term Var_Name)) _o_ty)) -> do
			r <- o
			return $ Right $ Command_Result r st
		 TeTy_Axiom (axiom_cast -> Just (Axiom_Term (o::IO Text) _o_ty)) -> do
			r <- o
			return $ Right $ Command_Result r st
		 _ ->
			return $ Right $ Command_Result n_te st

command_equiv :: Command
command_equiv (Command_Input str st) = do
	let ctx = context_from_env $ repl_state_env st
	return $
		case runIdentity $ read ((,) <$> parse_term <* (parse_token Token_Equal >> parse_token Token_Equal) <*> parse_term) str of
		 Left err -> Left $ Error_Parse str err
		 Right (x_te, y_te) -> Right $
			Command_Result (if equiv ctx x_te y_te then "True" else "False"::Text) st

command_normalize_dump
 :: (Context Var_Name -> Term Var_Name -> Term Var_Name)
 -> Command
command_normalize_dump norm (Command_Input str st) = do
	let ctx = context_from_env $ repl_state_env st
	return $
		case runIdentity $ read parse_term str of
		 Left err -> Left $ Error_Parse str err
		 Right te -> Right $ Command_Result (show $ norm ctx te) st

command_reset :: Command
command_reset (Command_Input _ st) =
	return $ Right $
		Command_Result (""::Text) st{repl_state_env = mempty}

command_type :: Command
command_type (Command_Input [] st) = do
	let env = repl_state_env st
	return $
		Right $ Command_Result
		 (
			foldr (flip (<>)) "" $
			List.intersperse "\n" $
			 (\(name, item) ->
				build name <> " : "
				 <> build (form_given $ env_item_type item))
			 <$> (Map.toList env)
		 )
		 st
command_type (Command_Input str st) = do
	let ctx = context_from_env $ repl_state_env st
	return $
		case runIdentity $ read parse_term str of
		 Left err -> Left $ Error_Parse str err
		 Right te ->
			case type_of ctx te of
			 Left err -> Left $ Error_Type err
			 Right ty -> Right $ Command_Result (normalize ctx ty) st

command_type_dump :: Command
command_type_dump (Command_Input [] st) = do
	let env = repl_state_env st
	return $
		Right $ Command_Result
		 (
			foldr (flip (<>)) "" $
			List.intersperse "\n" $
			 (\(name, item) ->
				build name <> " : "
				 <> build (show $ form_given $ env_item_type item))
			 <$> (Map.toList env)
		 )
		 st
command_type_dump (Command_Input str st) = do
	let ctx = context_from_env $ repl_state_env st
	return $
		case runIdentity $ read parse_term str of
		 Left err -> Left $ Error_Parse str err
		 Right te ->
			case type_of ctx te of
			 Left err -> Left $ Error_Type err
			 Right ty -> Right $ Command_Result (show $ normalize ctx ty) st

{-
command_echo :: String -> REPL ()
command_echo str =
	repl_read parse_term str repl_print_ln
-}

-- * Builtins

builtin :: Axioms -> [Text] -> Env
builtin =
	foldl $ \env str ->
		let ctx = context_from_env env in
		read_string parse_let str $ \(v, mty, te) ->
			let ety = case mty of
				 Just ty -> do
					_ty_ty <- type_of ctx ty
					const ty <$> check ctx (context_apply ctx ty) te
				 Nothing -> type_of ctx te in
			case ety of
			 Left err -> error $ show err
			 Right ty -> env_insert v te ty env
	where
	read_string p s k =
		case runIdentity $ read p s of
		 Right x -> k x
		 Left err -> error $
			 "Parsing_error:\n" ++ show err
			 ++ " Input: " ++ Text.unpack s

axioms :: Axioms
axioms =
	(axioms_monopoly <>) $
	(Map.fromList axioms_io <>) $
	(Map.fromList axioms_int <>) $
	(Map.fromList axioms_text <>) $
	Map.fromList
	 [ ("Unit", item $ Axiom_Type_Unit)
	 ]
	where
	item :: (Axiomable (Axiom ax), Typeable ax) => Axiom ax -> Env_Item
	item = env_item_from_axiom context
	{-
	axioms_arr =
	 [ -- ("Arr", item $ Axiom_Term_Abst)
	 ]
	-}
	axioms_text =
	 [ ("Text"      , item $ Axiom_Type_Text)
	 , ("text_empty", item $ axiom_term (""::Text))
	 , ("text_hello", item $ axiom_term ("Hello World!"::Text))
	 ] 
	axioms_int =
	 [ ("Int"      , item $ Axiom_Type_Integer)
	 , ("int_zero" , item $ axiom_term (0::Integer))
	 , ("int_one"  , item $ axiom_term (1::Integer))
	 , ("int_add"  , item $ axiom_term ((+)::Integer -> Integer -> Integer))
	 , ("int_neg"  , item $ axiom_term (negate::Integer -> Integer))
	 , ("int_sub"  , item $ axiom_term ((-)::Integer -> Integer -> Integer))
	 , ("int_mul"  , item $ axiom_term ((*)::Integer -> Integer -> Integer))
	 , ("int_div"  , item $ axiom_term (div::Integer -> Integer -> Integer))
	 ]
	axioms_io =
	 [ ("IO"       , item $ (Axiom_Type_IO::Axiom (IO A)))
	 , ("return_io", item $ return_io)
	 , ("bind_io"  , item $ bind_io)
	 , ("join_io"  , item $ join_io)
	 -- , ("return_io_text", item $ axiom_term $ (return::Text -> IO Text))
	 -- , ("return_io_int", item $ Axiom_Term $ (return::Int -> IO Int))
	 ]
		where
		-- | @return_io : ∀(A:*) -> A -> IO A@
		return_io :: Axiom (Axiom_Type_Abst)
		return_io =
			Axiom_Type_Abst "A" (TeTy_Axiom . ax_te) ax_ty
			where
				ax_te :: Type Var_Name
				 -> Axiom (Axiom_Term (Term Var_Name -> IO (Term Var_Name)))
				ax_te ty = Axiom_Term return (\ctx -> context_lift ctx <$> ty)
				ax_ty :: Abstraction (Suggest Var_Name) Type Var_Name
				ax_ty =
					(("A" =?) `abstract`) $
					axiom_type_of context $
					axiom_term (return::A -> IO A)
		
		-- | @bind_io : ∀(A:*) -> ∀(B:*) -> IO A -> (A -> IO B) -> IO B@
		bind_io :: Axiom Axiom_Type_Abst
		bind_io = ax1_te
			where
				ax1_te :: Axiom Axiom_Type_Abst
				ax1_te =
					Axiom_Type_Abst "A"
					 (\(Type_Abst _ _ ty_a_abst) ->
						TeTy_Axiom $
						Axiom_Type_Abst "B"
						 (TeTy_Axiom . ax0_te)
						 ty_a_abst)
					 ax1_ty
				ax1_ty =
					(("A" =?) `abstract`) $
					Type_Abst "B"
					 (Type_Sort (Type_Level_0, Type_Morphism_Mono)) $
					 ax0_ty
				
				ax0_te
				 :: Type Var_Name
				 -> Axiom (Axiom_Term
					 (  IO (Term Var_Name)
					 -> (Term Var_Name -> IO (Term Var_Name))
					 -> IO (Term Var_Name)
					 ))
				ax0_te ty =
					Axiom_Term (>>=) (\ctx -> context_lift ctx <$> ty)
				ax0_ty =
					(("B" =?) `abstract`) $
					axiom_type_of context $
					axiom_term ((>>=)::IO A -> (A -> IO B) -> IO B)
		
		-- | @join_io : ∀(A:*) -> IO (IO A) -> IO A@
		join_io :: Axiom Axiom_Type_Abst
		join_io = ax1_te
			where
				ax1_te :: Axiom Axiom_Type_Abst
				ax1_te =
					Axiom_Type_Abst "A"
					 (\ty_a ->
						TeTy_Axiom $
						Axiom_Term
						 (join::IO (IO (Term Var_Name)) -> IO (Term Var_Name))
						 (\ctx -> context_lift ctx <$> ty_a))
					 ax1_ty
				ax1_ty =
					(("A" =?) `abstract`) $
					axiom_type_of context $
					axiom_term (join::IO (IO A) -> IO A)

prelude :: Axioms -> Env
prelude axs =
	builtin axs $ concat $ (Text.unlines <$>) <$>
	 [ prelude_bool
	 , prelude_either
	 ]
	where
	prelude_bool :: [[Text]]
	prelude_bool =
	 [ ["Bool_Polytype : *p = (R:*) -> R -> R -> R"]
	 , ["Bool  : *m   = Monotype Bool_Polytype"]
	 , ["True  : Bool = monotype Bool_Polytype (λ(R:*) (True:R) (False:R) -> True)"]
	 , ["False : Bool = monotype Bool_Polytype (λ(R:*) (True:R) (False:R) -> False)"]
	 , [ "eq (x:Bool) (y:Bool) : Bool"
	   , "  = monotype Bool_Polytype"
	   , "    (λ(R:*) (True:R) (False:R) ->"
	   , "      polytype Bool_Polytype x R"
	   , "       (polytype Bool_Polytype y R True  False)"
	   , "       (polytype Bool_Polytype y R False True)"
	   , "    )"
	   ]
	 , [ "and (x:Bool) (y:Bool) : Bool"
	   , " = monotype Bool_Polytype"
	   , "   (λ(R:*) (True:R) (False:R) ->"
	   , "     polytype Bool_Polytype x R"
	   , "      (polytype Bool_Polytype y R True  False)"
	   , "      (polytype Bool_Polytype y R False False)"
	   , "   )"
	   ]
	 , [ "or (x:Bool) (y:Bool) : Bool"
	   , " = monotype Bool_Polytype"
	   , "   (λ(R:*) (True:R) (False:R) ->"
	   , "     polytype Bool_Polytype x R"
	   , "      (polytype Bool_Polytype y R True  True)"
	   , "      (polytype Bool_Polytype y R True False)"
	   , "   )"
	   ]
	 , [ "xor (x:Bool) (y:Bool) : Bool"
	   , " = monotype Bool_Polytype"
	   , "   (λ(R:*) (True:R) (False:R) ->"
	   , "     polytype Bool_Polytype x R"
	   , "      (polytype Bool_Polytype y R False  True)"
	   , "      (polytype Bool_Polytype y R True False)"
	   , "   )"
	   ]
	 , [ "not (x:Bool) : Bool"
	   , " = monotype Bool_Polytype"
	   , "   (λ(R:*) (True:R) (False:R) ->"
	   , "     polytype Bool_Polytype x R False True)"
	   ]
	 ]
	prelude_either :: [[Text]]
	prelude_either =
	 [ ["Either_Polytype (A:*) (B:*) : *p = (R:*) -> (A -> R) -> (B -> R) -> R"]
	 , ["Either  (A:*) (B:*) : *m = Monotype (Either_Polytype A B)"]
	 , [ "Left (A:*) (B:*) (x:A)"
	   , "     : Either A B"
	   , "     = monotype (Either_Polytype A B)"
	   , "       (λ(R:*) (Left:A -> R) (Right:B -> R) -> Left x)"
	   ]
	 , [ "Right (A:*) (B:*) (x:B)"
	   , "       : Either A B"
	   , "       = monotype (Either_Polytype A B)"
	   , "         (λ(R:*) (Left:A -> R) (Right:B -> R) -> Right x)"
	   ]
	 , [ "either (A:*) (B:*) (R:*) (l:A -> R) (r:B -> R) (e:Either A B) : R"
	   , "   = polytype (Either_Polytype A B) e R l r"
	   ]
	 ]