-- Copyright 2023 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

{-# LANGUAGE UndecidableInstances #-}

module Live.Eval (
  watchAndEvalFile, EvalServer, CellState (..), CellUpdate (..), CellsState, CellsUpdate,
  NodeList (..), NodeListUpdate (..), subscribeIO,
  CellStatus (..), nodeListAsUpdate, NodeId, evalFileNonInteractive) where

import Control.Concurrent
import Control.Monad
import Control.Monad.State.Strict
import Control.Monad.Writer.Strict
import qualified Data.Map.Strict as M
import Data.Aeson (ToJSON)
import Data.Functor ((<&>))
import Data.Maybe (fromJust)
import Data.Text (Text)
import Data.IORef
import Prelude hiding (span)
import GHC.Generics

import Actor
import IncState
import Types.Source
import TopLevel
import ConcreteSyntax
import MonadUtil
import Util (readFileText)

-- === Top-level interface ===

type EvalServer = StateServer CellsState

-- `watchAndEvalFile` returns the channel by which a client may
-- subscribe by sending a write-only view of its input channel.
watchAndEvalFile :: FilePath -> EvalConfig -> TopStateEx -> IO EvalServer
watchAndEvalFile fname opts env = do
  watcher <- launchFileWatcher fname
  parser <- launchCellParser watcher \source -> uModuleSourceBlocks $ parseUModule Main source
  launchDagEvaluator opts parser env

sourceBlockEvalFun :: EvalConfig -> Mailbox Outputs -> TopStateEx -> SourceBlock -> IO (ExitStatus, TopStateEx)
sourceBlockEvalFun cfg resultChan env block =
  evalSourceBlockIO cfg (send resultChan) env block

-- === Evaluating non-interactively to produce a standalone HTML page ===

evalFileNonInteractive :: FilePath -> EvalConfig -> TopStateEx -> IO CellsState
evalFileNonInteractive fname cfg initEnv = do
  envRef <- newIORef initEnv
  blocks <- parseSourceBlocks <$> readFileText fname
  cellStates <- forM blocks \block -> do
    if isInert block
      then return $ CellState block Inert mempty
      else do
        env <- readIORef envRef
        ((exitStatus, newEnv), outs) <- captureLogs \logger ->
          evalSourceBlockIO cfg logger env block
        writeIORef envRef newEnv
        return $ CellState block (exitStatusAsCellStatus exitStatus) outs
  runFreshNameT $ buildNodeList cellStates


-- === DAG diff state ===

-- We intend to make this an arbitrary Dag at some point but for now we just
-- assume that dependence is just given by the top-to-bottom ordering of blocks
-- within the file.

type NodeId = Int

data NodeList a = NodeList
  { orderedNodes :: [NodeId]
  , nodeMap      :: M.Map NodeId a }
  deriving (Show, Generic, Functor)

data NodeListUpdate s = NodeListUpdate
  { orderedNodesUpdate :: TailUpdate NodeId
  , nodeMapUpdate      :: MapUpdate NodeId s }
  deriving (Generic)

instance IncState s => Semigroup (NodeListUpdate s) where
  NodeListUpdate x1 y1 <> NodeListUpdate x2 y2 = NodeListUpdate (x1<>x2) (y1<>y2)

instance IncState s => Monoid (NodeListUpdate s) where
  mempty = NodeListUpdate mempty mempty

instance IncState s => IncState (NodeList s) where
  type Delta (NodeList s) = NodeListUpdate s
  applyDiff (NodeList m xs) (NodeListUpdate dm dxs) =
    NodeList (applyDiff m dm) (applyDiff xs dxs)

type Dag       a = NodeList (Unchanging a)
type DagUpdate a = NodeListUpdate (Unchanging a)

nodeListAsUpdate :: NodeList s -> NodeListUpdate s
nodeListAsUpdate (NodeList xs m)= NodeListUpdate (TailUpdate 0 xs) (MapUpdate $ fmap Create m)

emptyNodeList :: NodeList a
emptyNodeList = NodeList [] mempty

buildNodeList :: FreshNames NodeId m => [a] -> m (NodeList a)
buildNodeList vals = do
  nodeList <- forM vals \val -> do
    nodeId <- freshName
    return (nodeId, val)
  return $ NodeList (fst <$> nodeList) (M.fromList nodeList)

commonPrefixLength :: Eq a => [a] -> [a] -> Int
commonPrefixLength (x:xs) (y:ys) | x == y = 1 + commonPrefixLength xs ys
commonPrefixLength _ _ = 0

nodeListVals :: NodeList a -> [a]
nodeListVals nodes = orderedNodes nodes <&> \k -> fromJust $ M.lookup k (nodeMap nodes)

computeNodeListUpdate :: (Eq s, FreshNames NodeId m) => NodeList s -> [s] -> m (NodeListUpdate s)
computeNodeListUpdate nodes newVals = do
  let prefixLength = commonPrefixLength (nodeListVals nodes) newVals
  let oldTail = drop prefixLength $ orderedNodes nodes
  NodeList newTail nodesCreated <- buildNodeList $ drop prefixLength newVals
  let nodeUpdates = fmap Create nodesCreated <> M.fromList (fmap (,Delete) oldTail)
  return $ NodeListUpdate (TailUpdate (length oldTail) newTail) (MapUpdate nodeUpdates)

-- === Cell parser ===

-- This coarsely parses the full file into blocks and forms a DAG (for now a
-- trivial one assuming all top-to-bottom dependencies) of the results.

type CellParser = StateServer (Dag SourceBlock)

data CellParserMsg =
    Subscribe_CP (SubscribeMsg (Dag SourceBlock))
  | Update_CP (Overwrite Text)
  deriving (Show)

launchCellParser :: MonadIO m => FileWatcher -> (Text -> [SourceBlock]) -> m CellParser
launchCellParser fileWatcher parseCells =
  sliceMailbox Subscribe_CP <$> launchActor (cellParserImpl fileWatcher parseCells)

cellParserImpl :: FileWatcher -> (Text -> [SourceBlock]) -> ActorM CellParserMsg ()
cellParserImpl fileWatcher parseCells = runFreshNameT do
  Overwritable initContents <- subscribe Update_CP fileWatcher
  initNodeList <- buildNodeList $ fmap Unchanging $ parseCells initContents
  runIncServerT initNodeList $ messageLoop \case
    Subscribe_CP msg -> handleSubscribeMsg msg
    Update_CP NoChange -> return ()
    Update_CP (OverwriteWith newContents) -> do
      let newCells = fmap Unchanging $ parseCells newContents
      curNodeList <- getl It
      update =<< computeNodeListUpdate curNodeList newCells
      flushDiffs

-- === Dag evaluator ===

-- This is where we track the state of evaluation and decide what we needs to be
-- run and what needs to be killed.

type Evaluator = StateServer CellsState
newtype EvaluatorM a =
  EvaluatorM { runEvaluatorM' ::
                 IncServerT CellsState
                   (StateT EvaluatorState
                      (ActorM EvaluatorMsg)) a }
  deriving (Functor, Applicative, Monad, MonadIO, Actor (EvaluatorMsg))
deriving instance IncServer CellsState EvaluatorM

instance DefuncState EvaluatorMUpdate EvaluatorM where
  update = \case
    UpdateDagEU dag     -> EvaluatorM $ update dag
    UpdateCurJob status -> EvaluatorM $ lift $ modify \s -> s { curRunningJob = status }
    UpdateEnvs   envs   -> EvaluatorM $ lift $ modify \s -> s { prevEnvs      = envs}
    AppendEnv env -> do
      envs <- getl PrevEnvs
      update $ UpdateEnvs $ envs ++ [env]
    UpdateCellState nodeId cellUpdate -> update $ UpdateDagEU $ NodeListUpdate mempty $
      MapUpdate $ M.singleton nodeId $ Update cellUpdate

instance LabelReader EvaluatorMLabel EvaluatorM where
  getl l = case l of
    NodeListEM      -> EvaluatorM $ orderedNodes                <$> getl It
    NodeInfo nodeId -> EvaluatorM $ M.lookup nodeId <$> nodeMap <$> getl It
    PrevEnvs        -> EvaluatorM $ lift $ prevEnvs      <$> get
    CurRunningJob   -> EvaluatorM $ lift $ curRunningJob <$> get
    EvalCfg         -> EvaluatorM $ lift $ evaluatorCfg <$> get

data EvaluatorMUpdate =
   UpdateDagEU (NodeListUpdate CellState)
 | UpdateCellState NodeId CellUpdate
 | UpdateCurJob CurJobStatus
 | UpdateEnvs [TopStateEx]
 | AppendEnv TopStateEx

data EvaluatorMLabel a where
  NodeListEM    ::           EvaluatorMLabel [NodeId]
  NodeInfo      :: NodeId -> EvaluatorMLabel (Maybe CellState)
  PrevEnvs      ::           EvaluatorMLabel [TopStateEx]
  CurRunningJob ::           EvaluatorMLabel (CurJobStatus)
  EvalCfg       ::           EvaluatorMLabel EvalConfig

-- It's redundant to have both NodeId and TheadId but it defends against
-- possible GHC reuse of ThreadId (I don't know if that can actually happen)
type JobId = (ThreadId, NodeId)
type CurJobStatus = Maybe (JobId, CellIndex)

data EvaluatorState = EvaluatorState
  { evaluatorCfg  :: EvalConfig
  , prevEnvs      :: [TopStateEx]
  , curRunningJob :: CurJobStatus }

data CellStatus =
    Waiting
  | Running            -- TODO: split into compiling/running
  | Complete           -- completed without errors
  | CompleteWithErrors
  | Inert              -- doesn't require running at all
    deriving (Show, Generic)

data CellState  = CellState SourceBlock CellStatus Outputs
     deriving (Show, Generic)

data CellUpdate = CellUpdate (Overwrite CellStatus) Outputs deriving (Show, Generic)

type CellsState  = NodeList       CellState
type CellsUpdate = NodeListUpdate CellState

type CellIndex = Int -- index in the list of cells, not the NodeId

data JobUpdate =
    PartialJobUpdate   Outputs
  | JobComplete        (ExitStatus, TopStateEx)
    deriving (Show)

data EvaluatorMsg =
   SourceUpdate (DagUpdate SourceBlock)
 | JobUpdate JobId JobUpdate
 | Subscribe_E (SubscribeMsg CellsState)

initEvaluatorState :: EvalConfig -> TopStateEx -> EvaluatorState
initEvaluatorState cfg s = EvaluatorState cfg [s] Nothing

launchDagEvaluator :: MonadIO m => EvalConfig -> CellParser -> TopStateEx -> m Evaluator
launchDagEvaluator cfg cellParser env = do
  mailbox <- launchActor do
    let s = initEvaluatorState cfg env
    void $ flip runStateT s $ runIncServerT emptyNodeList $ runEvaluatorM' $
      dagEvaluatorImpl cellParser
  return $ sliceMailbox Subscribe_E mailbox

dagEvaluatorImpl :: CellParser -> EvaluatorM ()
dagEvaluatorImpl cellParser = do
  initDag <- subscribe SourceUpdate cellParser
  processDagUpdate (nodeListAsUpdate initDag) >> flushDiffs
  launchNextJob
  messageLoop \case
    Subscribe_E msg        -> handleSubscribeMsg msg
    SourceUpdate dagUpdate -> do
      processDagUpdate dagUpdate
      flushDiffs
    JobUpdate jobId jobUpdate -> do
      processJobUpdate jobId jobUpdate
      flushDiffs

processJobUpdate :: JobId -> JobUpdate -> EvaluatorM ()
processJobUpdate jobId jobUpdate = do
  getl CurRunningJob >>= \case
    Just (jobId', _) -> when (jobId == jobId') do
      let nodeId = snd jobId
      case jobUpdate of
        JobComplete (exitStatus, newEnv) -> do
          let newStatus = exitStatusAsCellStatus exitStatus
          update $ UpdateCellState nodeId $ CellUpdate (OverwriteWith newStatus) mempty
          update $ UpdateCurJob Nothing
          update $ AppendEnv newEnv
          launchNextJob
          flushDiffs
        PartialJobUpdate result -> update $ UpdateCellState nodeId $ CellUpdate NoChange result
    Nothing -> return () -- this job is a zombie

exitStatusAsCellStatus :: ExitStatus -> CellStatus
exitStatusAsCellStatus = \case
  ExitSuccess -> Complete
  ExitFailure -> CompleteWithErrors

nextCellIndex :: EvaluatorM Int
nextCellIndex = do
  envs <- getl PrevEnvs
  return $ length envs - 1

launchNextJob :: EvaluatorM ()
launchNextJob = do
  cellIndex <- nextCellIndex
  nodeList <- getl NodeListEM
  when (cellIndex < length nodeList) do -- otherwise we're all done
    curEnv <- (!! cellIndex) <$> getl PrevEnvs
    let nodeId = nodeList !! cellIndex
    CellState source _ _ <- fromJust <$> getl (NodeInfo nodeId)
    if isInert source
      then do
        update $ AppendEnv curEnv
        launchNextJob
      else launchJob cellIndex nodeId curEnv

launchJob :: CellIndex -> NodeId -> TopStateEx -> EvaluatorM ()
launchJob cellIndex nodeId env = do
  jobAction <- sourceBlockEvalFun <$> getl EvalCfg
  CellState source _ _ <- fromJust <$> getl (NodeInfo nodeId)
  mailbox <- selfMailbox id
  update $ UpdateCellState nodeId $ CellUpdate (OverwriteWith Running) mempty
  threadId <- liftIO $ forkIO do
    threadId <- myThreadId
    let jobId = (threadId, nodeId)
    let resultsMailbox = sliceMailbox (JobUpdate jobId . PartialJobUpdate) mailbox
    finalEnv <- jobAction resultsMailbox env source
    send mailbox $ JobUpdate jobId $ JobComplete finalEnv
  let jobId = (threadId, nodeId)
  update $ UpdateCurJob (Just (jobId, cellIndex))

computeNumValidCells :: TailUpdate NodeId -> EvaluatorM Int
computeNumValidCells tailUpdate = do
  let nDropped = numDropped tailUpdate
  nTotal <- length <$> getl NodeListEM
  return $ nTotal - nDropped

processDagUpdate :: DagUpdate SourceBlock -> EvaluatorM ()
processDagUpdate (NodeListUpdate tailUpdate mapUpdate) = do
  nValid <- computeNumValidCells tailUpdate
  envs <- getl PrevEnvs
  update $ UpdateEnvs $ take (nValid + 1) envs
  update $ UpdateDagEU $ NodeListUpdate tailUpdate $ mapUpdateMapWithKey mapUpdate
    (\_ (Unchanging source) -> initCellState source)
    (\_ () -> mempty)
  getl CurRunningJob >>= \case
    Nothing -> launchNextJob
    Just ((threadId, _), cellIndex)
      | (cellIndex >= nValid) -> do
          -- Current job is no longer valid. Kill it and restart.
          liftIO $ killThread threadId
          update $ UpdateCurJob Nothing
          launchNextJob
      | otherwise -> return () -- Current job is fine. Let it continue.

isInert :: SourceBlock -> Bool
isInert sb = case sbContents sb of
  TopDecl _   -> False
  Command _ _ -> False
  DeclareForeign _ _ _ -> False
  DeclareCustomLinearization _ _ _ -> False
  Misc misc -> case misc of
    GetNameType _  -> False
    ImportModule _ -> False
    QueryEnv _     -> False
    ProseBlock _  -> True
    CommentLine   -> True
    EmptyLines    -> True
  UnParseable _ _ -> True

initCellState :: SourceBlock -> CellState
initCellState source = do
  let status = if isInert source
        then Inert
        else Waiting
  CellState source status mempty

-- === ToJSON ===

instance ToJSON CellStatus
instance (IncState s, ToJSON s, ToJSON (Delta s)) => ToJSON (NodeListUpdate s)

-- === IncState and related instance ===

instance Semigroup CellUpdate where
  CellUpdate s o <> CellUpdate s' o' = CellUpdate (s<>s') (o<>o')

instance Monoid CellUpdate where
  mempty = CellUpdate mempty mempty

instance IncState CellState where
  type Delta CellState = CellUpdate
  applyDiff (CellState source status result) (CellUpdate status' result') =
    CellState source (fromOverwritable (applyDiff (Overwritable status) status')) (result <> result')