-- Interpreter for the Lamb programming language
-- Copyright (c) 2013 darkf
-- Licensed under the terms of the zlib license, see LICENSE for details

module Interp where
import Prelude hiding (lookup)
import qualified Data.Map as M
import qualified Data.ByteString.Char8 as BSC
import qualified Network.Socket as SO
import Data.List (intercalate)
import Control.Exception (try, SomeException)
import Control.Monad.IO.Class (liftIO)
import Control.Monad.Trans.State (StateT, runStateT, evalStateT, get, put)
import System.IO (Handle, hPutStr, hGetLine, hClose, hIsEOF, hSetBuffering,
		hSetBinaryMode, openBinaryFile, IOMode(..), BufferMode(NoBuffering), stdout, stdin)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import System.Directory (doesFileExist)
import System.FilePath (FilePath, splitExtension, takeBaseName, takeDirectory, (</>))
import System.Environment (getExecutablePath)
import AST
import Parser (parseProgram)

-- for Show
newtype BIF = BIF (Value -> InterpState Value)
instance Show BIF where show _ = "<built-in>"
instance Eq BIF where a == b = False
instance Ord BIF where compare a b = if a == b then EQ else LT

data Value = IntV Integer
		   | StrV String
		   | BoolV Bool
		   | StreamV Handle
		   | TupleV [Value]
		   | ListV [Value]
		   | DictV (M.Map Value Value)
		   | RefV (IORef Value)
		   | Builtin BIF
		   | FnV Env [(Pattern, AST)] -- closure pattern->body bindings
		   deriving (Eq)

instance Ord Value where
	compare (IntV a) (IntV b) = compare a b
	compare (StrV a) (StrV b) = compare a b
	compare (BoolV a) (BoolV b) = compare a b
	compare (TupleV a) (TupleV b) = compare a b
	compare (ListV a) (ListV b) = compare a b
	compare (StreamV a) (StreamV b) = if a == b then EQ else LT
	compare (Builtin a) (Builtin b) = compare a b
	compare (FnV a b) (FnV x y) = if a == x && b == y then EQ else LT
	compare (DictV a) (DictV b) = compare a b
	compare _ _ = error "compare: not valid"

type Env = [M.Map String Value] -- lexical environment (linked list)
type InterpState = StateT Env IO -- interpreter state (open handles, global env)

emptyEnv = [M.empty]
unitv = TupleV []

-- look up a binding from the bottom up
lookup :: Env -> String -> Maybe Value
lookup [] _ = Nothing
lookup (env:xs) name = maybe (lookup xs name) Just (M.lookup name env)

-- bind in the local environment
bind :: Env -> String -> Value -> Env
bind (env:xs) name value = (M.insert name value env):xs

instance Show Value where
	show (IntV i) = show i
	show (StrV s) = show s
	show (BoolV b) = show b
	show (TupleV []) = "(,)"
	show (TupleV v) = "(" ++ intercalate "," (map show v) ++ ")"
	show (ListV v) = show v
	show (DictV d) = "<dict " ++ show d ++ ">"
	show (FnV _ _) = "<fn>"
	show (StreamV _) = "<stream>"
	show (Builtin _) = "<built-in>"
	show (RefV _) = "<ref>"

-- value operators
(IntV l) +$ (IntV r) = IntV (l + r)
(StrV l) +$ (StrV r) = StrV (l ++ r)
(ListV l) +$ (ListV r) = ListV (l ++ r)
l +$ r = error $ "cannot + " ++ show l ++ " and " ++ show r

(IntV l) -$ (IntV r) = IntV (l - r)
l -$ r = error $ "cannot - " ++ show l ++ " and " ++ show r

(IntV l) *$ (IntV r) = IntV (l * r)
l *$ r = error $ "cannot * " ++ show l ++ " and " ++ show r

(IntV l) /$ (IntV r) = IntV (l `div` r)
l /$ r = error $ "cannot / " ++ show l ++ " and " ++ show r

(IntV l) <$ (IntV r) = BoolV (l < r)
l <$ r = error $ "cannot < " ++ show l ++ " and " ++ show r

(IntV l) >$ (IntV r) = BoolV (l > r)
l >$ r = error $ "cannot > " ++ show l ++ " and " ++ show r

l ==$ r = BoolV (l == r)
l !=$ r = BoolV (l /= r)

toDict :: M.Map String Value -> Value
toDict m =
	let wrapped = map (\(k,v) -> (StrV k, v)) $ M.toAscList m in
	DictV $ M.fromAscList wrapped

fromDict :: M.Map Value Value -> M.Map String Value
fromDict m =
	let unwrapped = map (\(StrV k,v) -> (k, v)) $ M.toAscList m in
	M.fromAscList unwrapped

