module NixTree.StorePath
  ( StoreName (..),
    storeNameToPath,
    storeNameToText,
    storeNameToShortText,
    storeNameToSplitShortText,
    storeNameToShortTextWithDisambiguation,
    NixPathSignature (..),
    StorePath (..),
    Installable (..),
    StoreEnv (..),
    StoreEnvOptions (..),
    withStoreEnv,
    seLookup,
    seAll,
    seGetRoots,
    seBottomUp,
    seFetchRefs,
    mkStoreName,
    storeEnvToDot,
  )
where

import Data.Aeson (FromJSON (..), Value (..), decode, eitherDecode, (.!=), (.:), (.:?))
import qualified Data.Aeson.Key as K
import qualified Data.Aeson.KeyMap as KM
import Data.Aeson.Types (Parser)
import qualified Data.ByteString.Lazy as BL
import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet as HS
import qualified Data.Set as Set
import qualified Data.Text as T
import qualified Dot
import System.FilePath.Posix (addTrailingPathSeparator, splitDirectories, (</>))
import System.IO (hPutStrLn)
import System.Process.Typed (proc, readProcessStdout_)

-- Technically these both are filepaths. However, most people use the default "/nix/store",
-- hence special casing it speeds things up.
data NixStore
  = NixStore
  | NixStoreCustom FilePath
  deriving stock (Show, Eq, Ord, Generic)
  deriving anyclass (NFData, Hashable)

mkNixStore :: FilePath -> NixStore
mkNixStore i' =
  let i = addTrailingPathSeparator i'
   in if i == "/nix/store/" then NixStore else NixStoreCustom i

unNixStore :: NixStore -> FilePath
unNixStore NixStore = "/nix/store/"
unNixStore (NixStoreCustom fp) = fp

getStoreDir :: Maybe FilePath -> IO NixStore
getStoreDir seoNixStore = do
  let prog = "nix-instantiate"
      args =
        ["--eval", "--expr", "(builtins.storeDir)", "--json"]
          ++ ( case seoNixStore of
                 Nothing -> []
                 Just url -> ["--option", "store", url]
             )
  out <-
    readProcessStdout_ (proc prog args)
      <&> fmap mkNixStore
      . decode @FilePath
  case out of
    Nothing -> fail $ "Error interpreting output of: " ++ show (prog, args)
    Just p -> return p

--------------------------------------------------------------------------------

data StoreName = StoreName NixStore Text
  deriving stock (Show, Eq, Ord, Generic)
  deriving anyclass (NFData, Hashable)

mkStoreName :: NixStore -> FilePath -> Maybe StoreName
mkStoreName ns path = do
  let ps = unNixStore ns
  guard $ ps `isPrefixOf` path
  let ds = splitDirectories (drop (length ps) path)
  sn <- listToMaybe ds
  return $ StoreName ns (toText sn)

storeNameToText :: StoreName -> Text
storeNameToText (StoreName _ n) = n

storeNameToPath :: StoreName -> FilePath
storeNameToPath (StoreName ns sn) = unNixStore ns </> toString sn

storeNameToShortText :: StoreName -> Text
storeNameToShortText = snd . storeNameToSplitShortText

storeNameToShortTextWithDisambiguation :: Int -> StoreName -> Text
storeNameToShortTextWithDisambiguation 0 sn = storeNameToShortText sn
storeNameToShortTextWithDisambiguation n (StoreName ns sn) =
  let (f, s) = storeNameToSplitShortText (StoreName ns sn)
   in T.take n f <> "...-" <> s

storeNameToSplitShortText :: StoreName -> (Text, Text)
storeNameToSplitShortText txt =
  case T.span (/= '-') $ storeNameToText txt of
    (f, s) | Just (c, s'') <- T.uncons s -> (T.snoc f c, s'')
    e -> e

--------------------------------------------------------------------------------

