## DB Migrate
## ==========
##
## Simple tool to manage database migrations.

import algorithm, json, times, os, strutils, docopt, db_postgres, sets,
  sequtils, logging

type
  DbMigrateConfig* = tuple[ driver, sqlDir, connectionString: string, logLevel: Level ]

proc createMigrationsTable(conn: DbConn): void =
  conn.exec(sql("""
CREATE TABLE IF NOT EXISTS migrations (
  id SERIAL PRIMARY KEY,
  name VARCHAR NOT NULL,
  run_at TIMESTAMP NOT NULL DEFAULT NOW())"""))

proc rollbackWithErr (pgConn: DbConn, errMsg: string): void =
  pgConn.exec(sql"ROLLBACK")
  dbError(errMsg)

proc loadConfig*(filename: string): DbMigrateConfig =
  ## Load DbMigrateConfig from a file.
  let cfg = json.parseFile(filename)

  var logLevel: Level
  if cfg.hasKey("logLevel"):
    let idx = find(LevelNames, cfg["logLevel"].getStr.toUpper)
    logLevel = if idx == -1: lvlInfo else: (Level)(idx)

  return (
    driver: if cfg.hasKey("driver"): cfg["driver"].getStr else: "postres",
    sqlDir: if cfg.hasKey("sqlDir"): cfg["sqlDir"].getStr else: "migrations",
    connectionString: cfg["connectionString"].getStr,
    logLevel: logLevel)

proc createMigration*(config: DbMigrateConfig, migrationName: string): seq[string] =
  ## Create a new set of database migration files.
  let timestamp = getTime().getLocalTime().format("yyyyMMddHHmmss")
  let filenamePrefix =  timestamp & "-" & migrationName

  let upFilename = joinPath(config.sqlDir, filenamePrefix & "-up.sql")
  let downFilename = joinPath(config.sqlDir, filenamePrefix & "-down.sql")

  let scriptDesc = migrationName & " (" & timestamp & ")"

  let upFile = open(upFilename, fmWrite)
  let downFile = open(downFilename, fmWrite)

  upFile.writeLine "-- UP script for " & scriptDesc
  downFile.writeLine "-- DOWN script for " & scriptDesc

  upFile.close()
  downFile.close()

  return @[upFilename, downFilename]

proc diffMigrations*(pgConn: DbConn, config: DbMigrateConfig):
  tuple[ run, notRun, missing: seq[string] ] =

  # Query the database to find out what migrations have been run.
  var migrationsRun = initSet[string]()

  for row in pgConn.fastRows(sql"SELECT * FROM migrations ORDER BY name", @[]):
    migrationsRun.incl(row[1])

  # Inspect the filesystem to see what migrations are available.
  var migrationsAvailable = initSet[string]()
  for filePath in walkFiles(joinPath(config.sqlDir, "*.sql")):
    var migrationName = filePath.extractFilename
    migrationName.removeSuffix("-up.sql")
    migrationName.removeSuffix("-down.sql")
    migrationsAvailable.incl(migrationName)

  # Diff with the list of migrations that we have in our migrations
  # directory.
  let migrationsInOrder =
    toSeq(migrationsAvailable.items).sorted(system.cmp)

  var migrationsNotRun = newSeq[string]()
  var missingMigrations = newSeq[string]()

  for migration in migrationsInOrder:
    if not migrationsRun.contains(migration):
      migrationsNotRun.add(migration)

    # if we've already seen some migrations that have not been run, but this
    # one has been, that means we have a gap and are missing migrations
    elif migrationsNotRun.len > 0:
      missingMigrations.add(migrationsNotRun)
      migrationsNotRun = newSeq[string]()

  return (run: toSeq(migrationsRun.items).sorted(system.cmp),
          notRun: migrationsNotRun,
          missing: missingMigrations)

proc readStatements*(filename: string): seq[SqlQuery] =
  let migrationSql = filename.readFile
  result = migrationSql.split(';').
    filter(proc(st: string): bool = st.strip.len > 0 and not st.startsWith("--")).
    map(proc(st: string): SqlQuery = sql(st & ";"))

proc up*(pgConn: DbConn, config: DbMigrateConfig, toRun: seq[string]): seq[string] =
  var migrationsRun = newSeq[string]()
  # Begin a transaction.
  pgConn.exec(sql"BEGIN")

  # Apply each of the migrations.
  for migration in toRun:
    info migration
    let filename = joinPath(config.sqlDir, migration & "-up.sql")

    if not filename.fileExists:
      pgConn.rollbackWithErr "Can not find UP file for " & migration &
        ". Expected '" & filename & "'."

    let statements = filename.readStatements

    try:
      for statement in statements: pgConn.exec(statement)
      pgConn.exec(sql"INSERT INTO migrations (name) VALUES (?);", migration)
    except DbError:
      pgConn.rollbackWithErr "Migration '" & migration & "' failed:\n\t" &
        getCurrentExceptionMsg()

    migrationsRun.add(migration)

  pgConn.exec(sql"COMMIT")

  return migrationsRun

