module Interpreter.Interpreter where

import Prelude hiding (map)

import Control.Concurrent.STM as STM
import Control.Exception (throw)
import Control.Monad
import Control.Monad.Catch (catch)
import Control.Monad.Loops (iterateWhile)
import Data.Coerce
import qualified Data.List.NonEmpty as NE
import Data.Map as M hiding (map)
import Data.Text as T hiding (index, map)
import qualified Data.Vector as V

import Common
import Compiler.AST.FunctionStatement
import Compiler.AST.Program
import Compiler.Lexer
import Control.Monad.State.Strict
import Interpreter.Common

lookupScope :: ScopeKey -> InterpretM Value
lookupScope key =
  ((lookupInTopScope key . isLocal) <$> get) >>= \case
    Just v -> pure v
    Nothing -> ((M.lookup key . isGlobalScope) <$> get) >>= \case
      Just v -> pure v
      Nothing ->
        throwErr $ SymbolNotFound (pack $ show key)

lookupInTopScope :: ScopeKey -> [Scope] -> Maybe Value
lookupInTopScope _ [] = Nothing
lookupInTopScope key (h: _) = M.lookup key h

evaluateExpression :: ExpressionWithLoc -> InterpretM Value
evaluateExpression exp'@(ExpressionWithLoc _ loc) =
  catch
    (catch (executeDebugStepable exp') rteHandler) peHandler
  where
    rteHandler (r :: RuntimeError) = pure $ ErrorValue $ (hReadable r) <> " at " <> hReadable loc
    peHandler (r :: ProgramError) = throwErr @(InterpretM Value) $ RuntimeErrorWithLoc (Right r) loc

evaluateExpression_ :: Expression -> InterpretM Value
evaluateExpression_ (EParan le) = evaluateExpression le
evaluateExpression_ (ENegated le) = evaluateExpression le >>= \case
  NumberValue n -> pure $ NumberValue $ negateValue n
  x             -> throwErr $ UnexpectedType ("Number", x)
evaluateExpression_ (ELiteral le) = evaluateLiteralExpression le
evaluateExpression_ (EVar idf)     = lookupScope (SkIdentifier idf) >>= \case
  BuiltIn (BuiltinVal v) -> pure v
  v                      -> pure v
evaluateExpression_ (ESubscripted subscript) = evaluateSubscriptedExpr subscript
evaluateExpression_ (EConditional boolEx ex1 ex2)  = evaluateExpression boolEx >>= \case
  BoolValue True  -> evaluateExpression ex1
  BoolValue False -> evaluateExpression ex2
  x               -> throwErr $ UnexpectedType ("Bool", x)
evaluateExpression_ (EOperator op e1 e2)  = do
  evaluateFn (FnOp op) [e1, e2] False >>= \case
    Just v  -> pure v
    Nothing -> throwErr MissingProcedureReturn
evaluateExpression_ (ECall iden exprs isTail)  =  evaluateFn (FnName iden) exprs isTail >>= \case
  Just v  -> pure v
  Nothing -> throwErr MissingProcedureReturn
evaluateExpression_ (EUnnamedFn args expr)  = do
  isLocal <$> get >>= \case
    [] -> pure $ UnnamedFnValue $ UnNamedFn args mempty expr
    (h: _) -> pure $ UnnamedFnValue $ UnNamedFn args h expr

popScope :: InterpretM ()
popScope = do
  (isLocal <$> get) >>= \case
    (_:rst) -> do
      modify (\is -> is { isLocal = rst })
    _ -> throwErr EmptyScopeStack

data FnId
  = FnOp Operator
  | FnName Identifier

evaluateCallback :: Callback -> [Value] -> InterpretM (Maybe Value)
evaluateCallback (CallbackUnNamed un) args = Just <$> evaluateUnnamedFn un args
evaluateCallback (CallbackNamed idf) args =
  evaluateProcedure (SkIdentifier idf) args

insertEmptyScope :: InterpretM ()
insertEmptyScope = insertScope mempty

insertScope :: Scope -> InterpretM ()
insertScope scope = modify $ mapLocal (\s -> scope : s)

evaluateUnnamedFn :: UnNamedFn -> [Value] -> InterpretM Value
evaluateUnnamedFn (UnNamedFn Nothing scope expr) _ = do
  insertScope scope
  x <- evaluateExpression expr
  popScope
  pure x
evaluateUnnamedFn (UnNamedFn (Just (NE.toList -> argNames)) scope expr) argsVals = do
  insertScope scope
  zipWithM_ (\a1 a2 -> insertBinding a1 a2) (SkIdentifier <$> argNames) argsVals -- @TODO Check argument counts
  r <- evaluateExpression expr
  popScope
  pure r

evaluateProcedure_ :: Value -> [Value] -> InterpretM (Maybe Value)
evaluateProcedure_ fnVal args = case fnVal of
  UnnamedFnValue un -> Just <$> evaluateUnnamedFn un args
  (ProcedureValue (FunctionDef _ argNames (NE.toList -> stms))) -> do
    insertEmptyScope
    zipWithM_ (\a1 a2 -> insertBinding a1 a2) (SkIdentifier <$> argNames) args -- @TODO Check argument counts
    executeStatements stms >>= \case
      ProcReturn False v -> do
        popScope
        pure $ Just v
      ProcReturn True v -> do
        -- Don't pop stack if the return was a tail call return
        -- because the stack was popped before entering the
        -- call.
        pure $ Just v
      ProcBreak -> do
        popScope
        pure Nothing
      ProcContinue -> do
        popScope
        pure Nothing
  (BuiltIn (BuiltinCall cb)) -> cb args
  (BuiltIn (BuiltinCallWithDoc (SomeBuiltin cb))) -> cb (toArgs args)
  a -> throwErr $ UnexpectedType ("Procedure", a)

evaluateProcedure :: ScopeKey -> [Value] -> InterpretM (Maybe Value)
evaluateProcedure sk args =
  lookupScope sk >>= (\x -> evaluateProcedure_ x args)

evaluateFn :: FnId -> [ExpressionWithLoc] -> Bool -> InterpretM (Maybe Value)
evaluateFn fnId argsExps isTail = do
  let sk = case fnId of
        FnOp op    -> SkOperator op
        FnName idf -> SkIdentifier idf
  args <- mapM (\x -> evaluateExpression x) argsExps
  if isTail
    then do
      fnVal <- lookupScope sk
      popScope
      evaluateProcedure_ fnVal args
    else evaluateProcedure sk args

evaluateSubscriptedExpr :: SubscriptedExpression -> InterpretM Value
evaluateSubscriptedExpr (EArraySubscript expr indexExpr) = evaluateExpression expr >>= \case
  ArrayValue v -> evaluateExpression indexExpr >>= \case
    NumberValue (NumberInt i) -> do
      let index :: Int = fromIntegral i
      if index <= V.length v && index >= 0 then (pure $ v V.! (index - 1)) else (throwErr $ ListIndexOutOfBounds index)
    a -> throwErr $ UnexpectedType ("Integer index", a)
  ObjectValue mp -> evaluateExpression indexExpr >>= \case
    StringValue key -> case M.lookup key mp of
      Just v  -> pure v
      Nothing -> throwErr $ KeyNotFound (pack $ show key)
    a -> throwErr $ UnexpectedType ("Property index", a)
  a -> throwErr $ UnexpectedType ("Array/Map", a)
evaluateSubscriptedExpr (EPropertySubscript expr (unIdentifer -> key)) = evaluateExpression expr >>= \case
  ObjectValue mp -> case M.lookup key mp of
      Just v  -> pure v
      Nothing -> throwErr $ KeyNotFound (pack $ show key)
  a -> throwErr $ UnexpectedType ("Map", a)

evaluateVar :: Subscript -> InterpretM Value
evaluateVar (NoSubscript idf) = lookupScope (SkIdentifier idf) >>= \case
  BuiltIn (BuiltinVal v) -> pure $ v
  v                      -> pure v
evaluateVar (SubscriptExpr sub expr) =
  -- Arrays are indexed from 1, not 0.
  evaluateExpression expr >>= \case
      NumberValue (NumberInt int) ->
        evaluateVar sub >>= \case
          ArrayValue v -> do
            let index :: Int = fromIntegral int
            if index <= V.length v && index >= 0 then (pure $ v V.! (index - 1)) else (throwErr $ ListIndexOutOfBounds index)
          a -> throwErr $ UnexpectedType ("Array/Object", a)
      StringValue key -> lookupInMapVar sub key
      a  -> throwErr $ UnexpectedType ("String/Integer container key", a)
evaluateVar (PropertySubscript sub idf) = lookupInMapVar sub (unIdentifer idf)

lookupInMapVar :: Subscript -> Text -> InterpretM Value
lookupInMapVar sub key = evaluateVar sub >>= \case
  ObjectValue mp -> case M.lookup key mp of
    Just v  -> pure v
    Nothing -> throwErr $ KeyNotFound (pack $ show key)
  a            ->  throwErr $ UnexpectedType ("Expecting Object Looking for key: " <> (T.pack $ show sub) <> ":" <> key, a)

evaluateLiteralExpression :: LiteralExpression -> InterpretM Value
evaluateLiteralExpression (LAtomic (LitString t)) = pure $ StringValue t
evaluateLiteralExpression (LAtomic (LitBytes t))  = pure $ BytesValue t
evaluateLiteralExpression (LAtomic (LitNumber n)) = pure $ NumberValue $ NumberInt n
evaluateLiteralExpression (LAtomic (LitFloat f))  = pure $ NumberValue $ NumberFractional (realToFrac f)
evaluateLiteralExpression (LAtomic (LitBool b))   = pure $ BoolValue b
evaluateLiteralExpression (LArray l)   = do
  v <- mapM (\x -> evaluateExpression x) l
  pure $ ArrayValue (V.fromList v)
evaluateLiteralExpression (LObject l)   = do
  v <- mapM (\x -> evaluateExpression x) l
  pure $ ObjectValue v

voidStm :: () -> InterpretM (Maybe Value)
voidStm _ = pure Nothing

executeStatements :: [FunctionStatementWithLoc] -> InterpretM ProcResult
executeStatements x = foldM (\a1 a2 -> fn a1 a2) ProcContinue x
  where
    fn :: ProcResult -> FunctionStatementWithLoc -> InterpretM ProcResult
    fn (ProcReturn tc x') _ = pure $ ProcReturn tc x'
    fn ProcBreak _       = pure ProcBreak
    fn ProcContinue fs   = executeStatement fs

modifyBinding :: Subscript -> Value -> InterpretM ()
modifyBinding (NoSubscript idf) val = insertBinding (SkIdentifier idf) val
modifyBinding (PropertySubscript sub (unIdentifer -> key)) val = do
  evaluateVar sub >>= \case
    ObjectValue v -> case M.lookup key v of
      Just _  -> modifyBinding sub (ObjectValue $ M.insert key val v)
      Nothing -> throwErr (KeyNotFound key)
    a -> throwErr $ UnexpectedType ("Map", a)
modifyBinding (SubscriptExpr sub expr) val = do
  evaluateVar sub >>= \case
    ArrayValue v -> evaluateExpression expr >>= \case
      NumberValue (NumberInt idx) -> do
        let index :: Int = fromIntegral idx
        if (index <= V.length v && index > 0)
          then modifyBinding sub (ArrayValue $ V.update v (V.fromList [(index - 1, val)]))
          else throwErr $ ListIndexOutOfBounds index
      a -> throwErr $ UnexpectedType ("Integer Index", a)
    ObjectValue v -> evaluateExpression expr >>= \case
      StringValue key -> case M.lookup key v of
        Just _  -> modifyBinding sub (ObjectValue $ M.insert key val v)
        Nothing -> throwErr (KeyNotFound key)
      a -> throwErr $ UnexpectedType ("String", a)
    a -> throwErr $ UnexpectedType ("Map", a)

class ToSource a => DebugStepable a b | a -> b where
  getLocation :: a -> Location
  execute :: a -> InterpretM b

instance DebugStepable FunctionStatementWithLoc ProcResult where
  getLocation (FunctionStatementWithLoc _ l) = l
  execute (FunctionStatementWithLoc fs _) = executeStatement_ fs

instance DebugStepable ExpressionWithLoc Value where
  getLocation (ExpressionWithLoc _ l) = l
  execute (ExpressionWithLoc exp' _) = evaluateExpression_ exp'

executeStatement :: FunctionStatementWithLoc -> InterpretM ProcResult
executeStatement fs@(FunctionStatementWithLoc _ loc) =
  catch
    (catch (executeDebugStepable fs) rteHandler) peHandler
    where
      rteHandler (r :: RuntimeError) = case r of
        CustomRTE msg -> throw (RuntimeErrorWithLoc (Left $ CustomRTE msg) loc)
        _ -> throw (RuntimeErrorWithLoc (Left r) loc)
      peHandler (r :: ProgramError) = throw (RuntimeErrorWithLoc (Right r) loc)

executeDebugStepable :: Show a => DebugStepable a b => a -> InterpretM b
executeDebugStepable dbs = do
  isRunMode <$> get >>= \case
    NormalMode _ -> do
      execute dbs
    DebugMode debugEnv@(DebugEnv { deInQueue = isDebugIn, deOutQueue = isDebugOut, deStepMode = stepMode }) -> do
      case stepMode of
        Continue -> execute dbs
        SingleStep -> do
          -- Send location of current instruction, and wait for command.
          sendDebugOut isDebugOut
          (liftIO $ atomically $ readTBQueue isDebugIn) >>= \case
              Run -> do
                modify (\is -> is { isRunMode = NormalMode (Just debugEnv) })
                execute dbs
              StepIn -> do
                modify (\is -> is { isRunMode = DebugMode (DebugEnv SingleStep isDebugIn isDebugOut) })
                execute dbs
              _ -> error "Unexpected debug command"
  where
    sendDebugOut debugOut = do
      is <- get
      let currentScope = case isLocal is of
            [] -> isGlobalScope is
            (scope : _) -> scope
      let
        dd = DebugState currentScope (getLocation dbs) (Just $ trimAndElipsis $ toSource dbs) (isThreadName is)
      liftIO $ atomically $ writeTBQueue debugOut $ DebugData dd

    trimAndElipsis (T.replace "\n" " " -> t) = if T.length t > 30 then T.take 30 t <> "..." else t

executeStatement_ :: FunctionStatement -> InterpretM ProcResult
executeStatement_ (FnComment _) = pure ProcContinue
executeStatement_ (Let sub exp') = do
  sourceValue <- evaluateExpression exp'
  modifyBinding sub sourceValue
  pure ProcContinue
executeStatement_ (Call iden args) = do
  _ <- evaluateFn (FnName iden) args False
  pure ProcContinue
executeStatement_ (IfThen expr stms) = evaluateExpression expr >>= \case
  BoolValue True -> executeStatements (NE.toList stms)
  BoolValue _    -> pure ProcContinue
  a              ->  throwErr $ UnexpectedType ("Bool", a)
executeStatement_ (If expr stms1 stms2) = evaluateExpression expr >>= \case
  BoolValue b -> case b of
    True  -> executeStatements (NE.toList stms1)
    False -> executeStatements (NE.toList stms2)

  a            ->  throwErr $ UnexpectedType ("Bool", a)
executeStatement_ (MultiIf expr stms1 elseifs mstms2) = evaluateExpression expr >>= \case
  BoolValue True -> executeStatements (NE.toList stms1)
  BoolValue False -> foldM executeElseIf Nothing elseifs >>= \case
    Just r  -> pure r
    Nothing -> case mstms2 of
      Just stms2 -> executeStatements (NE.toList stms2)
      Nothing    -> pure ProcContinue
  a            ->  throwErr $ UnexpectedType ("Bool", a)
  where
    executeElseIf a@(Just _) _ = pure a
    executeElseIf Nothing (bexpr, stms) = evaluateExpression bexpr >>= \case
      BoolValue True  -> Just <$> executeStatements (NE.toList stms)
      BoolValue False -> pure Nothing
      a               ->  throwErr $ UnexpectedType ("Bool", a)

executeStatement_ (Return eloc@(ExpressionWithLoc { elExpression = ECall idf args _ })) =
  -- TCO
  evaluateExpression (eloc { elExpression = ECall idf args True }) >>= pure . ProcReturn True
executeStatement_ (Return expr) = evaluateExpression expr >>= pure . ProcReturn False
executeStatement_ Break = pure ProcBreak
executeStatement_ (Loop (NE.toList -> stms)) = iterateWhile (\case
  ProcBreak    -> False
  ProcContinue -> True
  ProcReturn _ _ -> False) (executeStatements stms)

executeStatement_ (While exprBool (NE.toList -> stms)) = do
  r <- iterateWhile (\case
    ProcBreak    -> False
    ProcContinue -> True
    ProcReturn _ _ -> False)
    (evaluateExpression exprBool >>= \case
        BoolValue True  -> executeStatements stms
        BoolValue False -> pure ProcBreak
        a               -> throwErr $ UnexpectedType ("Bool", a))
  case r of
    ProcBreak -> pure ProcContinue
    a         -> pure a
executeStatement_ (For iden exprFrom exprTo (NE.toList -> stms)) = evaluateExpression exprFrom >>= \case
  NumberValue (NumberInt start) -> evaluateExpression exprTo >>= \case
    NumberValue (NumberInt end) -> do
      let
        fn :: ProcResult -> IntType -> InterpretM ProcResult
        fn ProcContinue current = do
          insertBinding (SkIdentifier iden) (NumberValue $ NumberInt $ current)
          executeStatements stms
        fn r _ = pure r
      foldM (\a1 a2 -> fn a1 a2) ProcContinue [start .. end] >>= \case
        ProcReturn tc v -> pure $ ProcReturn tc v
        _            -> pure ProcContinue
    a            ->  throwErr $ UnexpectedType ("Int", a)
  a            ->  throwErr $ UnexpectedType ("Int", a)

executeStatement_ (ForEach iden expr (NE.toList -> stms)) = evaluateExpression expr >>= \case
  ObjectValue map -> do
    foldM (\a1 (k, v) -> fn a1 (ObjectValue $ M.fromList [("key", StringValue k), ("value", v)])) ProcContinue (M.assocs map) >>= \case
      ProcReturn tc v -> pure $ ProcReturn tc v
      _            -> pure ProcContinue
  ArrayValue values -> do
    V.foldM (\a1 a2 -> fn a1 a2) ProcContinue values >>= \case
      ProcReturn tc v -> pure $ ProcReturn tc v
      _            -> pure ProcContinue
  a            ->  throwErr $ UnexpectedType ("Array/Object", a)
  where
    fn :: ProcResult -> Value -> InterpretM ProcResult
    fn ProcContinue current = do
      insertBinding (SkIdentifier iden) current
      executeStatements stms
    fn r _ = pure r

filter_ :: BuiltInFnWithDoc '[ '("list", V.Vector Value), '("callback", Callback)]
filter_ ((coerce -> v1) :> (coerce -> callback) :> _) =
  (\x -> Just $ ArrayValue x) <$> V.filterM fn v1
  where
    fn v = evaluateCallback callback [v] >>= \case
      Just (BoolValue x) -> pure x
      _           -> throwErr $ CustomRTE "Callback returned a non-bool value"