{-|
Mostly effectful functions to assemble artifacts.
-}
module B9.ArtifactGeneratorImpl where

import B9.ArtifactGenerator
import B9.B9Monad
import B9.B9Config
import B9.VmBuilder
import B9.Vm
import B9.DiskImageBuilder
import B9.ConfigUtils hiding (tell)
import B9.Content.StringTemplate
import B9.Content.Generator
import B9.Content.AST

import qualified Data.ByteString as B
import qualified Data.Text as T
import qualified Data.Text.Encoding as E
import Data.Data
import Data.Generics.Schemes
import Data.Generics.Aliases
import Data.List
import Data.Function
import Control.Arrow
import Control.Monad.IO.Class
import Control.Monad.Reader
import Control.Monad.Writer
import Control.Monad.Except
import System.FilePath
import System.Directory
import Text.Printf
import Text.Show.Pretty (ppShow)

-- | Return a list of relative paths for the /local/ files to be generated
-- by the ArtifactGenerator. This excludes 'Shared' and Transient image targets.
getArtifactOutputFiles :: ArtifactGenerator -> Either String [FilePath]
getArtifactOutputFiles g =
  concatMap getOutputs <$> evalArtifactGenerator undefined undefined [] g
  where
    getOutputs  (IG _ sgs a) =
      let
          toOutFile (AssemblyGeneratesOutputFiles fs) = fs
          toOutFile (AssemblyCopiesSourcesToDirectory pd) =
            let sourceFiles = sourceGeneratorOutputFile <$> sgs
            in (pd </>) <$> sourceFiles

      in getAssemblyOutput  a >>= toOutFile

-- | Run an artifact generator to produce the artifacts.
assemble :: ArtifactGenerator -> B9 [AssembledArtifact]
assemble artGen = do
  b9cfgEnvVars <- envVars <$> getConfig
  buildId      <- getBuildId
  buildDate    <- getBuildDate
  case evalArtifactGenerator buildId buildDate b9cfgEnvVars artGen of
    Left  err -> error err
    Right is  -> createAssembledArtifacts is

-- | Evaluate an 'ArtifactGenerator' into a list of low-level build instructions
-- that can be built with 'createAssembledArtifacts'.
evalArtifactGenerator
  :: String
  -> String
  -> BuildVariables
  -> ArtifactGenerator
  -> Either String [InstanceGenerator [SourceGenerator]]
evalArtifactGenerator buildId buildDate b9cfgEnvVars artGen =
  let
    ag = parseArtifactGenerator artGen
    e  = CGEnv
      ((buildDateKey, buildDate) : (buildIdKey, buildId) : b9cfgEnvVars)
      []
  in
    case execCGParser ag e of
      Left (CGError err) ->
        Left (printf "error parsing: %s: %s" (ppShow artGen) err)
      Right igs -> case execIGEnv `mapM` igs of
        Left err ->
          Left (printf "Failed to parse:\n%s\nError: %s" (ppShow artGen) err)
        Right is -> Right is

-- | Parse an artifacto generator inside a 'CGParser' monad.
parseArtifactGenerator :: ArtifactGenerator -> CGParser ()
parseArtifactGenerator g = case g of
  Sources srcs gs -> withArtifactSources srcs (mapM_ parseArtifactGenerator gs)
  Let bs gs -> withBindings bs (mapM_ parseArtifactGenerator gs)
  LetX bs gs -> withXBindings bs (mapM_ parseArtifactGenerator gs)
  EachT keySet valueSets gs -> do
    allBindings <- eachBindingSetT g keySet valueSets
    sequence_
      (flip withBindings (mapM_ parseArtifactGenerator gs) <$> allBindings)
  Each kvs gs -> do
    allBindings <- eachBindingSet g kvs
    sequence_ $ do
      b <- allBindings
      return (withBindings b (mapM_ parseArtifactGenerator gs))
  Artifact iid assembly -> writeInstanceGenerator iid assembly
  EmptyArtifact         -> return ()

-- | Execute a 'CGParser' action in an environment that contains a list of
-- 'ArtifactSource's.
withArtifactSources :: [ArtifactSource] -> CGParser () -> CGParser ()
withArtifactSources srcs =
  local (\ce -> ce { agSources = agSources ce ++ srcs })

withBindings :: [(String, String)] -> CGParser () -> CGParser ()
withBindings bs = local (addBindings bs)