-- some built-in functions

_fputbytes (TupleV [StreamV handle, StrV str]) = do
	io <- liftIO $ hPutStr handle str
	return unitv

_fputstr (TupleV [StreamV handle, StrV str]) = do
	io <- liftIO $ hPutStr handle str
	return unitv

_fgetline (StreamV handle) = do
	str <- liftIO $ hGetLine handle
	if last str == '\r' then -- remove trailing CR
		return . StrV $ init str
	else return $ StrV str

_freadbytes (TupleV [StreamV handle, IntV n]) = do
	str <- liftIO $ BSC.hGet handle (fromIntegral n :: Int)
	return . StrV $ BSC.unpack str

_fopen (TupleV [StrV path, StrV mode]) = do
	let mode' = case mode of
		"r" -> ReadMode
		"w" -> WriteMode
		"rw" -> ReadWriteMode
	handle <- liftIO $ openBinaryFile path mode'
	return $ StreamV handle

_feof (StreamV handle) = do
	isEof <- liftIO $ hIsEOF handle
	return $ BoolV isEof

_fclose (StreamV handle) = do
	liftIO $ hClose handle
	return unitv

_sockopen (TupleV [StrV host, IntV port]) = do
	liftIO $ SO.withSocketsDo $ do
		addr:_ <- SO.getAddrInfo Nothing (Just host) (Just $ show port)
		sock <- SO.socket (SO.addrFamily addr) SO.Stream SO.defaultProtocol
		SO.connect sock (SO.addrAddress addr)
		handle <- SO.socketToHandle sock ReadWriteMode
		hSetBuffering handle NoBuffering
		return $ StreamV handle

_putstr str@(StrV _) = _fputstr $ TupleV [StreamV stdout, str]
_putbytes str@(StrV _) = _fputbytes $ TupleV [StreamV stdout, str]
_getline (TupleV []) = _fgetline (StreamV stdin)

_print v = _putbytes $ StrV $ show v ++ "\n"
_repr v = return . StrV $ show v

_itos (IntV i) = return $ StrV $ show i
_itos v = error $ "itos: not an int: " ++ show v

_ref v = do
	value <- liftIO $ newIORef v
	return $ RefV value

_readRef (RefV r) = do
	value <- liftIO $ readIORef r
	return value

_setRef (TupleV [RefV r, v]) = do
	liftIO $ writeIORef r v
	return v

_loop args@(TupleV [fn@(FnV _ _), arg]) = do
	v <- apply fn arg
	if v /= BoolV False then
		_loop $ TupleV [fn, v]
	else return arg

_eval (TupleV [StrV code, DictV env]) = do
	let trySome :: IO a -> IO (Either SomeException a)
	    trySome = try
	    state = [fromDict env]
	ret <- liftIO. trySome $ evalStateT (evalString code) state
	case ret of
		Left err -> return $ TupleV [StrV "err", StrV (show err)]
		Right v -> return v