data NixVersion
  = NixOlder
  | Nix2_4
  | NixNewer
  | NixUnknown
  deriving (Show, Eq, Ord)

getNixVersion :: IO NixVersion
getNixVersion = do
  out <- decodeUtf8 . BL.toStrict <$> readProcessStdout_ (proc "nix" ["--version"])

  -- Parses strings like:
  --  nix (Nix) 2.6.0pre20211217_6e6e998
  --  nix (Nix) 2.5.1
  return . fromMaybe NixUnknown $ do
    -- get the last space delimited part
    ver <-
      out
        & T.splitOn " "
        & viaNonEmpty last

    -- split by ".", take the first two, and convert them to numbers
    (major, minor) <- do
      (maT, miT) <- case T.splitOn "." ver of
        p1 : p2 : _ -> Just (p1, p2)
        _ -> Nothing
      ma <- readMaybe @Natural (toString maT)
      mi <- readMaybe @Natural (toString miT)
      return (ma, mi)

    -- map it to the sum
    return $ case compare (major, minor) (2, 4) of
      LT -> NixOlder
      EQ -> Nix2_4
      GT -> NixNewer

--------------------------------------------------------------------------------

data StorePath ref payload = StorePath
  { spName :: StoreName,
    spSize :: Int,
    spRefs :: [ref],
    spPayload :: payload,
    spSignatures :: [NixPathSignature]
  }
  deriving (Show, Eq, Ord, Functor, Generic)

instance (NFData a, NFData b) => NFData (StorePath a b)

newtype Installable = Installable {installableToText :: Text}

--------------------------------------------------------------------------------

data PathInfoOptions = PathInfoOptions
  { pioIsRecursive :: Bool,
    pioIsDerivation :: Bool,
    pioIsImpure :: Bool,
    pioStoreURL :: Maybe FilePath,
    pioFile :: Maybe FilePath
  }

getPathInfo :: NixStore -> NixVersion -> PathInfoOptions -> NonEmpty Installable -> IO (NonEmpty (StorePath StoreName ()))
getPathInfo nixStore nixVersion options names = do
  infos <-
    eitherDecode @NixPathOutput
      <$> readProcessStdout_
        ( proc
            "nix"
            ( ["path-info", "--json"]
                ++ ["--impure" | options & pioIsImpure]
                ++ ["--recursive" | options & pioIsRecursive]
                ++ ["--derivation" | (options & pioIsDerivation) && nixVersion >= Nix2_4]
                ++ ( case options & pioStoreURL of
                       Nothing -> []
                       Just url -> ["--store", url]
                   )
                ++ ( case options & pioFile of
                       Nothing -> []
                       Just file -> ["--file", file]
                   )
                ++ (if nixVersion >= Nix2_4 then ["--extra-experimental-features", "nix-command flakes"] else [])
                ++ map (toString . installableToText) (toList names)
            )
        )
      >>= either (\e -> fail $ "Failed parsing nix path-info output: " ++ show e) return
      >>= mapM assertValidInfo . npoResults
      >>= maybe (fail "invariant violation: getPathInfo returned []") return . nonEmpty

  mapM infoToStorePath infos
  where
    infoToStorePath NixPathInfo {npiPath, npiNarSize, npiReferences, npiSignatures} = do
      name <- mkStoreNameIO npiPath
      refs <- filter (/= name) <$> mapM mkStoreNameIO npiReferences
      return $
        StorePath
          { spName = name,
            spRefs = refs,
            spSize = npiNarSize,
            spSignatures = npiSignatures,
            spPayload = ()
          }
    mkStoreNameIO p =
      maybe
        (fail $ "Failed parsing Nix store path: " ++ show p)
        return
        (mkStoreName nixStore p)

    assertValidInfo (NixPathInfoValid pathinfo) = return pathinfo
    assertValidInfo (NixPathInfoInvalid path) =
      fail $ "Invalid path: " ++ path ++ ". Make sure that it is built, or pass '--derivation' if you want to work on the derivation."

