{-# LANGUAGE RecordWildCards, PatternGuards, ViewPatterns #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE DeriveDataTypeable, GeneralizedNewtypeDeriving #-}

module Development.Shake.Database(
    Trace(..),
    Database, withDatabase,
    listDepends, lookupDependencies,
    Ops(..), build, Depends,
    progress,
    Stack, emptyStack, topStack, showStack, showTopStack,
    toReport, checkValid, listLive
    ) where

import Prelude(); import General.Prelude

import Development.Shake.Classes
import General.Binary
import Development.Shake.Pool
import Development.Shake.Value
import Development.Shake.Errors
import Development.Shake.Storage
import Development.Shake.Types
import Development.Shake.Special
import Development.Shake.Profile
import Development.Shake.Monad
import General.String
import General.Intern as Intern

import Numeric.Extra
import Control.Applicative
import Control.Exception
import Control.Monad.Extra
import Control.Concurrent.Extra
import qualified Data.HashSet as Set
import qualified Data.HashMap.Strict as Map
import Data.IORef.Extra
import Data.Maybe
import Data.List
import System.Time.Extra

type Map = Map.HashMap


---------------------------------------------------------------------
-- UTILITY TYPES

newtype Step = Step Word32 deriving (Eq,Ord,Show,Binary,NFData,Hashable,Typeable)

incStep (Step i) = Step $ i + 1


---------------------------------------------------------------------
-- CALL STACK

data Stack = Stack (Maybe Key) [Id] !(Set.HashSet Id)

showStack :: Database -> Stack -> IO [String]
showStack Database{..} (Stack _ xs _) = do
    status <- withLock lock $ readIORef status
    return $ reverse $ map (maybe "<unknown>" (show . fst) . flip Map.lookup status) xs

addStack :: Id -> Key -> Stack -> Stack
addStack x key (Stack _ xs set) = Stack (Just key) (x:xs) (Set.insert x set)

showTopStack :: Stack -> String
showTopStack = maybe "<unknown>" show . topStack

topStack :: Stack -> Maybe Key
topStack (Stack key _ _) = key

checkStack :: [Id] -> Stack -> Maybe Id
checkStack new (Stack _ old set)
    | bad:_ <- filter (`Set.member` set) new = Just bad
    | otherwise = Nothing

emptyStack :: Stack
emptyStack = Stack Nothing [] Set.empty


---------------------------------------------------------------------
-- CENTRAL TYPES

data Trace = Trace BS Float Float -- (message, start, end)
    deriving Show

instance NFData Trace where
    rnf (Trace a b c) = rnf a `seq` rnf b `seq` rnf c

-- | Invariant: The database does not have any cycles where a Key depends on itself
data Database = Database
    {lock :: Lock
    ,intern :: IORef (Intern Key)
    ,status :: IORef (Map Id (Key, Status))
    ,step :: Step
    ,journal :: Id -> (Key, Status {- Loaded or Missing -}) -> IO ()
    ,diagnostic :: String -> IO () -- logging function
    ,assume :: Maybe Assume
    }

data Status
    = Ready Result -- I have a value
    | Error SomeException -- I have been run and raised an error
    | Loaded Result -- Loaded from the database
    | Waiting Pending (Maybe Result) -- Currently checking if I am valid or building
    | Missing -- I am only here because I got into the Intern table
      deriving Show

data Result = Result
    {result :: Value -- the result associated with the Key
    ,built :: {-# UNPACK #-} !Step -- when it was actually run
    ,changed :: {-# UNPACK #-} !Step -- the step for deciding if it's valid
    ,depends :: [[Id]] -- dependencies
    ,execution :: {-# UNPACK #-} !Float -- how long it took when it was last run (seconds)
    ,traces :: [Trace] -- a trace of the expensive operations (start/end in seconds since beginning of run)
    } deriving Show


newtype Pending = Pending (IORef (IO ()))
    -- you must run this action when you finish, while holding DB lock
    -- after you have set the result to Error or Ready

instance Show Pending where show _ = "Pending"


statusType Ready{} = "Ready"
statusType Error{} = "Error"
statusType Loaded{} = "Loaded"
statusType Waiting{} = "Waiting"
statusType Missing{} = "Missing"

isError Error{} = True; isError _ = False
isWaiting Waiting{} = True; isWaiting _ = False
isReady Ready{} = True; isReady _ = False


-- All the waiting operations are only valid when isWaiting
type Waiting = Status

afterWaiting :: Waiting -> IO () -> IO ()
afterWaiting (Waiting (Pending p) _) act = modifyIORef' p (>> act)

newWaiting :: Maybe Result -> IO Waiting
newWaiting r = do ref <- newIORef $ return (); return $ Waiting (Pending ref) r

runWaiting :: Waiting -> IO ()
runWaiting (Waiting (Pending p) _) = join $ readIORef p

-- Wait for a set of actions to complete
-- If the action returns True, the function will not be called again
-- If the first argument is True, the thing is ended
waitFor :: [(a, Waiting)] -> (Bool -> a -> IO Bool) -> IO ()
waitFor ws@(_:_) act = do
    todo <- newIORef $ length ws
    forM_ ws $ \(k,w) -> afterWaiting w $ do
        t <- readIORef todo
        when (t /= 0) $ do
            b <- act (t == 1) k
            writeIORef' todo $ if b then 0 else t - 1


getResult :: Status -> Maybe Result
getResult (Ready r) = Just r
getResult (Loaded r) = Just r
getResult (Waiting _ r) = r
getResult _ = Nothing


---------------------------------------------------------------------
-- OPERATIONS

newtype Depends = Depends {fromDepends :: [Id]}
    deriving (NFData)


data Ops = Ops
    {stored :: Key -> IO (Maybe Value)
        -- ^ Given a Key, find the value stored on disk
    ,equal :: Key -> Value -> Value -> EqualCost
        -- ^ Given both Values, see if they are equal and how expensive that check was
    ,execute :: Stack -> Key -> Capture (Either SomeException (Value, [Depends], Seconds, [Trace]))
        -- ^ Given a stack and a key, either raise an exception or successfully build it
    }


-- | Return either an exception (crash), or (how much time you spent waiting, the value)
build :: Pool -> Database -> Ops -> Stack -> [Key] -> Capture (Either SomeException (Seconds,Depends,[Value]))
build pool Database{..} Ops{..} stack ks continue = do
    join $ withLock lock $ do
        is <- forM ks $ \k -> do
            is <- readIORef intern
            case Intern.lookup k is of
                Just i -> return i
                Nothing -> do
                    (is, i) <- return $ Intern.add k is
                    writeIORef' intern is
                    modifyIORef' status $ Map.insert i (k,Missing)
                    return i

        whenJust (checkStack is stack) $ \bad -> do
            status <- readIORef status
            uncurry errorRuleRecursion $ case Map.lookup bad status of
                Nothing -> (Nothing, Nothing)
                Just (k,_) -> (Just $ typeKey k, Just $ show k)

        vs <- mapM (reduce stack) is
        let errs = [e | Error e <- vs]
        if all isReady vs then
            return $ continue $ Right (0, Depends is, [result r | Ready r <- vs])
         else if not $ null errs then
            return $ continue $ Left $ head errs
         else do
            time <- offsetTime
            let done x = do
                    case x of
                        Left e -> addPoolPriority pool $ continue $ Left e
                        Right v -> addPool pool $ do dur <- time; continue $ Right (dur, Depends is, v)
                    return True
            waitFor (filter (isWaiting . snd) $ zip is vs) $ \finish i -> do
                s <- readIORef status
                case Map.lookup i s of
                    Just (_, Error e) -> done $ Left e -- on error make sure we immediately kick off our parent
                    Just (_, Ready{}) | finish -> done $ Right [result r | i <- is, let Ready r = snd $ fromJust $ Map.lookup i s]
                                      | otherwise -> return False
            return $ return ()
    where
        (#=) :: Id -> (Key, Status) -> IO Status
        i #= (k,v) = do
            s <- readIORef status
            writeIORef' status $ Map.insert i (k,v) s
            diagnostic $ maybe "Missing" (statusType . snd) (Map.lookup i s) ++ " -> " ++ statusType v ++ ", " ++ maybe "<unknown>" (show . fst) (Map.lookup i s)
            return v

        atom x = let s = show x in if ' ' `elem` s then "(" ++ s ++ ")" else s

        -- Rules for each eval* function
        -- * Must NOT lock
        -- * Must have an equal return to what is stored in the db at that point
        -- * Must not return Loaded

        reduce :: Stack -> Id -> IO Status
        reduce stack i = do
            s <- readIORef status
            case Map.lookup i s of
                Nothing -> err $ "interned value missing from database, " ++ show i
                Just (k, Missing) -> run stack i k Nothing
                Just (k, Loaded r) -> do
                    let out b = diagnostic $ "valid " ++ show b ++ " for " ++ atom k ++ " " ++ atom (result r)
                    let continue r = out True >> check stack i k r (depends r)
                    let rebuild = out False >> run stack i k (Just r)
                    case assume of
                        Just AssumeDirty -> rebuild
                        Just AssumeSkip -> continue r
                        _ -> do
                            s <- stored k
                            case s of
                                Just s -> case equal k (result r) s of
                                    NotEqual -> rebuild
                                    EqualCheap -> continue r
                                    EqualExpensive -> do
                                        -- warning, have the db lock while appending (may harm performance)
                                        r <- return r{result=s}
                                        journal i (k, Loaded r)
                                        i #= (k, Loaded r)
                                        continue r
                                _ -> rebuild
                Just (k, res) -> return res

        run :: Stack -> Id -> Key -> Maybe Result -> IO Waiting
        run stack i k r = do
            w <- newWaiting r
            addPool pool $ do
                let reply res =  do
                        ans <- withLock lock $ do
                            ans <- i #= (k, res)
                            runWaiting w
                            return ans
                        case ans of
                            Ready r -> do
                                diagnostic $ "result " ++ atom k ++ " = " ++ atom (result r)
                                journal i (k, Loaded r) -- leave the DB lock before appending
                            Error _ -> do
                                diagnostic $ "result " ++ atom k ++ " = error"
                                journal i (k, Missing)
                            _ -> return ()
                let norm = execute (addStack i k stack) k $ \res ->
                        reply $ case res of
                            Left err -> Error err
                            Right (v,deps,(doubleToFloat -> execution),traces) ->
                                let c | Just r <- r, result r == v = changed r
                                      | otherwise = step
                                in Ready Result{result=v,changed=c,built=step,depends=map fromDepends deps,..}

                case r of
                    Just r | assume == Just AssumeClean -> do
                            v <- stored k
                            case v of
                                Just v -> reply $ Ready r{result=v}
                                Nothing -> norm
                    _ -> norm
            i #= (k, w)

        check :: Stack -> Id -> Key -> Result -> [[Id]] -> IO Status
        check stack i k r [] =
            i #= (k, Ready r)
        check stack i k r (ds:rest) = do
            vs <- mapM (reduce (addStack i k stack)) ds
            let ws = filter (isWaiting . snd) $ zip ds vs
            if any isError vs || any (> built r) [changed | Ready Result{..} <- vs] then
                run stack i k $ Just r
             else if null ws then
                check stack i k r rest
             else do
                self <- newWaiting $ Just r
                waitFor ws $ \finish d -> do
                    s <- readIORef status
                    let buildIt = do
                            b <- run stack i k $ Just r
                            afterWaiting b $ runWaiting self
                            return True
                    case Map.lookup d s of
                        Just (_, Error{}) -> buildIt
                        Just (_, Ready r2)
                            | changed r2 > built r -> buildIt
                            | finish -> do
                                res <- check stack i k r rest
                                if not $ isWaiting res
                                    then runWaiting self
                                    else afterWaiting res $ runWaiting self
                                return True
                            | otherwise -> return False
                i #= (k, self)


---------------------------------------------------------------------
-- PROGRESS

progress :: Database -> IO Progress
progress Database{..} = do
    s <- readIORef status
    return $ foldl' f mempty $ map snd $ Map.elems s
    where
        g = floatToDouble

        f s (Ready Result{..}) = if step == built
            then s{countBuilt = countBuilt s + 1, timeBuilt = timeBuilt s + g execution}
            else s{countSkipped = countSkipped s + 1, timeSkipped = timeSkipped s + g execution}
        f s (Loaded Result{..}) = s{countUnknown = countUnknown s + 1, timeUnknown = timeUnknown s + g execution}
        f s (Waiting _ r) =
            let (d,c) = timeTodo s
                t | Just Result{..} <- r = let d2 = d + g execution in d2 `seq` (d2,c)
                  | otherwise = let c2 = c + 1 in c2 `seq` (d,c2)
            in s{countTodo = countTodo s + 1, timeTodo = t}
        f s _ = s


---------------------------------------------------------------------
-- QUERY DATABASE

-- | Given a map of representing a dependency order (with a show for error messages), find an ordering for the items such
--   that no item points to an item before itself.
--   Raise an error if you end up with a cycle.
dependencyOrder :: (Eq a, Hashable a) => (a -> String) -> Map a [a] -> [a]
-- Algorithm:
--    Divide everyone up into those who have no dependencies [Id]
--    And those who depend on a particular Id, Dep :-> Maybe [(Key,[Dep])]
--    Where d :-> Just (k, ds), k depends on firstly d, then remaining on ds
--    For each with no dependencies, add to list, then take its dep hole and
--    promote them either to Nothing (if ds == []) or into a new slot.
--    k :-> Nothing means the key has already been freed
dependencyOrder shw status = f (map fst noDeps) $ Map.map Just $ Map.fromListWith (++) [(d, [(k,ds)]) | (k,d:ds) <- hasDeps]
    where
        (noDeps, hasDeps) = partition (null . snd) $ Map.toList status

        f [] mp | null bad = []
                | otherwise = error $ unlines $
                    "Internal invariant broken, database seems to be cyclic" :
                    map ("    " ++) bad ++
                    ["... plus " ++ show (length badOverflow) ++ " more ..." | not $ null badOverflow]
            where (bad,badOverflow) = splitAt 10 $ [shw i | (i, Just _) <- Map.toList mp]

        f (x:xs) mp = x : f (now++xs) later
            where Just free = Map.lookupDefault (Just []) x mp
                  (now,later) = foldl' g ([], Map.insert x Nothing mp) free

        g (free, mp) (k, []) = (k:free, mp)
        g (free, mp) (k, d:ds) = case Map.lookupDefault (Just []) d mp of
            Nothing -> g (free, mp) (k, ds)
            Just todo -> (free, Map.insert d (Just $ (k,ds) : todo) mp)


-- | Eliminate all errors from the database, pretending they don't exist
resultsOnly :: Map Id (Key, Status) -> Map Id (Key, Result)
resultsOnly mp = Map.map (\(k, v) -> (k, let Just r = getResult v in r{depends = map (filter (isJust . flip Map.lookup keep)) $ depends r})) keep
    where keep = Map.filter (isJust . getResult . snd) mp

removeStep :: Map Id (Key, Result) -> Map Id (Key, Result)
removeStep = Map.filter (\(k,_) -> k /= stepKey)

toReport :: Database -> IO [ProfileEntry]
toReport Database{..} = do
    status <- fmap (removeStep . resultsOnly) $ readIORef status
    let order = let shw i = maybe "<unknown>" (show . fst) $ Map.lookup i status
                in dependencyOrder shw $ Map.map (concat . depends . snd) status
        ids = Map.fromList $ zip order [0..]

        steps = let xs = Set.toList $ Set.fromList $ concat [[changed, built] | (_,Result{..}) <- Map.elems status]
                in Map.fromList $ zip (sortBy (flip compare) xs) [0..]

        f (k, Result{..}) = ProfileEntry
            {prfName = show k
            ,prfBuilt = fromStep built
            ,prfChanged = fromStep changed
            ,prfDepends = mapMaybe (`Map.lookup` ids) (concat depends)
            ,prfExecution = floatToDouble execution
            ,prfTraces = map fromTrace traces
            }
            where fromStep i = fromJust $ Map.lookup i steps
                  fromTrace (Trace a b c) = ProfileTrace (unpack a) (floatToDouble b) (floatToDouble c)
    return [maybe (err "toReport") f $ Map.lookup i status | i <- order]


checkValid :: Database -> (Key -> IO (Maybe Value)) -> (Key -> Value -> Value -> EqualCost) -> [(Key, Key)] -> IO ()
checkValid Database{..} stored equal missing = do
    status <- readIORef status
    intern <- readIORef intern
    diagnostic "Starting validity/lint checking"

    -- Do not use a forM here as you use too much stack space
    bad <- (\f -> foldM f [] (Map.toList status)) $ \seen (i,v) -> case v of
        (key, Ready Result{..}) -> do
            now <- stored key
            let good = maybe False ((==) EqualCheap . equal key result) now
            diagnostic $ "Checking if " ++ show key ++ " is " ++ show result ++ ", " ++ if good then "passed" else "FAILED"
            return $ [(key, result, now) | not good && not (specialAlwaysRebuilds result)] ++ seen
        _ -> return seen
    unless (null bad) $ do
        let n = length bad
        errorStructured
            ("Lint checking error - " ++ (if n == 1 then "value has" else show n ++ " values have")  ++ " changed since being depended upon")
            (intercalate [("",Just "")] [ [("Key", Just $ show key),("Old", Just $ show result),("New", Just $ maybe "<missing>" show now)]
                                        | (key, result, now) <- bad])
            ""

    bad <- return [(parent,key) | (parent, key) <- missing, isJust $ Intern.lookup key intern]
    unless (null bad) $ do
        let n = length bad
        errorStructured
            ("Link checking error - " ++ (if n == 1 then "value" else show n ++ " values") ++ " did not have " ++ (if n == 1 then "its" else "their") ++ " creation tracked")
            (intercalate [("",Just "")] [ [("Rule", Just $ show parent), ("Created", Just $ show key)] | (parent,key) <- bad])
            ""

    diagnostic "Validity/lint check passed"


listLive :: Database -> IO [Key]
listLive Database{..} = do
    diagnostic "Listing live keys"
    status <- readIORef status
    return [k | (k, Ready{}) <- Map.elems status]


listDepends :: Database -> Depends -> IO [Key]
listDepends Database{..} (Depends xs) =
    withLock lock $ do
        status <- readIORef status
        return $ map (fst . fromJust . flip Map.lookup status) xs

lookupDependencies :: Database -> Key -> IO [Key]
lookupDependencies Database{..} k = do
    withLock lock $ do
        intern <- readIORef intern
        status <- readIORef status
        let Just i = Intern.lookup k intern
        let Just (_, Ready r) = Map.lookup i status
        return $ map (fst . fromJust . flip Map.lookup status) $ concat $ depends r


---------------------------------------------------------------------
-- STORAGE

-- To simplify journaling etc we smuggle the Step in the database, with a special StepKey
newtype StepKey = StepKey ()
    deriving (Show,Eq,Typeable,Hashable,Binary,NFData)

stepKey :: Key
stepKey = newKey $ StepKey ()

toStepResult :: Step -> Result
toStepResult i = Result (newValue i) i i [] 0 []

fromStepResult :: Result -> Step
fromStepResult = fromValue . result


withDatabase :: ShakeOptions -> (String -> IO ()) -> (Database -> IO a) -> IO a
withDatabase opts diagnostic act = do
    registerWitness $ StepKey ()
    registerWitness $ Step 0
    witness <- currentWitness
    withStorage opts diagnostic witness $ \mp2 journal -> do
        let mp1 = Intern.fromList [(k, i) | (i, (k,_)) <- Map.toList mp2]

        (mp1, stepId) <- case Intern.lookup stepKey mp1 of
            Just stepId -> return (mp1, stepId)
            Nothing -> do
                (mp1, stepId) <- return $ Intern.add stepKey mp1
                return (mp1, stepId)

        intern <- newIORef mp1
        status <- newIORef mp2
        let step = case Map.lookup stepId mp2 of
                        Just (_, Loaded r) -> incStep $ fromStepResult r
                        _ -> Step 1
        journal stepId (stepKey, Loaded $ toStepResult step)
        lock <- newLock
        act Database{assume=shakeAssume opts,..}


instance BinaryWith Witness Result where
    putWith ws (Result x1 x2 x3 x4 x5 x6) = putWith ws x1 >> put x2 >> put x3 >> put (BinList $ map BinList x4) >> put (BinFloat x5) >> put (BinList x6)
    getWith ws = (\x1 x2 x3 (BinList x4) (BinFloat x5) (BinList x6) -> Result x1 x2 x3 (map fromBinList x4) x5 x6) <$>
        getWith ws <*> get <*> get <*> get <*> get <*> get

instance Binary Trace where
    put (Trace a b c) = put a >> put (BinFloat b) >> put (BinFloat c)
    get = (\a (BinFloat b) (BinFloat c) -> Trace a b c) <$> get <*> get <*> get

instance BinaryWith Witness Status where
    putWith ctx Missing = putWord8 0
    putWith ctx (Loaded x) = putWord8 1 >> putWith ctx x
    putWith ctx x = err $ "putWith, Cannot write Status with constructor " ++ statusType x
    getWith ctx = do i <- getWord8; if i == 0 then return Missing else fmap Loaded $ getWith ctx