{-# LANGUAGE ScopedTypeVariables #-}
-- | Reordering or pruning the tree in order to prefer or make certain choices.
module Distribution.Solver.Modular.Preference
    ( avoidReinstalls
    , deferSetupChoices
    , deferWeakFlagChoices
    , enforceManualFlags
    , enforcePackageConstraints
    , enforceSingleInstanceRestriction
    , firstGoal
    , preferBaseGoalChoice
    , preferLinked
    , preferPackagePreferences
    , preferReallyEasyGoalChoices
    , requireInstalled
    , sortGoals
    ) where

import Prelude ()
import Distribution.Client.Compat.Prelude

import Data.Function (on)
import qualified Data.List as L
import qualified Data.Map as M
import Control.Monad.Reader hiding (sequence)
import Data.Traversable (sequence)

import Distribution.Solver.Types.Flag
import Distribution.Solver.Types.InstalledPreference
import Distribution.Solver.Types.LabeledPackageConstraint
import Distribution.Solver.Types.OptionalStanza
import Distribution.Solver.Types.PackageConstraint
import Distribution.Solver.Types.PackagePath
import Distribution.Solver.Types.PackagePreferences
import Distribution.Solver.Types.Variable

import Distribution.Solver.Modular.Dependency
import Distribution.Solver.Modular.Flag
import Distribution.Solver.Modular.Package
import qualified Distribution.Solver.Modular.PSQ as P
import Distribution.Solver.Modular.Tree
import Distribution.Solver.Modular.Version
import qualified Distribution.Solver.Modular.ConflictSet as CS
import qualified Distribution.Solver.Modular.WeightedPSQ as W

-- | Update the weights of children under 'PChoice' nodes. 'addWeights' takes a
-- list of weight-calculating functions in order to avoid sorting the package
-- choices multiple times. Each function takes the package name, sorted list of
-- children's versions, and package option. 'addWeights' prepends the new
-- weights to the existing weights, which gives precedence to preferences that
-- are applied later.
addWeights :: [PN -> [Ver] -> POption -> Weight] -> Tree d c -> Tree d c
addWeights fs = trav go
  where
    go :: TreeF d c (Tree d c) -> TreeF d c (Tree d c)
    go (PChoiceF qpn@(Q _ pn) rdm x cs) =
      let sortedVersions = L.sortBy (flip compare) $ L.map version (W.keys cs)
          weights k = [f pn sortedVersions k | f <- fs]

          elemsToWhnf :: [a] -> ()
          elemsToWhnf = foldr seq ()
      in  PChoiceF qpn rdm x
          -- Evaluate the children's versions before evaluating any of the
          -- subtrees, so that 'sortedVersions' doesn't hold onto all of the
          -- subtrees (referenced by cs) and cause a space leak.
          (elemsToWhnf sortedVersions `seq`
             W.mapWeightsWithKey (\k w -> weights k ++ w) cs)
    go x                            = x

addWeight :: (PN -> [Ver] -> POption -> Weight) -> Tree d c -> Tree d c
addWeight f = addWeights [f]

version :: POption -> Ver
version (POption (I v _) _) = v

-- | Prefer to link packages whenever possible.
preferLinked :: Tree d c -> Tree d c
preferLinked = addWeight (const (const linked))
  where
    linked (POption _ Nothing)  = 1
    linked (POption _ (Just _)) = 0

-- Works by setting weights on choice nodes. Also applies stanza preferences.
preferPackagePreferences :: (PN -> PackagePreferences) -> Tree d c -> Tree d c
preferPackagePreferences pcs =
    preferPackageStanzaPreferences pcs .
    addWeights [
          \pn _  opt -> preferred pn opt

        -- Note that we always rank installed before uninstalled, and later
        -- versions before earlier, but we can change the priority of the
        -- two orderings.
        , \pn vs opt -> case preference pn of
                          PreferInstalled -> installed opt
                          PreferLatest    -> latest vs opt
        , \pn vs opt -> case preference pn of
                          PreferInstalled -> latest vs opt
                          PreferLatest    -> installed opt
        ]
  where
    -- Prefer packages with higher version numbers over packages with
    -- lower version numbers.
    latest :: [Ver] -> POption -> Weight
    latest sortedVersions opt =
      let l = length sortedVersions
          index = fromMaybe l $ L.findIndex (<= version opt) sortedVersions
      in  fromIntegral index / fromIntegral l

    preference :: PN -> InstalledPreference
    preference pn =
      let PackagePreferences _ ipref _ = pcs pn
      in  ipref

    -- | Prefer versions satisfying more preferred version ranges.
    preferred :: PN -> POption -> Weight
    preferred pn opt =
      let PackagePreferences vrs _ _ = pcs pn
      in fromIntegral . negate . L.length $
         L.filter (flip checkVR (version opt)) vrs

    -- Prefer installed packages over non-installed packages.
    installed :: POption -> Weight
    installed (POption (I _ (Inst _)) _) = 0
    installed _                          = 1

-- | Traversal that tries to establish package stanza enable\/disable
-- preferences. Works by reordering the branches of stanza choices.
preferPackageStanzaPreferences :: (PN -> PackagePreferences) -> Tree d c -> Tree d c
preferPackageStanzaPreferences pcs = trav go
  where
    go (SChoiceF qsn@(SN (PI (Q pp pn) _) s) rdm gr _tr ts)
      | primaryPP pp && enableStanzaPref pn s =
          -- move True case first to try enabling the stanza
          let ts' = W.mapWeightsWithKey (\k w -> weight k : w) ts
              weight k = if k then 0 else 1
          -- defer the choice by setting it to weak
          in  SChoiceF qsn rdm gr (WeakOrTrivial True) ts'
    go x = x

    enableStanzaPref :: PN -> OptionalStanza -> Bool
    enableStanzaPref pn s =
      let PackagePreferences _ _ spref = pcs pn
      in  s `elem` spref

-- | Helper function that tries to enforce a single package constraint on a
-- given instance for a P-node. Translates the constraint into a
-- tree-transformer that either leaves the subtree untouched, or replaces it
-- with an appropriate failure node.
processPackageConstraintP :: forall d c. QPN
                          -> ConflictSet
                          -> I
                          -> LabeledPackageConstraint
                          -> Tree d c
                          -> Tree d c
processPackageConstraintP qpn c i (LabeledPackageConstraint (PackageConstraint scope prop) src) r =
    if constraintScopeMatches scope qpn
    then go i prop
    else r
  where
    go :: I -> PackageProperty -> Tree d c
    go (I v _) (PackagePropertyVersion vr)
        | checkVR vr v  = r
        | otherwise     = Fail c (GlobalConstraintVersion vr src)
    go _       PackagePropertyInstalled
        | instI i       = r
        | otherwise     = Fail c (GlobalConstraintInstalled src)
    go _       PackagePropertySource
        | not (instI i) = r
        | otherwise     = Fail c (GlobalConstraintSource src)
    go _       _        = r

-- | Helper function that tries to enforce a single package constraint on a
-- given flag setting for an F-node. Translates the constraint into a
-- tree-transformer that either leaves the subtree untouched, or replaces it
-- with an appropriate failure node.
processPackageConstraintF :: forall d c. QPN
                          -> Flag
                          -> ConflictSet
                          -> Bool
                          -> LabeledPackageConstraint
                          -> Tree d c
                          -> Tree d c
processPackageConstraintF qpn f c b' (LabeledPackageConstraint (PackageConstraint scope prop) src) r =
    if constraintScopeMatches scope qpn
    then go prop
    else r
  where
    go :: PackageProperty -> Tree d c
    go (PackagePropertyFlags fa) =
        case L.lookup f fa of
          Nothing            -> r
          Just b | b == b'   -> r
                 | otherwise -> Fail c (GlobalConstraintFlag src)
    go _                             = r

-- | Helper function that tries to enforce a single package constraint on a
-- given flag setting for an F-node. Translates the constraint into a
-- tree-transformer that either leaves the subtree untouched, or replaces it
-- with an appropriate failure node.
processPackageConstraintS :: forall d c. QPN
                          -> OptionalStanza
                          -> ConflictSet
                          -> Bool
                          -> LabeledPackageConstraint
                          -> Tree d c
                          -> Tree d c
processPackageConstraintS qpn s c b' (LabeledPackageConstraint (PackageConstraint scope prop) src) r =
    if constraintScopeMatches scope qpn
    then go prop
    else r
  where
    go :: PackageProperty -> Tree d c
    go (PackagePropertyStanzas ss) =
        if not b' && s `elem` ss then Fail c (GlobalConstraintFlag src)
                                 else r
    go _                               = r

-- | Traversal that tries to establish various kinds of user constraints. Works
-- by selectively disabling choices that have been ruled out by global user
-- constraints.
enforcePackageConstraints :: M.Map PN [LabeledPackageConstraint]
                          -> Tree d c
                          -> Tree d c
enforcePackageConstraints pcs = trav go
  where
    go (PChoiceF qpn@(Q _ pn) rdm gr                    ts) =
      let c = varToConflictSet (P qpn)
          -- compose the transformation functions for each of the relevant constraint
          g = \ (POption i _) -> foldl (\ h pc -> h . processPackageConstraintP qpn c i pc)
                                       id
                                       (M.findWithDefault [] pn pcs)
      in PChoiceF qpn rdm gr        (W.mapWithKey g ts)
    go (FChoiceF qfn@(FN (PI qpn@(Q _ pn) _) f) rdm gr tr m d ts) =
      let c = varToConflictSet (F qfn)
          -- compose the transformation functions for each of the relevant constraint
          g = \ b -> foldl (\ h pc -> h . processPackageConstraintF qpn f c b pc)
                           id
                           (M.findWithDefault [] pn pcs)
      in FChoiceF qfn rdm gr tr m d (W.mapWithKey g ts)
    go (SChoiceF qsn@(SN (PI qpn@(Q _ pn) _) f) rdm gr tr   ts) =
      let c = varToConflictSet (S qsn)
          -- compose the transformation functions for each of the relevant constraint
          g = \ b -> foldl (\ h pc -> h . processPackageConstraintS qpn f c b pc)
                           id
                           (M.findWithDefault [] pn pcs)
      in SChoiceF qsn rdm gr tr     (W.mapWithKey g ts)
    go x = x

-- | Transformation that tries to enforce the rule that manual flags can only be
-- set by the user.
--
-- If there are no constraints on a manual flag, this function prunes all but
-- the default value. If there are constraints, then the flag is allowed to have
-- the values specified by the constraints. Note that the type used for flag
-- values doesn't need to be Bool.
--
-- This function makes an exception for the case where there are multiple goals
-- for a single package (with different qualifiers), and flag constraints for
-- manual flag x only apply to some of those goals. In that case, we allow the
-- unconstrained goals to use the default value for x OR any of the values in
-- the constraints on x (even though the constraints don't apply), in order to
-- allow the unconstrained goals to be linked to the constrained goals. See
-- https://github.com/haskell/cabal/issues/4299.
--
-- This function does not enforce any of the constraints, since that is done by
-- 'enforcePackageConstraints'.
enforceManualFlags :: M.Map PN [LabeledPackageConstraint] -> Tree d c -> Tree d c
enforceManualFlags pcs = trav go
  where
    go (FChoiceF qfn@(FN (PI (Q _ pn) _) fn) rdm gr tr Manual d ts) =
        FChoiceF qfn rdm gr tr Manual d $
          let -- A list of all values specified by constraints on 'fn',
              -- regardless of scope.
              flagConstraintValues :: [Bool]
              flagConstraintValues =
                  [ flagVal
                  | let lpcs = M.findWithDefault [] pn pcs
                  , (LabeledPackageConstraint (PackageConstraint _ (PackagePropertyFlags fa)) _) <- lpcs
                  , (fn', flagVal) <- fa
                  , fn' == fn ]

              -- Prune flag values that are not the default and do not match any
              -- of the constraints.
              restrictToggling :: Eq a => a -> [a] -> a -> Tree d c -> Tree d c
              restrictToggling flagDefault constraintVals flagVal r =
                  if flagVal `elem` constraintVals || flagVal == flagDefault
                  then r
                  else Fail (varToConflictSet (F qfn)) ManualFlag

      in W.mapWithKey (restrictToggling d flagConstraintValues) ts
    go x                                                            = x

-- | Require installed packages.
requireInstalled :: (PN -> Bool) -> Tree d c -> Tree d c
requireInstalled p = trav go
  where
    go (PChoiceF v@(Q _ pn) rdm gr cs)
      | p pn      = PChoiceF v rdm gr (W.mapWithKey installed cs)
      | otherwise = PChoiceF v rdm gr                         cs
      where
        installed (POption (I _ (Inst _)) _) x = x
        installed _ _ = Fail (varToConflictSet (P v)) CannotInstall
    go x          = x

-- | Avoid reinstalls.
--
-- This is a tricky strategy. If a package version is installed already and the
-- same version is available from a repo, the repo version will never be chosen.
-- This would result in a reinstall (either destructively, or potentially,
-- shadowing). The old instance won't be visible or even present anymore, but
-- other packages might have depended on it.
--
-- TODO: It would be better to actually check the reverse dependencies of installed
-- packages. If they're not depended on, then reinstalling should be fine. Even if
-- they are, perhaps this should just result in trying to reinstall those other
-- packages as well. However, doing this all neatly in one pass would require to
-- change the builder, or at least to change the goal set after building.
avoidReinstalls :: (PN -> Bool) -> Tree d c -> Tree d c
avoidReinstalls p = trav go
  where
    go (PChoiceF qpn@(Q _ pn) rdm gr cs)
      | p pn      = PChoiceF qpn rdm gr disableReinstalls
      | otherwise = PChoiceF qpn rdm gr cs
      where
        disableReinstalls =
          let installed = [ v | (_, POption (I v (Inst _)) _, _) <- W.toList cs ]
          in  W.mapWithKey (notReinstall installed) cs

        notReinstall vs (POption (I v InRepo) _) _ | v `elem` vs =
          Fail (varToConflictSet (P qpn)) CannotReinstall
        notReinstall _ _ x =
          x
    go x          = x

-- | Sort all goals using the provided function.
sortGoals :: (Variable QPN -> Variable QPN -> Ordering) -> Tree d c -> Tree d c
sortGoals variableOrder = trav go
  where
    go (GoalChoiceF rdm xs) = GoalChoiceF rdm (P.sortByKeys goalOrder xs)
    go x                    = x

    goalOrder :: Goal QPN -> Goal QPN -> Ordering
    goalOrder = variableOrder `on` (varToVariable . goalToVar)

    varToVariable :: Var QPN -> Variable QPN
    varToVariable (P qpn)                    = PackageVar qpn
    varToVariable (F (FN (PI qpn _) fn))     = FlagVar qpn fn
    varToVariable (S (SN (PI qpn _) stanza)) = StanzaVar qpn stanza

-- | Always choose the first goal in the list next, abandoning all
-- other choices.
--
-- This is unnecessary for the default search strategy, because
-- it descends only into the first goal choice anyway,
-- but may still make sense to just reduce the tree size a bit.
firstGoal :: Tree d c -> Tree d c
firstGoal = trav go
  where
    go (GoalChoiceF rdm xs) = GoalChoiceF rdm (P.firstOnly xs)
    go x                    = x
    -- Note that we keep empty choice nodes, because they mean success.

-- | Transformation that tries to make a decision on base as early as
-- possible by pruning all other goals when base is available. In nearly
-- all cases, there's a single choice for the base package. Also, fixing
-- base early should lead to better error messages.
preferBaseGoalChoice :: Tree d c -> Tree d c
preferBaseGoalChoice = trav go
  where
    go (GoalChoiceF rdm xs) = GoalChoiceF rdm (P.filterIfAnyByKeys isBase xs)
    go x                    = x

    isBase :: Goal QPN -> Bool
    isBase (Goal (P (Q _pp pn)) _) = unPN pn == "base"
    isBase _                       = False

-- | Deal with setup dependencies after regular dependencies, so that we can
-- will link setup dependencies against package dependencies when possible
deferSetupChoices :: Tree d c -> Tree d c
deferSetupChoices = trav go
  where
    go (GoalChoiceF rdm xs) = GoalChoiceF rdm (P.preferByKeys noSetup xs)
    go x                    = x

    noSetup :: Goal QPN -> Bool
    noSetup (Goal (P (Q (PackagePath _ns (QualSetup _)) _)) _) = False
    noSetup _                                                  = True

-- | Transformation that tries to avoid making weak flag choices early.
-- Weak flags are trivial flags (not influencing dependencies) or such
-- flags that are explicitly declared to be weak in the index.
deferWeakFlagChoices :: Tree d c -> Tree d c
deferWeakFlagChoices = trav go
  where
    go (GoalChoiceF rdm xs) = GoalChoiceF rdm (P.prefer noWeakFlag (P.prefer noWeakStanza xs))
    go x                    = x

    noWeakStanza :: Tree d c -> Bool
    noWeakStanza (SChoice _ _ _ (WeakOrTrivial True)   _) = False
    noWeakStanza _                                        = True

    noWeakFlag :: Tree d c -> Bool
    noWeakFlag (FChoice _ _ _ (WeakOrTrivial True) _ _ _) = False
    noWeakFlag _                                          = True

-- | Transformation that prefers goals with lower branching degrees.
--
-- When a goal choice node has at least one goal with zero or one children, this
-- function prunes all other goals. This transformation can help the solver find
-- a solution in fewer steps by allowing it to backtrack sooner when it is
-- exploring a subtree with no solutions. However, each step is more expensive.
preferReallyEasyGoalChoices :: Tree d c -> Tree d c
preferReallyEasyGoalChoices = trav go
  where
    go (GoalChoiceF rdm xs) = GoalChoiceF rdm (P.filterIfAny zeroOrOneChoices xs)
    go x                    = x

-- | Monad used internally in enforceSingleInstanceRestriction
--
-- For each package instance we record the goal for which we picked a concrete
-- instance. The SIR means that for any package instance there can only be one.
type EnforceSIR = Reader (Map (PI PN) QPN)

-- | Enforce ghc's single instance restriction
--
-- From the solver's perspective, this means that for any package instance
-- (that is, package name + package version) there can be at most one qualified
-- goal resolving to that instance (there may be other goals _linking_ to that
-- instance however).
enforceSingleInstanceRestriction :: Tree d c -> Tree d c
enforceSingleInstanceRestriction = (`runReader` M.empty) . cata go
  where
    go :: TreeF d c (EnforceSIR (Tree d c)) -> EnforceSIR (Tree d c)

    -- We just verify package choices.
    go (PChoiceF qpn rdm gr cs) =
      PChoice qpn rdm gr <$> sequence (W.mapWithKey (goP qpn) cs)
    go _otherwise =
      innM _otherwise

    -- The check proper
    goP :: QPN -> POption -> EnforceSIR (Tree d c) -> EnforceSIR (Tree d c)
    goP qpn@(Q _ pn) (POption i linkedTo) r = do
      let inst = PI pn i
      env <- ask
      case (linkedTo, M.lookup inst env) of
        (Just _, _) ->
          -- For linked nodes we don't check anything
          r
        (Nothing, Nothing) ->
          -- Not linked, not already used
          local (M.insert inst qpn) r
        (Nothing, Just qpn') -> do
          -- Not linked, already used. This is an error
          return $ Fail (CS.union (varToConflictSet (P qpn)) (varToConflictSet (P qpn'))) MultipleInstances