-- MIT License, Copyright (c) 2022 Marvin Borner
module Eval
  ( evalMain
  ) where

import           Control.Concurrent
import           Control.DeepSeq                ( deepseq )
import           Control.Exception
import           Control.Monad                  ( when )
import           Control.Monad.State
import qualified Control.Monad.State.Strict    as StrictState
import qualified Data.BitString                as Bit
import qualified Data.ByteString.Lazy          as Byte
import           Data.Function                  ( on )
import           Data.Functor
import           Data.List
import qualified Data.Map                      as M
import           Data.Maybe
import           Data.Time.Clock
import           Paths_bruijn
import           System.Clock
import           System.Console.Haskeline
import           System.Directory
import           System.FilePath.Posix          ( takeBaseName
                                                , takeDirectory
                                                )
import           System.Mem
import           Target
import           Text.Megaparsec         hiding ( State
                                                , try
                                                )

import           Binary
import           Config
import           Conversion
import           Error
import           Helper
import           Humanification
import           Optimizer
import           Parser
import           Reducer

data EnvState = EnvState
  { _env   :: Environment
  , _conf  :: EvalConf
  , _cache :: EnvCache
  }

type M = StrictState.StateT EnvState IO

entryFunction :: Identifier
entryFunction = NormalFunction "main"

-- why isn't this in Prelude??
split :: (Eq a) => [a] -> [a] -> [[a]]
split _  [] = []
split [] x  = map (: []) x
split a@(_ : _) b@(c : _)
  | Just suffix <- a `stripPrefix` b = [] : split a suffix
  | null rest                        = [[c]]
  | otherwise                        = (c : head rest) : tail rest
  where rest = split a $ tail b