addBindings :: [(String, String)] -> CGEnv -> CGEnv
addBindings bs ce =
  let addBinding env (k, v) = nubBy ((==) `on` fst) ((k, subst env v) : env)
      newEnv = foldl addBinding (agEnv ce) bs
  in  ce { agEnv = newEnv }

withXBindings :: [(String, [String])] -> CGParser () -> CGParser ()
withXBindings bs cp = (`local`cp) `mapM_` (addBindings <$> allXBindings bs)
 where
  allXBindings ((k, vs):rest) =
    [ (k, v) : c | v <- vs, c <- allXBindings rest ]
  allXBindings [] = [[]]

eachBindingSetT
  :: ArtifactGenerator
  -> [String]
  -> [[String]]
  -> CGParser [[(String, String)]]
eachBindingSetT g vars valueSets = if all ((==length vars) . length) valueSets
  then return (zip vars <$> valueSets)
  else cgError
    ( printf
      "Error in 'Each' binding during artifact generation in:\n '%s'.\n\nThe variable list\n%s\n has %i entries, but this binding set\n%s\n\nhas a different number of entries!\n"
      (ppShow g)
      (ppShow vars)
      (length vars)
      (ppShow (head (dropWhile ((==length vars) . length) valueSets)))
    )

eachBindingSet
  :: ArtifactGenerator -> [(String, [String])] -> CGParser [[(String, String)]]
eachBindingSet g kvs = do
  checkInput
  return bindingSets
 where
  bindingSets = transpose [ repeat k `zip` vs | (k, vs) <- kvs ]
  checkInput  = when
    (1 /= length (nub $ length . snd <$> kvs))
    ( cgError
      ( printf
        "Error in 'Each' binding: \n%s\nAll value lists must have the same length!"
        (ppShow g)
      )
    )