proc down*(pgConn: DbConn, config: DbMigrateConfig, migrationsToDown: seq[string]): seq[string] =
  var migrationsDowned = newSeq[string]()

  pgConn.exec(sql"BEGIN")

  for migration in migrationsToDown:
    info migration

    let filename = joinPath(config.sqlDir, migration & "-down.sql")

    if not filename.fileExists:
      pgConn.rollbackWithErr "Can not find DOWN file for " & migration &
        ". Expected '" & filename & "'."

    let statements = filename.readStatements

    try:
      for statement in statements: pgConn.exec(statement)
      pgConn.exec(sql"DELETE FROM migrations WHERE name = ?;", migration)
    except DbError:
      pgConn.rollbackWithErr "Migration '" & migration & "' failed:\n\t" &
        getCurrentExceptionMsg()

    migrationsDowned.add(migration)

  pgConn.exec(sql"COMMIT")
  return migrationsDowned

when isMainModule:
  let doc = """
Usage:
  db_migrate [options] create <migration-name>
  db_migrate [options] up [<count>]
  db_migrate [options] down [<count>]
  db_migrate [options] init <schema-name>
  db_migrate (-V | --version)

Options:

  -c --config <config-file>   Use the given configuration file (defaults to
                              "database.json").

  -q --quiet                  Suppress log information.

  -v --verbose                Print detailed log information.

     --very-verbose           Print very detailed log information.

  -V --version                Print the tools version information.
"""

  # Parse arguments
  let args = docopt(doc, version = "db-migrate 0.2.1")

  let exitErr = proc(msg: string): void =
    fatal("db_migrate: " & msg)
    quit(QuitFailure)

  # Load configuration file
  let configFilename =
    if args["--config"]: $args["<config-file>"]
    else: "database.json"

  var config: DbMigrateConfig
  try:
    config = loadConfig(configFilename)
  except IOError:
    exitErr "Cannot open config file: " & configFilename
  except:
    exitErr "Error parsing config file: " & configFilename & "\L\t" & getCurrentExceptionMsg()


  logging.addHandler(newConsoleLogger())
  if args["--quiet"]: logging.setLogFilter(lvlError)
  elif args["--very-verbose"]: logging.setLogFilter(lvlAll)
  elif args["--verbose"]: logging.setlogFilter(lvlDebug)
  else: logging.setLogFilter(config.logLevel)

  # Check for migrations directory
  if not existsDir config.sqlDir:
    try:
      warn "SQL directory '" & config.sqlDir &
        "' does not exist and will be created."
      createDir config.sqlDir
    except IOError:
      exitErr "Unable to create directory: " & config.sqlDir & ":\L\T" & getCurrentExceptionMsg()

  # Execute commands
  if args["create"]:
    try:
      let filesCreated = createMigration(config, $args["<migration-name>"])
      info "Created new migration files:"
      for filename in filesCreated: info "\t" & filename
    except IOError:
      exitErr "Unable to create migration scripts: " & getCurrentExceptionMsg()

  else:
    let pgConn: DbConn =
      try: open("", "", "", config.connectionString)
      except DbError:
        exitErr "Unable to connect to the database: " & getCurrentExceptionMsg()
        (DbConn)(nil)

    pgConn.createMigrationsTable

    let (run, notRun, missing) = diffMigrations(pgConn, config)

    # Make sure we have no gaps (database is in an unknown state)
    if missing.len > 0:
      exitErr "Database is in an inconsistent state. Migrations have been " &
        "run that are not sequential."

    if args["up"]:
      try:
        let count = if args["<count>"]: parseInt($args["<count>"]) else: high(int)
        let toRun = if count < notRun.len: notRun[0..<count] else: notRun
        let migrationsRun = pgConn.up(config, toRun)
        if count < high(int): info "Went up " & $(migrationsRun.len) & "."
      except DbError:
        exitErr "Unable to migrate database: " & getCurrentExceptionMsg()

    elif args["down"]:
      try:
        let count = if args["<count>"]: parseInt($args["<count>"]) else: 1
        let toRun = if count < run.len: run.reversed[0..<count] else: run.reversed
        let migrationsRun = pgConn.down(config, toRun)
        info "Went down " & $(migrationsRun.len) & "."
      except DbError:
        exitErr "Unable to migrate database: " & getCurrentExceptionMsg()

    elif args["init"]: discard

    let newResults = diffMigrations(pgConn, config)
    if newResults.notRun.len > 0:
      info "Database is behind by " & $(newResults.notRun.len) & " migrations."
    else: info "Database is up to date."