module Interpreter.Interpreter where

import Prelude hiding (map)

import Control.Concurrent.MVar
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)  = pure $ UnnamedFnValue $ UnNamedFn args 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 = modify $ mapLocal (\s -> mempty : s)

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

evaluateProcedure :: ScopeKey -> [Value] -> InterpretM (Maybe Value)
evaluateProcedure sk args = do
  lookupScope sk >>= \case
    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)

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
  when isTail popScope
  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 :: DebugStepable a b => a -> InterpretM b
executeDebugStepable dbs = do
  isRunMode <$> get >>= \case
    NormalMode -> do
      execute dbs
    DebugMode -> do
      isStepMode <$> get >>= \case
        Continue -> execute dbs
        SingleStep -> do
          isDebugOut <$> get >>= sendDebugOut
          isDebugIn <$> get >>= (liftIO . takeMVar) >>= \case
            Run -> do
              modify (\is -> is { isRunMode = NormalMode })
              execute dbs
            StepIn -> execute dbs
            AddWatch _ -> execute dbs
            StepOver -> do
              modify (\is -> is { isStepMode = Continue })
              r <- execute dbs
              modify (\is -> is { isStepMode = SingleStep })
              pure r
  where
    sendDebugOut debugOut = do
      currentScope <- isLocal <$> get >>= \case
        [] -> isGlobalScope <$> get
        (scope : _) -> pure scope
      let
        dd = DebugState currentScope (getLocation dbs) (Just $ toSource dbs)
      liftIO $ putMVar debugOut $ DebugData dd

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"