_eval (TupleV [code@(StrV _), (ListV env)]) =
	let env' = map (\(TupleV [k,v]) -> (k,v)) env in
	_eval (TupleV [code, DictV $ M.fromList env'])
_eval _ = error "eval: invalid args (want code and environment)"

-- returns a dictionary of a new environment with only the standard 
-- default-imported functions
_newStdEnv (TupleV []) = do
	let [stdEnv] = initialState
	return $ toDict stdEnv

_globals (TupleV []) = do
	env <- get
	return $ toDict (last env)

_locals (TupleV []) = do
	locals:_ <- get
	return $ toDict locals

-- import a module name as a module
_Import (StrV modname) = do
	env <- get -- save current state
	put initialState
	(path,modname) <- liftIO $ findModule modname -- find the module file
	evalFile path -- evaluate the module file
	[modenv] <- get -- get the module env
	let [initialEnv] = initialState
	let modenv' = M.difference modenv initialEnv -- subtract prelude stuff
	let mod = toDict modenv'
	let env' = bind env modname mod -- bind it
	put env' -- restore state
	return mod -- return module value
	
	where
		findModule :: FilePath -> IO (FilePath, String)
		findModule modname = do
			execPath <- fmap takeDirectory getExecutablePath
			findModuleIn [".", execPath </> "mods"] -- search paths for modules
			where
				findModuleIn [] = error $ "module " ++ modname ++ " couldn't be found"
				findModuleIn (dir:xs) = do
					let path = dir </> modname ++ ".lamb"
					exists <- doesFileExist path
					if exists then return (path, takeBaseName path)
					else findModuleIn xs

bif = Builtin . BIF
initialState = [M.fromList [
                            ("id", FnV emptyEnv [(VarP "x", Var "x")]),
                            ("loop", bif _loop),
                            ("ref!", bif _ref),
                            ("readRef!", bif _readRef),
                            ("setRef!", bif _setRef),
                            ("repr", bif _repr),
                            ("stdout", StreamV stdout),
                            ("stdin", StreamV stdin),
                            ("print", bif _print),
                            ("putstr", bif _putstr),
                            ("putstrln", bif (\x -> _putstr $ x +$ StrV "\n")),
                            ("getline", bif _getline),
                            ("fgetline", bif _fgetline),
                            ("putbytes", bif _putbytes),
                            ("fputbytes", bif _fputbytes),
                            ("fputstr", bif _fputstr),
                            ("freadbytes", bif _freadbytes),
                            ("feof", bif _feof),
                            ("fclose", bif _fclose),
                            ("fopen", bif _fopen),
                            ("sockopen", bif _sockopen),
                            ("itos", bif _itos),
                            ("globals", bif _globals),
                            ("locals", bif _locals),
                            ("newStdEnv", bif _newStdEnv),
                            ("eval", bif _eval),
                            ("import", bif _Import)]]

eval :: AST -> InterpState Value

eval (IntConst i) = return $ IntV i
eval (StrConst s) = return $ StrV s
eval (BoolConst b) = return $ BoolV b

eval (Block body) = foldr1 (>>) $ map eval body

eval (Cons a b) = do
	a' <- eval a
	b' <- eval b
	case b' of
		ListV v' -> return $ ListV $ a':v'
		_ -> error "cons: RHS must be a list"

eval (ListConst v) = mapM eval v >>= return . ListV
eval (TupleConst v) = mapM eval v >>= return . TupleV

eval (IfExpr c t e) = eval c >>= \cond ->
	case cond of
		BoolV True -> eval t
		BoolV False -> eval e
		_ -> error "if: condition must be a boolean"

eval (Var var) = get >>= \env ->
	maybe (error $ "unbound variable " ++ var) return (lookup env var)

eval (Defun name fn) = do
	env <- get
	case lookup env name of
            Nothing -> -- bind new fn
                    eval fn >>= \fn' ->
                    	put (bind env name fn') >> return fn'
            Just oldfn -> -- add pattern to old fn
                    let FnV cls oldpats = oldfn
                        Lambda [(pat, body)] = fn
                        newfn = FnV cls (oldpats ++ [(pat, body)]) in
                        put (bind env name newfn) >> return newfn

eval (Def pat v') = do
	v <- eval v'
	locals:xs <- get
	case patternBindings pat v of
		Nothing -> error $ "pattern binding doesn't satisfy: " ++ show v ++ " with " ++ show pat
		Just bindings -> do
			put $ (M.union bindings locals):xs -- update our local bindings
			return v

eval (Lambda pats) = do
	env <- get
	if length env == 1 then -- if in global env just use [], denoting the current global scope
		return $ FnV [] pats
	else return $ FnV env pats

eval (Add l r) = do { l <- eval l; r <- eval r; return $ l +$ r }
eval (Sub l r) = do { l <- eval l; r <- eval r; return $ l -$ r }
eval (Mul l r) = do { l <- eval l; r <- eval r; return $ l *$ r }
eval (Div l r) = do { l <- eval l; r <- eval r; return $ l /$ r }

eval (Equals l r) = do { l <- eval l; r <- eval r; return $ l ==$ r }
eval (NotEquals l r) = do { l <- eval l; r <- eval r; return $ l !=$ r }
eval (LessThan l r) = do { l <- eval l; r <- eval r; return $ l <$ r }
eval (GreaterThan l r) = do { l <- eval l; r <- eval r; return $ l >$ r }

eval (Access left (Var right)) = do
	lhs <- eval left
	case lhs of
		DictV dict ->
			case M.lookup (StrV right) dict of
				Just (FnV [] fn) -> -- use the module's global scope
					return $ FnV (mapToEnv dict) fn
				Just v -> return v
				Nothing -> return $ TupleV [StrV "nothing"]
		_ -> error $ "op/: need a dict, got " ++ show lhs
	where
		mapToEnv :: M.Map Value Value -> Env
		mapToEnv m = [M.fromAscList $ map (\(StrV k,v) -> (k,v)) (M.toAscList m)]
eval (Access _ _) = error "op/: RHS must be an identifier"

eval (Call lhs arg) = do
	env <- get
	v <- eval lhs
	case v of
	  fn@(FnV cls _) -> do
	  	arg' <- eval arg
	  	let cls' = if cls == [] then [last env] else cls -- if [], use current global env
	  	put cls' -- enter closure env
	  	v <- apply fn arg'
	  	put env -- restore env
	  	return v
	  fn@(Builtin _) -> eval arg >>= apply fn
	  _ -> error $ "call: " ++ show v ++ " is not a function"

eval x = error $ "eval: unhandled: " ++ show x

patternBindings :: Pattern -> Value -> Maybe (M.Map String Value)
patternBindings (VarP n) v = Just $ M.fromList [(n, v)]

patternBindings (IntP n) (IntV v)
	| v == n = Just M.empty
	| otherwise = Nothing
patternBindings (IntP n) _ = Nothing

patternBindings (BoolP b) (BoolV v)
	| v == b = Just M.empty
	| otherwise = Nothing

patternBindings (StrP x) (StrV y)
	| x == y = Just M.empty
	| otherwise = Nothing
patternBindings (StrP _) _ = Nothing

-- cons on strings
patternBindings (ConsP x (ListP [])) (StrV (y:[])) = patternBindings x (StrV [y])
-- "xy":xs pattern
patternBindings (ConsP (StrP xp) xsp) (StrV str) =
	let len = length xp in
	if take len str == xp then -- matches
		patternBindings xsp $ StrV (drop len str) -- match the rest of the string
	else Nothing -- no match
patternBindings (ConsP xp xsp) (StrV (x:xs)) =
	do
		xe <- patternBindings xp (StrV [x])
		xse <- patternBindings xsp $ StrV xs
		Just $ M.union xe xse

-- cons on lists
patternBindings (ConsP x (ListP [])) (ListV (y:[])) = patternBindings x y
patternBindings (ConsP xp xsp) (ListV (x:xs)) =
	do
		xe <- patternBindings xp x
		xse <- patternBindings xsp $ ListV xs
		Just $ M.union xe xse
patternBindings (ConsP _ _) _ = Nothing

-- lists
patternBindings (ListP []) (ListV (x:_)) = Nothing -- not enough patterns
patternBindings (ListP (_:_)) (ListV []) = Nothing -- not enough values
patternBindings (ListP []) (ListV []) = Just M.empty -- base case
patternBindings (ListP (x:xs)) (ListV (y:ys)) =
	do
		env <- patternBindings x y
		env' <- patternBindings (ListP xs) (ListV ys)
		Just $ M.union env' env
patternBindings (ListP _) _ = Nothing -- not a list

-- tuples
patternBindings (TupleP []) (TupleV (x:_)) = Nothing -- not enough patterns
patternBindings (TupleP (_:_)) (TupleV []) = Nothing -- not enough values
patternBindings (TupleP []) (TupleV []) = Just M.empty -- base case
patternBindings (TupleP (x:xs)) (TupleV (y:ys)) =
	do
		env <- patternBindings x y
		env' <- patternBindings (TupleP xs) (TupleV ys)
		Just $ M.union env' env
patternBindings (TupleP _) _ = Nothing -- not a tuple

-- applies a function
apply :: Value -> Value -> InterpState Value
apply (FnV _ pats) arg =
	apply' pats
	where
		apply' [] = error $ "argument " ++ show arg ++ " doesn't satisfy any patterns"
		apply' ((pat, body):xs) =
			case patternBindings pat arg of
				Just bindings -> -- satisfies
					do
						env <- get
						let newenv = bindings:env
						put newenv
						eval body
				Nothing -> -- doesn't satisfy this pattern
					apply' xs

apply (Builtin (BIF fn)) arg = fn arg

-- some helper programs for evaluation

-- sets up stdin/stdout for binary mode and makes them unbuffered
initIO :: IO ()
initIO = do
       hSetBinaryMode stdin True
       hSetBinaryMode stdout True
       hSetBuffering stdout NoBuffering

-- Takes an interpreter state and evaluates it with the empty initial state.
interpret :: InterpState a -> IO a
interpret state = evalStateT state initialState

evalProgram :: [AST] -> InterpState Value
evalProgram nodes = foldr1 (>>) $ map eval nodes

evalString :: String -> InterpState Value
evalString program =
	case parseProgram program of
		Left err -> error $ show err
		Right prg -> evalProgram prg

isLiterate :: FilePath -> Bool
isLiterate path = snd (splitExtension path) == ".lilamb"

-- Takes the lines of a literate program and returns the lines for a new executable program
-- from lines beginning with four spaces.
parseLiterate :: [String] -> [String]
parseLiterate lns = [drop 4 line | line <- lns, take 4 line == "    "]

evalFile :: FilePath -> InterpState Value
evalFile path = do
	contents <- liftIO $ if path == "-" then getContents else readFile path
	if isLiterate path then
		evalString . unlines . parseLiterate . lines $ contents
	else evalString contents

evalFileV :: FilePath -> IO Value
evalFileV = interpret . evalFile