writeInstanceGenerator :: InstanceId -> ArtifactAssembly -> CGParser ()
writeInstanceGenerator (IID iidStrT) assembly = do
  env@(CGEnv bindings _) <- ask
  iid <- either (throwError . CGError) (return . IID) (substE bindings iidStrT)
  let env'       = addBindings [(instanceIdKey, iidStr)] env
      IID iidStr = iid
  tell [IG iid env' assembly]

-- | Monad for creating Instance generators.
newtype CGParser a =
          CGParser
            { runCGParser :: WriterT [InstanceGenerator CGEnv] (ReaderT CGEnv (Either CGError)) a }
  deriving (Functor, Applicative, Monad, MonadReader CGEnv, MonadWriter [InstanceGenerator CGEnv], MonadError CGError)

data CGEnv = CGEnv { agEnv :: [(String, String)], agSources :: [ArtifactSource] }
  deriving (Read, Show, Eq)

data InstanceGenerator e = IG InstanceId e ArtifactAssembly
  deriving (Read, Show, Typeable, Data, Eq)

newtype CGError = CGError String
  deriving (Read, Show, Typeable, Data, Eq)

cgError :: String -> CGParser a
cgError msg = throwError (CGError msg)

execCGParser :: CGParser () -> CGEnv -> Either CGError [InstanceGenerator CGEnv]
execCGParser = runReaderT . execWriterT . runCGParser

execIGEnv
  :: InstanceGenerator CGEnv
  -> Either String (InstanceGenerator [SourceGenerator])
execIGEnv (IG iid (CGEnv env sources) assembly) =
  IG iid <$> sourceGens <*> pure (substAssembly env assembly)
  where sourceGens = join <$> mapM (toSourceGen env) sources

substAssembly :: [(String, String)] -> ArtifactAssembly -> ArtifactAssembly
substAssembly env = everywhere gsubst
 where
  gsubst :: Data a => a -> a
  gsubst =
    mkT substAssembly_ `extT` substImageTarget env `extT` substVmScript env

  substAssembly_ (CloudInit ts f) = CloudInit ts (sub f)
  substAssembly_ vm               = vm

  sub = subst env

toSourceGen
  :: [(String, String)] -> ArtifactSource -> Either String [SourceGenerator]
toSourceGen env src = case src of
  FromFile t (Source conv f) -> do
    t' <- substE env t
    f' <- substE env f
    return [SG env (SGFiles [Source conv f']) KeepPerm t']
  FromContent t c -> do
    t' <- substE env t
    return [SG env (SGContent c) KeepPerm t']
  Concatenation t src' -> do
    sgs <- join <$> mapM (toSourceGen env) src'
    t'  <- substE env t
    let froms = join (sgGetFroms <$> sgs)
    return [SG env (SGFiles froms) KeepPerm t']
  SetPermissions o g a src' -> do
    sgs <- join <$> mapM (toSourceGen env) src'
    mapM (setSGPerm o g a) sgs
  FromDirectory fromDir src' -> do
    sgs      <- join <$> mapM (toSourceGen env) src'
    fromDir' <- substE env fromDir
    return (setSGFromDirectory fromDir' <$> sgs)
  IntoDirectory toDir src' -> do
    sgs    <- join <$> mapM (toSourceGen env) src'
    toDir' <- substE env toDir
    return (setSGToDirectory toDir' <$> sgs)

createAssembledArtifacts
  :: [InstanceGenerator [SourceGenerator]] -> B9 [AssembledArtifact]
createAssembledArtifacts igs = do
  buildDir <- getBuildDir
  let outDir = buildDir </> "artifact-instances"
  ensureDir (outDir ++ "/")
  generated <- generateSources outDir `mapM` igs
  createTargets `mapM` generated

generateSources
  :: FilePath
  -> InstanceGenerator [SourceGenerator]
  -> B9 (InstanceGenerator FilePath)
generateSources outDir (IG iid sgs assembly) = do
  uiid@(IID uiidStr) <- generateUniqueIID iid
  dbgL (printf "generating sources for %s" uiidStr)
  let instanceDir = outDir </> uiidStr
  traceL (printf "generating sources for %s:\n%s\n" uiidStr (ppShow sgs))
  generateSourceTo instanceDir `mapM_` sgs
  return (IG uiid instanceDir assembly)

createTargets :: InstanceGenerator FilePath -> B9 AssembledArtifact
createTargets (IG uiid@(IID uiidStr) instanceDir assembly) = do
  targets <- createTarget uiid instanceDir assembly
  dbgL (printf "assembled artifact %s" uiidStr)
  return (AssembledArtifact uiid targets)

generateUniqueIID :: InstanceId -> B9 InstanceId
generateUniqueIID (IID iid) = do
  buildId <- getBuildId
  return (IID (printf "%s-%s" iid buildId))

generateSourceTo :: FilePath -> SourceGenerator -> B9 ()
generateSourceTo instanceDir (SG env sgSource p to) = do
  let toAbs = instanceDir </> to
  ensureDir toAbs
  result <- case sgSource of
    SGFiles froms -> do
      sources <- mapM (sgReadSourceFile env) froms
      return (mconcat sources)
    SGContent c -> withEnvironment env (render c)
  traceL (printf "rendered: \n%s\n" (T.unpack (E.decodeUtf8 result)))
  liftIO (B.writeFile toAbs result)
  sgChangePerm toAbs p


sgReadSourceFile :: [(String, String)] -> SourceFile -> B9 B.ByteString
sgReadSourceFile env = withEnvironment env . readTemplateFile

sgChangePerm :: FilePath -> SGPerm -> B9 ()
sgChangePerm _ KeepPerm = return ()
sgChangePerm f (SGSetPerm (o, g, a)) =
  cmd (printf "chmod 0%i%i%i '%s'" o g a f)

-- | Internal data type simplifying the rather complex source generation by
--   bioling down 'ArtifactSource's to a flat list of uniform 'SourceGenerator's.
data SourceGenerator = SG [(String, String)] SGSource SGPerm FilePath
  deriving (Read, Show, Eq)

-- | Return the (internal-)output file of the source file that is generated.
sourceGeneratorOutputFile :: SourceGenerator -> FilePath
sourceGeneratorOutputFile (SG _ _ _ f) = f

data SGSource = SGFiles [SourceFile]
              | SGContent Content
  deriving (Read, Show, Eq)

data SGPerm = SGSetPerm (Int, Int, Int)
            | KeepPerm
  deriving (Read, Show, Typeable, Data, Eq)


sgGetFroms :: SourceGenerator -> [SourceFile]
sgGetFroms (SG _ (SGFiles fs) _ _) = fs
sgGetFroms _                       = []

setSGPerm
  :: Int -> Int -> Int -> SourceGenerator -> Either String SourceGenerator
setSGPerm o g a (SG env from KeepPerm dest) =
  Right (SG env from (SGSetPerm (o, g, a)) dest)
setSGPerm o g a sg
  | o < 0 || o > 7 = Left
    (printf "Bad 'owner' permission %i in \n%s" o (ppShow sg))
  | g < 0 || g > 7 = Left
    (printf "Bad 'group' permission %i in \n%s" g (ppShow sg))
  | a < 0 || a > 7 = Left
    (printf "Bad 'all' permission %i in \n%s" a (ppShow sg))
  | otherwise = Left
    (printf "Permission for source already defined:\n %s" (ppShow sg))

setSGFromDirectory :: FilePath -> SourceGenerator -> SourceGenerator
setSGFromDirectory fromDir (SG e (SGFiles fs) p d) = SG
  e
  (SGFiles (setSGFrom <$> fs))
  p
  d
  where setSGFrom (Source t f) = Source t (fromDir </> f)
setSGFromDirectory _fromDir sg = sg

setSGToDirectory :: FilePath -> SourceGenerator -> SourceGenerator
setSGToDirectory toDir (SG e fs p d) = SG e fs p (toDir </> d)

-- | Create the actual target, either just a mountpoint, or an ISO or VFAT
-- image.
createTarget
  :: InstanceId -> FilePath -> ArtifactAssembly -> B9 [ArtifactTarget]
createTarget iid instanceDir (VmImages imageTargets vmScript) = do
  dbgL (printf "Creating VM-Images in '%s'" instanceDir)
  success <- buildWithVm iid imageTargets instanceDir vmScript
  let err_msg      = printf "Error creating 'VmImages' for instance '%s'" iidStr
      (IID iidStr) = iid
  unless success (errorL err_msg >> error err_msg)
  return [VmImagesTarget]
createTarget _ instanceDir (CloudInit types outPath) = mapM create_ types
 where
  create_ CI_DIR = do
    let ciDir = outPath
    ensureDir (ciDir ++ "/")
    dbgL (printf "creating directory '%s'" ciDir)
    files <- getDirectoryFiles instanceDir
    traceL (printf "copying files: %s" (show files))
    liftIO
      ( mapM_
        ( \(f, t) -> do
          ensureDir t
          copyFile f t
        )
        (((instanceDir</>) &&& (ciDir</>)) <$> files)
      )
    infoL (printf "CREATED CI_DIR: '%s'" (takeFileName ciDir))
    return (CloudInitTarget CI_DIR ciDir)

  create_ CI_ISO = do
    buildDir <- getBuildDir
    let isoFile = outPath <.> "iso"
        tmpFile = buildDir </> takeFileName isoFile
    ensureDir tmpFile
    dbgL
      ( printf
        "creating cloud init iso temp image '%s', destination file: '%s"
        tmpFile
        isoFile
      )
    cmd
      ( printf "genisoimage -output '%s' -volid cidata -rock -d '%s' 2>&1"
               tmpFile
               instanceDir
      )
    dbgL (printf "moving iso image '%s' to '%s'" tmpFile isoFile)
    ensureDir isoFile
    liftIO (copyFile tmpFile isoFile)
    infoL (printf "CREATED CI_ISO IMAGE: '%s'" (takeFileName isoFile))
    return (CloudInitTarget CI_ISO isoFile)

  create_ CI_VFAT = do
    buildDir <- getBuildDir
    let vfatFile = outPath <.> "vfat"
        tmpFile  = buildDir </> takeFileName vfatFile
    ensureDir tmpFile
    files <- map (instanceDir</>) <$> getDirectoryFiles instanceDir
    dbgL (printf "creating cloud init vfat image '%s'" tmpFile)
    traceL (printf "adding '%s'" (show files))
    cmd (printf "truncate --size 2M '%s'" tmpFile)
    cmd (printf "mkfs.vfat -n cidata '%s' 2>&1" tmpFile)
    cmd
      (  unwords (printf "mcopy -oi '%s' " tmpFile : (printf "'%s'" <$> files))
      ++ " ::"
      )
    dbgL (printf "moving vfat image '%s' to '%s'" tmpFile vfatFile)
    ensureDir vfatFile
    liftIO (copyFile tmpFile vfatFile)
    infoL (printf "CREATED CI_VFAT IMAGE: '%s'" (takeFileName vfatFile))
    return (CloudInitTarget CI_ISO vfatFile)