content: {-# LANGUAGE CPP, OverloadedStrings #-}
  
  -- | Functions for creating and running database migrations. You should
  -- probably be using the `pg_migrate` executable to run migrations, however
  -- these functions are exposed for developers that want to integrate migrations
  -- more tightly into their applications or utilities.
  
  module Database.PostgreSQL.Migrate
    ( initializeDb
    , runMigrationsForDir
    , runRollbackForDir
    , compileMigrationsForDir
    , dumpDb
    , newMigration
    , defaultMigrationsDir
    , getDirectoryMigrations
    , MigrationDetails(..)
    ) where
  
  import Control.Exception (bracket)
  import Control.Monad
  import Data.List
  import Data.Time
  import Database.PostgreSQL.Simple hiding (connect)
  import qualified Data.ByteString.Char8 as S8
  import Database.PostgreSQL.Migrations
  import Language.Haskell.Exts
    (parseFile, fromParseResult,
     Module(..), ModuleHead(..), ModuleName(..),
     Decl(..), Pat(..), Name(..))
  import Language.Haskell.Exts.Pretty (prettyPrint)
  import System.Exit
  import GHC.IO.Handle
  import System.Process
  import System.Directory
  import System.FilePath
  import System.Environment
  import System.IO
  import System.IO.Temp (withTempDirectory)
  #if !MIN_VERSION_time(1,5,0)
  import System.Locale
  #endif
  
  import Paths_postgresql_orm
  
  -- | The default relative path containing migrations: @\"db\/migrations\"@
  defaultMigrationsDir :: FilePath
  defaultMigrationsDir = "db" </> "migrations"
  
  -- | Dumps the database schema to the given file handle.
  --
  -- This is a wrapper around the utility /pg_dump/ that comes with postgresql.
  -- Therefore, /pg_dump/ must be installed on the system.
  dumpDb :: Handle -> IO ExitCode
  dumpDb outputFile = do
    let opts = ["--schema-only", "-O", "-x"]
    e <- getEnvironment
    let args = case lookup "DATABASE_URL" e of
                 Just dburl -> dburl:opts
                 Nothing -> opts
    (_, out, err, ph) <- runInteractiveProcess "pg_dump" args Nothing Nothing
    exitCode <- waitForProcess ph
    if exitCode /= ExitSuccess then do
      S8.hGetContents err >>= S8.hPut stderr
      else do
        raw <- S8.hGetContents out
        let clean = S8.concat $ intersperse "\n" $
                      filter ((/= "--") . (S8.take 2)) $
                      S8.lines raw
        S8.hPut outputFile clean
    return exitCode
  
  -- | Initializes the database by creating a \"schema-migrations\" table.
  -- This table must exist before running any migrations.
  initializeDb :: IO ()
  initializeDb = bracket connectEnv close $ \conn ->
    void $ execute_ conn
      "create table if not exists schema_migrations (version VARCHAR(28))"
  
  -- | Runs all new migrations in a given directory and dumps the
  -- resulting schema to a file \"schema.sql\" in the migrations
  -- directory.
  --
  -- Determining which migrations to run is done by querying the database for the
  -- largest version in the /schema_migrations/ table, and choosing all
  -- migrations in the given directory with higher versions.
  runMigrationsForDir :: Handle -- ^ Log output (probably stdout)
                      -> FilePath -- ^ Path to directory containing migrations
                      -> IO ExitCode
  runMigrationsForDir logOut dir = bracket connectEnv close $ \conn -> do
    res <- query_ conn
            "select version from schema_migrations order by version desc limit 1"
    let latestVersion = case res of
                          [] -> ""
                          (Only latest):_ -> latest
    migrations <- getDirectoryMigrations dir >>=
                      return . (dropWhile (isVersion (<= latestVersion)))
    go migrations
    where go [] = withFile (dir </> ".." </> "schema.sql") WriteMode dumpDb
          go (mig@(MigrationDetails _ _ name):fs) = do
                hPutStrLn logOut $ "=== Running Migration " ++ name
                exitCode <- runMigration mig
                if exitCode == ExitSuccess then do
                  hPutStrLn logOut "=== Success"
                  go fs
                  else do
                    hPutStrLn logOut "=== Migration Failed!"
                    return exitCode
  
  -- | Run a migration. The returned exit code denotes the success or failure of
  -- the migration.
  runMigration :: MigrationDetails -> IO ExitCode
  runMigration (MigrationDetails file vers _) = do
    rawSystem "runghc"
      [file, "up", vers, "--with-db-commit"]
  
  runRollbackForDir :: FilePath -> IO ExitCode
  runRollbackForDir dir = bracket connectEnv close $ \conn -> do
    res <- query_ conn
            "select version from schema_migrations order by version desc limit 1"
    case res of
      [] -> do
        putStrLn "=== DB Fully Rolled Back!"
        return ExitSuccess
      (Only latest):_ -> do
        (Just (mig@(MigrationDetails _ _ name))) <-
                    getDirectoryMigrations dir >>=
                      return . (find (isVersion (== latest)))
        putStrLn $ "=== Running Rollback " ++ name
        exitCode <- runRollback mig
        if exitCode == ExitSuccess then do
          putStrLn "=== Success"
          withFile (dir </> ".." </> "schema.sql") WriteMode dumpDb
          else do
            putStrLn "=== Migration Failed!"
            return exitCode
  
  -- | Run a migration. The returned exit code denotes the success or failure of
  -- the migration.
  runRollback :: MigrationDetails -> IO ExitCode
  runRollback (MigrationDetails file vers _) = do
    rawSystem "runghc"
      [file, "down", vers, "--with-db-commit"]
  
  data MigrationDetails = MigrationDetails { migrationPath :: FilePath
                                           , migrationVersion :: String
                                           , migrationName :: String }
                                           deriving (Show)
  
  getDirectoryMigrations :: FilePath -> IO [MigrationDetails]
  getDirectoryMigrations dir = do
    files0 <- getDirectoryContents dir
    let files = filter (('.' /=) . head) $ sort files0
    return $ map (splitFileVersionName dir) files
  
  splitFileVersionName :: FilePath -> FilePath -> MigrationDetails
  splitFileVersionName dir file =
    let fileName = takeBaseName file
        parts    = foldr (\chr (hd:result) ->
                            if chr == '_' then
                              "":hd:result
                              else ((chr:hd):result))
                         [""] fileName
        vers  = head parts
        name     = concat $ intersperse "_" $ tail parts
    in MigrationDetails (dir </> file) vers name
  
  isVersion :: (String -> Bool) -> MigrationDetails -> Bool
  isVersion cond (MigrationDetails _ v _) = cond v
  
  newMigration :: FilePath -> FilePath -> IO ()
  newMigration baseName dir = do
    now <- getZonedTime
    let filePath = (formatTime defaultTimeLocale "%Y%m%d%H%M%S" now) ++
                   "_" ++ baseName ++ ".hs"
    origFile <- getDataFileName "static/migration.hs"
    copyFile origFile (dir </> filePath)
  
  
  compileMigrationsForDir :: FilePath -> FilePath -> IO ExitCode
  compileMigrationsForDir dir exe = do
    migrations <- getDirectoryMigrations dir
  
    -- Make temporary working directory.
    withTempDirectory "." "migration-compile-" $ \tmpdir -> do
  --  ($ "migration-compile") $ \tmpdir -> do
  
      -- Iterate over migrations in input directory, copying migration
      -- file to working directory, modifying as:
      --   - Add "module Migration<datestamp> (up, down) where" line at
      --     beginning of file, after any LANGUAGE options.
      --   - Remove any main function.
      --   - Collect (label, MigrationYYMMDD) pairs.
      moduleNames <- mapM (fixModule tmpdir) migrations
  
      -- Copy main program text to working directory, including
      -- migration list.
      makeMain tmpdir $ zip moduleNames migrations
  
      -- Copy compiler utilities.
      utils <- getDataFileName "static/CompilerUtils.hs"
      copyFile utils (tmpdir </> "CompilerUtils.hs")
  
      -- Compile migrater, writing executable outside temporary
      -- directory.
      cwd <- getCurrentDirectory
      system $ "cd " ++ tmpdir ++ "; ghc -o " ++ (cwd </> exe) ++ " Main.hs"
  
  -- | Turn migration scripts into valid modules.
  fixModule :: FilePath -> MigrationDetails -> IO String
  fixModule tmpdir (MigrationDetails path ver _) = do
    let modulename = "Migration" ++ ver
    modin <- parseFile path
    let modout = removeMain $ addModuleHeader modulename $ fromParseResult modin
    writeFile (tmpdir </> modulename <.> "hs") $ prettyPrint modout
    return modulename
  
  -- | Add "MigrationYYYYMMDDHHMMSS" module header.
  addModuleHeader :: String -> Module l -> Module l
  addModuleHeader name (Module sinfo _ pragmas imports decls) =
    Module sinfo (Just header) pragmas imports decls
    where header = ModuleHead sinfo (ModuleName sinfo name) Nothing Nothing
  addModuleHeader _ _ = error "Something wrong..."
  
  -- | Remove main function.
  removeMain :: Module l -> Module l
  removeMain (Module sinfo header pragmas imports decls) =
    Module sinfo header pragmas imports $ filter (not . isMain) decls
    where isMain :: Decl l -> Bool
          isMain (TypeSig _ [Ident _ "main"] _) = True
          isMain (PatBind _ (PVar _ (Ident _ "main")) _ _) = True
          isMain _ = False
  removeMain _ = error "Something wrong..."
  
  -- | Write main program.
  makeMain :: FilePath -> [(String, MigrationDetails)] -> IO ()
  makeMain tmpdir migrations =
    withFile (tmpdir </> "Main.hs") WriteMode $ \h -> do
      hPutStrLn h "module Main where\n"
      hPutStrLn h "import CompilerUtils\n"
      forM_ migrations $ \(modname, _) -> hPutStrLn h $ "import qualified " ++ modname
      hPutStrLn h "\nmigrations :: MigrationMap"
      hPutStrLn h "migrations = fromList"
      let migs = zipWith (++) ("  [ " : repeat "  , ") $ map doone migrations
      hPutStrLn h $ unlines migs
      hPutStrLn h "  ]\n"
      hPutStrLn h "main :: IO ()"
      hPutStrLn h "main = compiledMain migrations"
      where doone (modname, MigrationDetails _ ver name) =
              "(\"" ++ ver ++ "\",\n" ++
              "     Migration \"" ++ ver ++ "_" ++ name ++ "\" " ++
              modname ++ ".up " ++ modname ++ ".down)"
  