--------------------------------------------------------------------------------

data StoreEnv payload = StoreEnv
  { sePaths :: HashMap StoreName (StorePath StoreName payload),
    seRoots :: NonEmpty StoreName
  }
  deriving (Functor, Generic, NFData)

data StoreEnvOptions = StoreEnvOptions
  { seoIsDerivation :: Bool,
    seoIsImpure :: Bool,
    seoStoreURL :: Maybe String,
    seoFile :: Maybe String
  }

withStoreEnv ::
  forall m a.
  (MonadIO m) =>
  StoreEnvOptions ->
  NonEmpty Installable ->
  (StoreEnv () -> m a) ->
  m a
withStoreEnv StoreEnvOptions {seoIsDerivation, seoIsImpure, seoStoreURL, seoFile} names cb = do
  nixStore <- liftIO $ getStoreDir seoStoreURL

  -- See: https://github.com/utdemir/nix-tree/issues/12
  nixVersion <- liftIO getNixVersion

  when (seoIsDerivation && nixVersion < Nix2_4) $
    liftIO $
      hPutStrLn stderr "Warning: --derivation flag is ignored on Nix versions older than 2.4."

  roots <-
    liftIO $
      getPathInfo
        nixStore
        nixVersion
        (PathInfoOptions {pioIsDerivation = seoIsDerivation, pioIsRecursive = False, pioIsImpure = seoIsImpure, pioStoreURL = seoStoreURL, pioFile = seoFile})
        names

  paths <-
    liftIO $
      getPathInfo
        nixStore
        nixVersion
        (PathInfoOptions {pioIsDerivation = seoIsDerivation, pioIsRecursive = True, pioIsImpure = seoIsImpure, pioStoreURL = seoStoreURL, pioFile = seoFile})
        (Installable . toText . storeNameToPath . spName <$> roots)

  let env =
        StoreEnv
          ( paths
              & toList
              & map (\p@StorePath {spName} -> (spName, p))
              & HM.fromList
          )
          (roots <&> spName)
  cb env

seLookup :: StoreEnv a -> StoreName -> StorePath StoreName a
seLookup StoreEnv {sePaths} name =
  fromMaybe
    (error $ "invariant violation, StoreName not found: " <> show name)
    (HM.lookup name sePaths)

seAll :: StoreEnv a -> NonEmpty (StorePath StoreName a)
seAll StoreEnv {sePaths} = case HM.elems sePaths of
  [] -> error "invariant violation: no paths"
  (x : xs) -> x :| xs

seGetRoots :: StoreEnv a -> NonEmpty (StorePath StoreName a)
seGetRoots env@StoreEnv {seRoots} =
  fmap (seLookup env) seRoots

seFetchRefs ::
  StoreEnv a ->
  (StoreName -> Bool) ->
  NonEmpty StoreName ->
  [StorePath StoreName a]
seFetchRefs env predicate =
  fst
    . foldl'
      (\(acc, visited) name -> go acc visited name)
      ([], HS.empty)
  where
    go acc visited name
      | HS.member name visited = (acc, visited)
      | not (predicate name) = (acc, visited)
      | otherwise =
          let sp@StorePath {spRefs} = seLookup env name
           in foldl'
                (\(acc', visited') name' -> go acc' visited' name')
                (sp : acc, HS.insert name visited)
                spRefs

seBottomUp ::
  forall a b.
  (StorePath (StorePath StoreName b) a -> b) ->
  StoreEnv a ->
  StoreEnv b