-- TODO: Force naming convention for namespaces/files
loadFile :: EvalConf -> EnvCache -> Environment -> IO EnvState
loadFile conf cache env = do
  f <- try $ readFile (_path conf) :: IO (Either IOError String)
  case f of
    Left exception ->
      print
          (ContextualError (ImportError $ show (exception :: IOError))
                           (Context "" $ _nicePath conf)
          )
        >> pure (EnvState env conf cache)
    Right f' -> eval
      (filter (not . null) $ split "\n\n" f')
      (EnvState
        env
        (conf { _isRepl = False, _evalPaths = _path conf : _evalPaths conf })
        cache
      )

evalFun :: Identifier -> Environment -> EvalState (Failable Expression)
evalFun fun (Environment sub) = state $ \env@(Environment e) ->
  let lookup' env' = case M.lookup fun env' of
        Nothing                    -> Left $ UndefinedIdentifier fun
        Just (EnvDef { _exp = x }) -> Right x
      matching n
        | null e = "<no idea>"
        | otherwise = snd $ minimumBy (compare `on` fst) $ map
          (\f -> (levenshtein (functionName f) n, show f))
          (M.keys e)
      suggest (Left u@(UndefinedIdentifier n)) =
        Left $ SuggestSolution u (matching $ functionName n)
      suggest x = x
  in  case lookup' sub of -- search in sub env
        s@(Right _) -> (s, env)
        _           -> (suggest $ lookup' e, env) -- search in global env

evalAbs :: Expression -> Environment -> EvalState (Failable Expression)
evalAbs e sub = evalExp e sub <&> fmap Abstraction

evalApp
  :: Expression -> Expression -> Environment -> EvalState (Failable Expression)
evalApp f g sub = evalExp f sub >>= \case
  Left  e  -> pure $ Left e
  Right f' -> fmap (Application f') <$> evalExp g sub

-- TODO: This could be nicer and more performant (current is redundantly recursive)
evalMixfix :: [Mixfix] -> Environment -> EvalState (Failable Expression)
evalMixfix m sub = resolve (mixfixKind m) mixfixArgs
 where
  longestMatching [] = pure $ Left $ UnmatchedMixfix (mixfixKind m) m
  longestMatching x  = evalFun (MixfixFunction x) sub >>= \case
    Left  _ -> longestMatching $ init x
    Right _ -> pure $ Right $ Function $ MixfixFunction x
  holeCount f = length [ h | h@MixfixNone <- f ]
  resolve f args
    | null [ s | s@(MixfixSome _) <- f ] = evalExp (foldl1 Application args) sub
    | otherwise = longestMatching f >>= \case
      e@(Left _) -> pure e
      Right l@(Function (MixfixFunction l')) ->
        let splitted = take (holeCount l') args
            chainRst = drop (length l') m
        in  case chainRst of
              [] -> evalExp (foldl1 Application $ l : splitted) sub
              _  -> evalExp
                ( MixfixChain
                $ MixfixExpression (foldl1 Application $ l : splitted)
                : chainRst
                )
                sub
      _ -> invalidProgramState
  mixfixArgs = [ a | (MixfixExpression a) <- m ]
  mixfixKind = map $ \case
    MixfixOperator i -> MixfixSome $ functionName i
    _                -> MixfixNone

evalPrefix
  :: Identifier -> Expression -> Environment -> EvalState (Failable Expression)
evalPrefix p e = evalExp $ Application (Function p) e

evalQuote :: Expression -> Environment -> EvalState (Failable Expression)
evalQuote f sub = evalExp f sub >>= \case
  Left  e  -> pure $ Left e
  Right f' -> pure $ Right $ quotify 0 f'
   where
    base l r = Abstraction $ Abstraction $ Abstraction $ Application l r
    quotify n (Abstraction e) = base (Bruijn 0) (quotify (n + 1) e)
    quotify n (Application l r) =
      base (Application (Bruijn 1) (quotify (n + 1) l)) (quotify (n + 1) r)
    quotify _ (Bruijn i) = base (Bruijn 2) (decimalToUnary $ fromIntegral i)
    quotify n (Unquote (Bruijn i)) = Bruijn $ n * 3 + i
    quotify n (Unquote e         ) = quotify n e
    quotify _ _                    = invalidProgramState

evalUnquote :: Expression -> Environment -> EvalState (Failable Expression)
evalUnquote f sub = evalExp f sub >>= \case
  Left  e  -> pure $ Left e
  Right f' -> pure $ Right $ Unquote $ reduceNoIO f'

evalExp :: Expression -> Environment -> EvalState (Failable Expression)
evalExp idx@(Bruijn      _  ) = const $ pure $ Right idx
evalExp (    Function    fun) = evalFun fun
evalExp (    Abstraction e  ) = evalAbs e
evalExp (    Application f g) = evalApp f g
evalExp (    MixfixChain es ) = evalMixfix es
evalExp (    Prefix p e     ) = evalPrefix p e
evalExp (    Quote   e      ) = evalQuote e
evalExp (    Unquote e      ) = evalUnquote e

evalDefinition
  :: Identifier -> Expression -> Environment -> EvalState (Failable Expression)
evalDefinition i e sub = evalExp e sub >>= \case
  Left e' -> pure $ Left e'
  Right f ->
    modify
        (\(Environment s) -> Environment
          $ M.insert i (EnvDef f (Environment M.empty) defaultFlags) s
        )
      >> pure (Right f)

evalTest
  :: Expression -> Expression -> Environment -> EvalState (Failable Command)
evalTest e1 e2 sub = evalExp e1 sub >>= \case
  Left  err -> pure $ Left err
  Right e1' -> fmap (Test e1') <$> evalExp e2 sub

evalSubEnv :: [Instruction] -> EnvState -> IO EnvState
evalSubEnv [] s = return s
evalSubEnv (instr : is) s =
  handleInterrupt (putStrLn "<aborted>" >> return s)
    $ evalInstruction instr s (evalSubEnv is)

fullPath :: String -> String -> IO String
fullPath current path = do
  lib    <- getDataFileName path -- TODO: Use actual lib directory
  exists <- doesFileExist lib
  pure $ if exists then lib else takeDirectory current ++ "/" ++ path

evalCommand :: String -> EnvState -> Command -> IO EnvState
evalCommand inp s@(EnvState env@(Environment envDefs) conf cache) = \case
  Input path -> do
    full <- fullPath (_path conf) path
    if full `elem` _evalPaths conf
      then
        print
            (ContextualError (ImportError path) (Context inp $ _nicePath conf))
          >> pure s
      -- paths can never be cached upon input since they may depend on context
      -- BUT inputs can be loaded from import cache since they would've had missing definitions!
      else if M.member path (_imported cache)
        then
          let (Environment env') = fromJust $ M.lookup path (_imported cache)
          in  pure $ s { _env = Environment $ M.union env' envDefs }
        else do
          EnvState env' _ _ <- loadFile
            (conf { _nicePath = path, _path = full })
            cache -- TODO: Fix wrong `within` in import error
            env
          pure $ EnvState env' conf cache -- import => _isRepl = False
  Watch path ->
    let
      monitor mtime = do
        threadDelay 100000 -- TODO: fix watch
        full <- fullPath (_path conf) path
        t    <- getModificationTime full
        if t > mtime
          then
            putStrLn "\ESC[2Jreload"
            >> evalCommand inp s (Input full)
            >> monitor t
          else monitor t
    in  getCurrentTime >>= monitor
  Import path namespace -> do
    -- TODO: Merge with Input (very similar)
    full <- fullPath (_path conf) path
    if full `elem` _evalPaths conf
      then
        print
            (ContextualError (ImportError path) (Context inp $ _nicePath conf))
          >> pure s
      else if M.member path (_imported cache)
        then -- load from cache
          let
            (Environment env') = fromJust $ M.lookup path (_imported cache)
            prefix | null namespace   = takeBaseName path ++ "."
                   | namespace == "." = ""
                   | otherwise        = namespace ++ "."
            rewriteKeys "" = id
            rewriteKeys p  = M.mapKeys $ \f -> NamespacedFunction p f
            rewriteFuns =
              M.map $ \d -> d { _flags = (_flags d) { _isImported = True } }
            filterImported =
              M.filter $ \(EnvDef { _flags = f }) -> not $ _isImported f
            env'' = rewriteFuns $ rewriteKeys prefix $ filterImported env'
          in
            pure $ s { _env = Environment $ M.union env'' envDefs }
        else do
          EnvState (Environment env') _ cache' <- loadFile
            (conf { _nicePath = path, _path = full })
            cache -- TODO: Fix wrong `within` in import error
            (Environment M.empty)
          let
            cache'' = cache
              { _imported = M.insert path (Environment env')
                              $ M.union (_imported cache) (_imported cache')
              }
            prefix | null namespace   = takeBaseName path ++ "."
                   | namespace == "." = ""
                   | otherwise        = namespace ++ "."
            rewriteKeys "" = id
            rewriteKeys p  = M.mapKeys $ \f -> NamespacedFunction p f
            rewriteFuns =
              M.map $ \d -> d { _flags = (_flags d) { _isImported = True } }
            filterImported =
              M.filter $ \(EnvDef { _flags = f }) -> not $ _isImported f
            env'' = rewriteFuns $ rewriteKeys prefix $ filterImported env'
          pure $ EnvState (Environment $ M.union env'' envDefs) conf cache'' -- import => _isRepl = False
  Test e1 e2
    | _evalTests conf
    -> let (res, _) = evalTest e1 e2 (Environment M.empty) `runState` env
       in  case res of
             Left err ->
               print (ContextualError err $ Context inp $ _nicePath conf)
                 >> pure s
             Right (Test e1' e2') -> do
               lhs <- reduce conf e1'
               rhs <- reduce conf e2'
               when (lhs /= rhs) (print $ FailedTest e1 e2 lhs rhs)
               when (lhs == rhs && _isVerbose conf) (print $ PassedTest e1 e2)
               pure s
             _ -> pure s
    | otherwise
    -> pure s
  ClearState -> do
    -- TODO: actually free memory :/
    putStr "Currently allocated: "
    getAllocationCounter >>= putStr . show . (0 -)
    putStrLn " Byte"
    performGC
    pure $ EnvState (Environment M.empty) conf (EnvCache M.empty)
  Length e -> do
    let (res, _) = evalExp e (Environment M.empty) `runState` env
    case res of
      Left  err -> print err
      Right e'  -> do
        print $ length $ toBinary e'
        red <- optimizedReduce conf e'
        print $ length $ toBinary red
    pure s
  Blc e -> do
    let (res, _) = evalExp e (Environment M.empty) `runState` env
    case res of
      Left  err -> print err
      Right e'  -> do
        putStrLn $ toBinary e'
        red <- optimizedReduce conf e'
        putStrLn $ toBinary red
    pure s
  Jot str -> do
    let e        = fromJot str
    let (res, _) = evalExp e (Environment M.empty) `runState` env
    case res of
      Left  err -> print err
      Right e'  -> do
        print e
        print e'
        print $ length $ toBinary e'
        red <- optimizedReduce conf e'
        print red
        print $ length $ toBinary red
    pure s
  Time e -> do
    start <- getTime Monotonic
    let (res, _) = evalExp e (Environment M.empty) `runState` env
    end <- case res of
      Left  err -> print err >> getTime Monotonic
      Right e'  -> do -- TODO: should timing not include optimization time?
        red <- optimizedReduce conf e'
        deepseq red (getTime Monotonic)
    let roundSecs x = fromIntegral (round $ x * 1e6 :: Integer) / 1e6 :: Double
    putStr
      $ show
      $ roundSecs
      $ (fromIntegral $ toNanoSecs $ diffTimeSpec start end :: Double)
      / 1e9
    putStrLn " seconds"
    pure s

showResult :: Expression -> Environment -> IO ()
showResult reduced env =
  let matching   = matchingFunctions reduced env
      humanified = humanifyExpression reduced
  in  putStrLn
        $  "*> "
        <> show reduced
        <> (if null humanified then "" else "\n?> " <> humanified)
        <> (if null matching then "" else "\n#> " <> matching)

evalInstruction
  :: Instruction -> EnvState -> (EnvState -> IO EnvState) -> IO EnvState
evalInstruction (ContextualInstruction instr inp) s@(EnvState env conf _) rec =
  case instr of
    Define i e sub -> do
      EnvState subEnv _     _ <- evalSubEnv sub s
      (        res    , env') <- pure $ evalDefinition i e subEnv `runState` env
      case res of
        Left err ->
          print (ContextualError err $ Context inp $ _nicePath conf) >> pure s -- don't continue
        Right _
          | _isRepl conf -> putStrLn (show i <> " = " <> show e)
          >> return s { _env = env' }
          | otherwise -> rec s { _env = env' }
    Evaluate e ->
      let (res, _) = evalExp e (Environment M.empty) `runState` env
      in  (case res of
            Left  err -> print err >> rec s
            Right e'  -> do
              red <- optimizedReduce conf e'
              showResult red env
              rec s
          )
    Commands cs -> yeet (pure s) cs >>= rec
     where
        -- TODO: sus
      yeet s' []        = s'
      yeet s' (c : cs') = do
        s'' <- s'
        yeet (evalCommand inp s'' c) cs'
    _ -> rec s
evalInstruction instr s rec =
  evalInstruction (ContextualInstruction instr "<unknown>") s rec

eval :: [String] -> EnvState -> IO EnvState
eval []   s = return s
eval [""] s = return s
eval (block : bs) s@(EnvState _ conf _) =
  handleInterrupt (putStrLn "<aborted>" >> return s)
    $ case parse blockParser "" block of
        Left err ->
          print
              (ContextualError (SyntaxError $ printBundle err)
                               (Context "" $ _nicePath conf)
              )
            >> eval bs s
        Right instr -> evalInstruction instr s (eval bs)
 where
  blockParser | _isRepl conf = parseReplLine
              | otherwise    = parseBlock 0

dumpFile :: EvalConf -> (a -> IO ()) -> (Expression -> a) -> IO ()
dumpFile conf wr conv = do
  EnvState (Environment env) _ _ <- loadFile conf
                                             (EnvCache M.empty)
                                             (Environment M.empty)
  case M.lookup entryFunction env of
    Nothing -> print $ ContextualError (UndefinedIdentifier entryFunction)
                                       (Context "" (_nicePath conf))
    Just EnvDef { _exp = e } -> toTarget conf e >>= wr . conv

evalFileConf :: EvalConf -> IO ()
evalFileConf conf = do
  EnvState (Environment env) _ _ <- loadFile conf
                                             (EnvCache M.empty)
                                             (Environment M.empty)
  arg <- encodeStdin
  case M.lookup entryFunction env of
    Nothing -> print $ ContextualError (UndefinedIdentifier entryFunction)
                                       (Context "" (_nicePath conf))
    Just EnvDef { _exp = e } -> do
      red <- optimizedReduce conf { _hasArg = True } (Application e arg)
      showResult red (Environment env)

exec :: EvalConf -> (String -> IO (Either IOError a)) -> (a -> String) -> IO ()
exec conf rd conv = do
  f   <- rd (_path conf)
  arg <- encodeStdin
  case f of
    Left  exception -> print (exception :: IOError)
    Right f'        -> do
      red <- optimizedReduce conf { _hasArg = True } (Application e arg)
      showResult red (Environment M.empty)
      where e = fromBinary $ conv f'

repl :: EnvState -> InputT M ()
repl (EnvState env conf cache) =
  handleInterrupt (return $ Just "")
                  (withInterrupt $ getInputLine "\ESC[36mλ\ESC[0m ")
    >>= \case
      -- TODO: Add non-parser error support for REPL
          Nothing   -> return ()
          Just line -> do
            -- setting imported [] for better debugging
            s' <- liftIO
              $ eval [line] (EnvState env conf cache { _imported = M.empty })
            lift $ StrictState.put s'
            repl s'

lookupCompletion :: String -> M [Completion]
lookupCompletion str = do
  (EnvState (Environment env) _ _) <- StrictState.get
  return $ map (\s -> Completion s s False) $ filter
    (\s -> str `isPrefixOf` s)
    (map functionName (M.keys env))

completionSettings :: String -> Settings M
completionSettings history = Settings
  { complete       = completeWord Nothing " \n" lookupCompletion
  , historyFile    = Just history
  , autoAddHistory = True
  }

runRepl :: EvalConf -> IO ()
runRepl conf = do
  config  <- getDataFileName "config"
  history <- getDataFileName "history"
  prefs   <- readPrefs config
  let looper = runInputTWithPrefs
        prefs
        (completionSettings history)
        (withInterrupt $ repl $ EnvState (Environment M.empty)
                                         conf
                                         (EnvCache M.empty)
        )
  StrictState.evalStateT
    looper
    (EnvState (Environment M.empty) conf (EnvCache M.empty))

evalMain :: Args -> IO ()
evalMain args = do
  let conf = argsToConf args
  case _argMode args of
    ArgEval | _isRepl conf -> runRepl conf
    ArgEval | otherwise    -> evalFileConf conf
    ArgDumpBblc            -> dumpFile conf
                                       (Byte.putStr . Bit.realizeBitStringLazy)
                                       (toBitString . toBinary)
    ArgDumpBlc -> dumpFile conf putStrLn toBinary
    ArgEvalBblc ->
      exec conf (try . Byte.readFile) (fromBitString . Bit.bitStringLazy)
    ArgEvalBlc -> exec conf (try . readFile) id