seBottomUp f StoreEnv {sePaths, seRoots} =
  StoreEnv
    { sePaths = snd $ execState (mapM_ go seRoots) (sePaths, HM.empty),
      seRoots
    }
  where
    unsafeLookup k m =
      fromMaybe
        (error $ "invariant violation: name doesn't exists: " <> show k)
        (HM.lookup k m)
    go ::
      StoreName ->
      State
        ( HashMap StoreName (StorePath StoreName a),
          HashMap StoreName (StorePath StoreName b)
        )
        (StorePath StoreName b)
    go name = do
      processed <- gets snd
      case name `HM.lookup` processed of
        Just sp -> return sp
        Nothing -> do
          sp@StorePath {spName, spRefs} <- unsafeLookup name <$> gets fst
          refs <- mapM go spRefs
          let new = sp {spPayload = f sp {spRefs = refs}}
          modify $ bimap (HM.delete spName) (HM.insert spName new)
          return new

--------------------------------------------------------------------------------

storeEnvToDot :: StoreEnv a -> Text
storeEnvToDot env =
  seBottomUp go env
    & seGetRoots
    & toList
    & map spPayload
    & mconcat
    & render
  where
    go sp =
      fromList [Set.singleton (spName sp, spName ref) <> spPayload ref | ref <- spRefs sp]
        & mconcat

    render :: Set (StoreName, StoreName) -> Text
    render edges =
      Dot.DotGraph
        Dot.Strict
        Dot.Directed
        Nothing
        [ Dot.StatementEdge
            ( Dot.EdgeStatement
                (Dot.ListTwo (Dot.EdgeNode (mkNodeId from)) (Dot.EdgeNode (mkNodeId to)) [])
                []
            )
          | (from, to) <- toList edges
        ]
        & Dot.encode

    mkNodeId :: StoreName -> Dot.NodeId
    mkNodeId = fromString . toString . storeNameToShortText

--------------------------------------------------------------------------------

data NixPathInfo = NixPathInfo
  { npiPath :: FilePath,
    npiNarSize :: Int,
    npiReferences :: [FilePath],
    npiSignatures :: [NixPathSignature]
  }

data NixPathSignature = NixPathSignature
  { npsKeyName :: Text,
    npsSignature :: Text
  }
  deriving (Show, Eq, Ord, NFData, Generic)

instance FromJSON NixPathSignature where
  parseJSON (String t) =
    case T.splitOn ":" t of
      [key, sig]
        | not (T.null key) && not (T.null sig) ->
            return $ NixPathSignature key sig
      _ -> fail "Expecting a string in the form of 'key:sig'."
  parseJSON _ = fail "Expecting a string."

data NixPathInfoResult
  = NixPathInfoValid NixPathInfo
  | NixPathInfoInvalid FilePath

parse2_18 :: Value -> Parser NixPathInfoResult
parse2_18 (Object obj) =
  ( NixPathInfoValid
      <$> ( NixPathInfo
              <$> obj .: "path"
              <*> obj .: "narSize"
              <*> obj .: "references"
              <*> (obj .:? "signatures" .!= [])
          )
  )
    <|> ( do
            path <- obj .: "path"
            valid <- obj .: "valid"
            guard (not valid)
            return $ NixPathInfoInvalid path
        )
parse2_18 _ = fail "Expecting an object."

parse2_19 :: (FilePath, Value) -> Parser NixPathInfoResult
parse2_19 (path, Object obj) =
  NixPathInfoValid
    <$> ( NixPathInfo
            path
            <$> obj .: "narSize"
            <*> obj .: "references"
            <*> (obj .:? "signatures" .!= [])
        )
parse2_19 (path, Null) = return $ NixPathInfoInvalid path
parse2_19 (_, _) = fail "Expecting an object or null"

newtype NixPathOutput = NixPathOutput
  { npoResults :: [NixPathInfoResult]
  }

instance FromJSON NixPathOutput where
  parseJSON (Array a) = NixPathOutput <$> mapM parse2_18 (toList a)
  parseJSON (Object o) = NixPathOutput <$> mapM (parse2_19 . first K.toString) (KM.toList o)
  parseJSON _ = fail "Expecting an array (nix<=2.18) or an object with mapping from path to info (nix>=2.19)."