content: {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TypeFamilies #-}
  
  -- | This helper module is intended for use by the backend creators
  module Database.Groundhog.Generic.Migration
    ( Column (..),
      Reference (..),
      QualifiedName,
      TableInfo (..),
      UniqueDefInfo,
      AlterColumn (..),
      AlterTable (..),
      AlterDB (..),
      MigrationPack (..),
      SchemaAnalyzer (..),
      migrateRecursively,
      migrateSchema,
      migrateEntity,
      migrateList,
      getAlters,
      defaultMigConstr,
      showReferenceAction,
      readReferenceAction,
    )
  where
  
  import Control.Arrow ((&&&), (***))
  import Control.Monad (when)
  import Control.Monad.Trans.Class (lift)
  import Control.Monad.Trans.State (gets, modify)
  import Data.Either (lefts)
  import Data.Function (on)
  import Data.List (group, intercalate)
  import qualified Data.Map as Map
  import Data.Maybe (fromMaybe, isNothing, mapMaybe)
  import Database.Groundhog.Core
  import Database.Groundhog.Generic
  import Database.Groundhog.Generic.Sql (flatten, mainTableName, tableName)
  
  data Column = Column
    { colName :: String,
      colNull :: Bool,
      colType :: DbTypePrimitive,
      colDefault :: Maybe String
    }
    deriving (Eq, Show)
  
  data Reference = Reference
    { referencedTableName :: QualifiedName,
      -- | child column, parent column
      referencedColumns :: [(String, String)],
      referenceOnDelete :: Maybe ReferenceActionType,
      referenceOnUpdate :: Maybe ReferenceActionType
    }
    deriving (Show)
  
  -- | Schema-qualified name of a table of another database object
  type QualifiedName = (Maybe String, String)
  
  -- | Either column name or expression
  type UniqueDefInfo = UniqueDef' String (Either String String)
  
  data TableInfo = TableInfo
    { tableColumns :: [Column],
      tableUniques :: [UniqueDefInfo],
      -- | constraint name and reference
      tableReferences :: [(Maybe String, Reference)]
    }
    deriving (Show)
  
  data AlterColumn
    = Type DbTypePrimitive
    | IsNull
    | NotNull
    | Default String
    | NoDefault
    | UpdateValue String
    deriving (Show)
  
  data AlterTable
    = AddUnique UniqueDefInfo
    | DropConstraint String
    | DropIndex String
    | AddReference Reference
    | DropReference String
    | DropColumn String
    | AddColumn Column
    | AlterColumn Column [AlterColumn]
    deriving (Show)
  
  data AlterDB
    = AddTable String
    | -- | Qualified table name, create statement, structure of table from DB, structure of table from datatype, alters
      AlterTable QualifiedName String TableInfo TableInfo [AlterTable]
    | -- | Qualified trigger name, qualified table name
      DropTrigger QualifiedName QualifiedName
    | -- | Qualified trigger name, qualified table name, body
      AddTriggerOnDelete QualifiedName QualifiedName String
    | -- | Qualified trigger name, qualified table name, field name, body
      AddTriggerOnUpdate QualifiedName QualifiedName (Maybe String) String
    | -- | Statement which creates the function
      CreateOrReplaceFunction String
    | -- | Qualified function name
      DropFunction QualifiedName
    | -- | Schema name, if not exists
      CreateSchema String Bool
    deriving (Show)
  
  data MigrationPack conn = MigrationPack
    { compareTypes :: DbTypePrimitive -> DbTypePrimitive -> Bool,
      compareRefs :: (Maybe String, Reference) -> (Maybe String, Reference) -> Bool,
      compareUniqs :: UniqueDefInfo -> UniqueDefInfo -> Bool,
      compareDefaults :: String -> String -> Bool,
      migTriggerOnDelete :: QualifiedName -> [(String, String)] -> Action conn (Bool, [AlterDB]),
      migTriggerOnUpdate :: QualifiedName -> [(String, String)] -> Action conn [(Bool, [AlterDB])],
      migConstr :: EntityDef -> ConstructorDef -> Action conn (Bool, SingleMigration),
      escape :: String -> String,
      autoincrementedKeyTypeName :: String,
      mainTableId :: String,
      defaultPriority :: Int,
      -- | Sql pieces for the create table statement that add constraints and alterations for running after the table is created
      addUniquesReferences :: [UniqueDefInfo] -> [Reference] -> ([String], [AlterTable]),
      showSqlType :: DbTypePrimitive -> String,
      showColumn :: Column -> String,
      showAlterDb :: AlterDB -> SingleMigration,
      defaultReferenceOnDelete :: ReferenceActionType,
      defaultReferenceOnUpdate :: ReferenceActionType
    }
  
  mkColumns :: DbTypePrimitive -> (String, DbType) -> ([Column] -> [Column])
  mkColumns listAutoKeyType = go ""
    where
      go prefix (fname, typ) acc =
        case typ of
          DbEmbedded (EmbeddedDef flag ts) _ -> foldr (go prefix') acc ts where prefix' = if flag then "" else prefix ++ fname ++ [delim]
          DbList _ _ -> Column fullName False listAutoKeyType Nothing : acc
          DbTypePrimitive t nullable def _ -> Column fullName nullable t def : acc
        where
          fullName = prefix ++ fname
  
  mkReferences :: DbTypePrimitive -> (String, DbType) -> [Reference]
  mkReferences autoKeyType field = concat $ traverseDbType f field
    where
      f (DbEmbedded _ ref) cols = maybe [] (pure . mkReference cols) ref
      f (DbList lName _) cols = [mkReference cols (Right ((Nothing, lName), ["id"]), Nothing, Nothing)]
      f (DbTypePrimitive _ _ _ ref) cols = maybe [] (pure . mkReference cols) ref
      mkReference :: [String] -> ParentTableReference -> Reference
      mkReference cols (parent, onDel, onUpd) = case parent of
        Left (e, Nothing) -> Reference (entitySchema e, entityName e) (zipWith' (,) cols [keyName]) onDel onUpd
          where
            keyName = case constructors e of
              [cDef] -> fromMaybe (error "mkReferences: autokey name is Nothing") $ constrAutoKeyName cDef
              _ -> "id"
        Left (e, Just uName) -> Reference (entitySchema e, entityName e) (zipWith' (,) cols (map colName parentColumns)) onDel onUpd
          where
            cDef = case constructors e of
              [cDef'] -> cDef'
              _ -> error "mkReferences: datatype with unique key cannot have more than one constructor"
            uDef = findOne "unique" uniqueDefName (Just uName) $ constrUniques cDef
            fields = map (\(fName, _) -> findOne "field" fst fName $ constrParams cDef) $ getUniqueFields uDef
            parentColumns = foldr (mkColumns autoKeyType) [] fields
        Right (parentTable, parentColumns) -> Reference parentTable (zipWith' (,) cols parentColumns) onDel onUpd
      zipWith' _ [] [] = []
      zipWith' g (x : xs) (y : ys) = g x y : zipWith' g xs ys
      zipWith' _ _ _ = error "mkReferences: the lists have different length"
  
  traverseDbType :: (DbType -> [String] -> a) -> (String, DbType) -> [a]
  traverseDbType f field = snd $ go "" field
    where
      go prefix (fname, typ) =
        case typ of
          t@(DbEmbedded (EmbeddedDef flag ts) _) -> (cols, f t cols : xs)
            where
              prefix' = if flag then "" else prefix ++ fname ++ [delim]
              (cols, xs) = concatMap' (go prefix') ts
          t@DbList {} -> ([name], [f t [name]])
          t@DbTypePrimitive {} -> ([name], [f t [name]])
        where
          name = prefix ++ fname
          concatMap' g xs = concat *** concat $ unzip $ map g xs
  
  -- | Create migration for a given entity and all entities it depends on.
  -- The stateful Map is used to avoid duplicate migrations when an entity type
  -- occurs several times in a datatype
  migrateRecursively ::
    (PersistBackend m, PersistEntity v) =>
    -- | migrate schema
    (String -> m SingleMigration) ->
    -- | migrate entity
    (EntityDef -> m SingleMigration) ->
    -- | migrate list
    (DbType -> m SingleMigration) ->
    -- | initial entity
    v ->
    Migration m
  migrateRecursively migS migE migL v = result
    where
      result = migEntity $ entityDef proxy v
      proxy = (undefined :: t m a -> proxy (Conn m)) result
      go l@(DbList lName t) = f ("list " ++ lName) (migL l) (go t)
      go (DbEmbedded (EmbeddedDef _ ts) ref) = mapM_ (go . snd) ts >> migRef ref
      go (DbTypePrimitive _ _ _ ref) = migRef ref
      allSubtypes = map snd . concatMap constrParams . constructors
      migRef ref = case ref of
        Just (Left (e, _), _, _) -> migEntity e
        _ -> pure ()
      migEntity e = do
        case entitySchema e of
          Just name -> f ("schema " ++ name) (migS name) (pure ())
          Nothing -> pure ()
        f ("entity " ++ mainTableName id e) (migE e) (mapM_ go (allSubtypes e))
      f name mig cont = do
        a <- gets (Map.lookup name)
        when (isNothing a) $
          lift mig >>= modify . Map.insert name >> cont
  
  migrateSchema :: SchemaAnalyzer conn => MigrationPack conn -> String -> Action conn SingleMigration
  migrateSchema MigrationPack {..} schema = do
    x <- schemaExists schema
    pure $
      if x
        then Right []
        else showAlterDb $ CreateSchema schema False
  
  migrateEntity :: (SchemaAnalyzer conn, PersistBackendConn conn) => MigrationPack conn -> EntityDef -> Action conn SingleMigration
  migrateEntity m@MigrationPack {..} e = do
    autoKeyType <- fmap getDefaultAutoKeyType phantomDb
    let name = entityName e
        constrs = constructors e
        mainTableQuery = "CREATE TABLE " ++ escape name ++ " (" ++ mainTableId ++ " " ++ autoincrementedKeyTypeName ++ ", discr INTEGER NOT NULL)"
        expectedMainStructure = TableInfo [Column "id" False autoKeyType Nothing, Column "discr" False DbInt32 Nothing] [UniqueDef Nothing (UniquePrimary True) [Left "id"]] []
  
    if isSimple constrs
      then do
        x <- analyzeTable (entitySchema e, name)
        -- check whether the table was created for multiple constructors before
        case x of
          Just old
            | null $ getAlters m old expectedMainStructure ->
              pure $
                Left
                  ["Datatype with multiple constructors was truncated to one constructor. Manual migration required. Datatype: " ++ name]
          _ -> fmap snd $ migConstr e $ head constrs
      else do
        mainStructure <- analyzeTable (entitySchema e, name)
        let constrTable c = name ++ [delim] ++ constrName c
        res <- mapM (migConstr e) constrs
        pure $ case mainStructure of
          Nothing ->
            -- no constructor tables can exist if there is no main data table
            let orphans = filter (fst . fst) $ zip res constrs
             in if null orphans
                  then mergeMigrations $ Right [(False, defaultPriority, mainTableQuery)] : map snd res
                  else Left $ map (\(_, c) -> "Orphan constructor table found: " ++ constrTable c) orphans
          Just mainStructure' ->
            if null $ getAlters m mainStructure' expectedMainStructure
              then
                let -- the datatype had also many constructors before
                    -- check whether any new constructors appeared and increment older discriminators, which were shifted by newer constructors inserted not in the end
                    updateDiscriminators = go 0 . map (head &&& length) . group . map fst $ res
                      where
                        go acc ((False, n) : (True, n2) : xs) = (False, defaultPriority, "UPDATE " ++ escape name ++ " SET discr = discr + " ++ show n ++ " WHERE discr >= " ++ show acc) : go (acc + n + n2) xs
                        go acc ((True, n) : xs) = go (acc + n) xs
                        go _ _ = []
                 in mergeMigrations $ Right updateDiscriminators : map snd res
              else Left ["Unexpected structure of main table for Datatype: " ++ name ++ ". Table info: " ++ show mainStructure']
  
  migrateList :: (SchemaAnalyzer conn, PersistBackendConn conn) => MigrationPack conn -> DbType -> Action conn SingleMigration
  migrateList m@MigrationPack {..} (DbList mainName t) = do
    autoKeyType <- fmap getDefaultAutoKeyType phantomDb
    let valuesName = mainName ++ delim : "values"
        (valueCols, valueRefs) = (($ []) . mkColumns autoKeyType) &&& mkReferences autoKeyType $ ("value", t)
        refs' = Reference (Nothing, mainName) [("id", "id")] (Just Cascade) Nothing : valueRefs
        expectedMainStructure = TableInfo [Column "id" False autoKeyType Nothing] [UniqueDef Nothing (UniquePrimary True) [Left "id"]] []
        mainQuery = "CREATE TABLE " ++ escape mainName ++ " (id " ++ autoincrementedKeyTypeName ++ ")"
        (addInCreate, addInAlters) = addUniquesReferences [] refs'
        expectedValuesStructure = TableInfo valueColumns [] (map (\x -> (Nothing, x)) refs')
        valueColumns = Column "id" False autoKeyType Nothing : Column "ord" False DbInt32 Nothing : valueCols
        valuesQuery = "CREATE TABLE " ++ escape valuesName ++ " (" ++ intercalate ", " (map showColumn valueColumns ++ addInCreate) ++ ")"
    -- TODO: handle case when outer entity has a schema
    mainStructure <- analyzeTable (Nothing, mainName)
    valuesStructure <- analyzeTable (Nothing, valuesName)
    let triggerMain = []
    (_, triggerValues) <- migTriggerOnDelete (Nothing, valuesName) $ mkDeletes escape ("value", t)
    pure $ case (mainStructure, valuesStructure) of
      (Nothing, Nothing) ->
        let rest = [AlterTable (Nothing, valuesName) valuesQuery expectedValuesStructure expectedValuesStructure addInAlters]
         in mergeMigrations $ map showAlterDb $ [AddTable mainQuery, AddTable valuesQuery] ++ rest ++ triggerMain ++ triggerValues
      (Just mainStructure', Just valuesStructure') ->
        let f name a b =
              if null $ getAlters m a b
                then []
                else ["List table " ++ name ++ " error. Expected: " ++ show b ++ ". Found: " ++ show a]
            errors = f mainName mainStructure' expectedMainStructure ++ f valuesName valuesStructure' expectedValuesStructure
         in if null errors then Right [] else Left errors
      (_, Nothing) -> Left ["Found orphan main list table " ++ mainName]
      (Nothing, _) -> Left ["Found orphan list values table " ++ valuesName]
  migrateList _ t = fail $ "migrateList: expected DbList, got " ++ show t
  
  getAlters ::
    MigrationPack m ->
    -- | From database
    TableInfo ->
    -- | From datatype
    TableInfo ->
    [AlterTable]
  getAlters m@MigrationPack {..} (TableInfo oldColumns oldUniques oldRefs) (TableInfo newColumns newUniques newRefs) = tableAlters
    where
      (oldOnlyColumns, newOnlyColumns, commonColumns) = matchElements ((==) `on` colName) oldColumns newColumns
      (oldOnlyUniques, newOnlyUniques, commonUniques) = matchElements compareUniqs oldUniques newUniques
      (oldOnlyRefs, newOnlyRefs, _) = matchElements compareRefs oldRefs newRefs
      primaryColumns = lefts $ concatMap uniqueDefFields $ filter ((== UniquePrimary True) . uniqueDefType) oldUniques
  
      colAlters = mapMaybe (\(a, b) -> mkAlterColumn b $ migrateColumn m a b) (filter ((`notElem` primaryColumns) . colName . fst) commonColumns)
      mkAlterColumn col alters = if null alters then Nothing else Just $ AlterColumn col alters
      tableAlters =
        map dropUnique oldOnlyUniques
          ++ map (DropReference . fromMaybe (error "getAlters: old reference does not have name") . fst) oldOnlyRefs
          ++ map (DropColumn . colName) oldOnlyColumns
          ++ map AddColumn newOnlyColumns
          ++ colAlters
          ++ map AddUnique newOnlyUniques
          ++ concatMap (uncurry migrateUniq) commonUniques
          ++ map (AddReference . snd) newOnlyRefs
  
  -- from database, from datatype
  migrateColumn :: MigrationPack m -> Column -> Column -> [AlterColumn]
  migrateColumn MigrationPack {..} (Column _ isNull1 type1 def1) (Column _ isNull2 type2 def2) = modDef ++ modNull ++ modType
    where
      modNull = case (isNull1, isNull2) of
        (False, True) -> [IsNull]
        (True, False) -> case def2 of
          Nothing -> [NotNull]
          Just s -> [UpdateValue s, NotNull]
        _ -> []
      modType = if compareTypes type1 type2 then [] else [Type type2]
      modDef = case (def1, def2) of
        (Nothing, Nothing) -> []
        (Just def1', Just def2') | compareDefaults def1' def2' -> []
        _ -> [maybe NoDefault Default def2]
  
  -- from database, from datatype
  migrateUniq :: UniqueDefInfo -> UniqueDefInfo -> [AlterTable]
  migrateUniq u1@(UniqueDef _ _ cols1) u2@(UniqueDef _ _ cols2) =
    if haveSameElems (==) cols1 cols2
      then []
      else [dropUnique u1, AddUnique u2]
  
  dropUnique :: UniqueDefInfo -> AlterTable
  dropUnique (UniqueDef name typ _) =
    case typ of
      UniqueConstraint -> DropConstraint name'
      UniqueIndex -> DropIndex name'
      UniquePrimary _ -> DropConstraint name'
    where
      name' = fromMaybe (error "dropUnique: constraint which should be dropped does not have a name") name
  
  defaultMigConstr :: (SchemaAnalyzer conn, PersistBackendConn conn) => MigrationPack conn -> EntityDef -> ConstructorDef -> Action conn (Bool, SingleMigration)
  defaultMigConstr m@MigrationPack {..} e constr = do
    let simple = isSimple $ constructors e
        name = entityName e
        qualifiedCName = (entitySchema e, if simple then name else name ++ [delim] ++ constrName constr)
    autoKeyType <- fmap getDefaultAutoKeyType phantomDb
    tableStructure <- analyzeTable qualifiedCName
    let dels = concatMap (mkDeletes escape) $ constrParams constr
    (triggerExisted, delTrigger) <- migTriggerOnDelete qualifiedCName dels
    updTriggers <- concatMap snd <$> migTriggerOnUpdate qualifiedCName dels
  
    let (expectedTableStructure, (addTable, addInAlters)) =
          ( case constrAutoKeyName constr of
              Nothing -> (TableInfo columns uniques (mkRefs refs), f [] columns uniques refs)
              Just keyName ->
                let keyColumn = Column keyName False autoKeyType Nothing
                 in if simple
                      then
                        ( TableInfo (keyColumn : columns) (uniques ++ [UniqueDef Nothing (UniquePrimary True) [Left keyName]]) (mkRefs refs),
                          f [escape keyName ++ " " ++ autoincrementedKeyTypeName] columns uniques refs
                        )
                      else
                        let columns' = keyColumn : columns
                            refs' = refs ++ [Reference (entitySchema e, name) [(keyName, mainTableId)] (Just Cascade) Nothing]
                            uniques' = uniques ++ [UniqueDef Nothing UniqueConstraint [Left keyName]]
                         in (TableInfo columns' uniques' (mkRefs refs'), f [] columns' uniques' refs')
          )
          where
            (columns, refs) = foldr (mkColumns autoKeyType) [] &&& concatMap (mkReferences autoKeyType) $ constrParams constr
            uniques = map (\u -> u {uniqueDefFields = concatMap (either (map Left . ($ []) . flatten id) (pure . Right)) $ uniqueDefFields u}) $ constrUniques constr
            f autoKey cols uniqs refs' = (addTable', addInAlters')
              where
                (addInCreate, addInAlters') = addUniquesReferences uniqs refs'
                items = autoKey ++ map showColumn cols ++ addInCreate
                addTable' = "CREATE TABLE " ++ tableName escape e constr ++ " (" ++ intercalate ", " items ++ ")"
            mkRefs = map (\r -> (Nothing, r))
  
        (migErrs, constrExisted, mig) = case tableStructure of
          Nothing ->
            let rest = AlterTable qualifiedCName addTable expectedTableStructure expectedTableStructure addInAlters
             in ([], False, [AddTable addTable, rest])
          Just oldTableStructure ->
            let alters = getAlters m oldTableStructure expectedTableStructure
                alterTable =
                  if null alters
                    then []
                    else [AlterTable qualifiedCName addTable oldTableStructure expectedTableStructure alters]
             in ([], True, alterTable)
        -- this can happen when an ephemeral field was added. Consider doing something else except throwing an error
        allErrs =
          if constrExisted == triggerExisted || (constrExisted && null dels)
            then migErrs
            else ("Both trigger and constructor table must exist: " ++ show qualifiedCName) : migErrs
    pure
      ( constrExisted,
        if null allErrs
          then mergeMigrations $ map showAlterDb $ mig ++ delTrigger ++ updTriggers
          else Left allErrs
      )
  
  -- on delete removes all ephemeral data
  -- returns column name and delete statement for the referenced table
  mkDeletes :: (String -> String) -> (String, DbType) -> [(String, String)]
  mkDeletes esc = concat . traverseDbType f
    where
      f (DbList ref _) [col] = [(col, "DELETE FROM " ++ esc ref ++ " WHERE id=old." ++ esc col ++ ";")]
      f _ _ = []
  
  showReferenceAction :: ReferenceActionType -> String
  showReferenceAction NoAction = "NO ACTION"
  showReferenceAction Restrict = "RESTRICT"
  showReferenceAction Cascade = "CASCADE"
  showReferenceAction SetNull = "SET NULL"
  showReferenceAction SetDefault = "SET DEFAULT"
  
  readReferenceAction :: String -> Maybe ReferenceActionType
  readReferenceAction c = case c of
    "NO ACTION" -> Just NoAction
    "RESTRICT" -> Just Restrict
    "CASCADE" -> Just Cascade
    "SET NULL" -> Just SetNull
    "SET DEFAULT" -> Just SetDefault
    _ -> Nothing
  
  class PersistBackendConn conn => SchemaAnalyzer conn where
    schemaExists ::
      (PersistBackend m, Conn m ~ conn) =>
      -- | Schema name
      String ->
      m Bool
    getCurrentSchema :: (PersistBackend m, Conn m ~ conn) => m (Maybe String)
    listTables ::
      (PersistBackend m, Conn m ~ conn) =>
      -- | Schema name
      Maybe String ->
      m [String]
    listTableTriggers ::
      (PersistBackend m, Conn m ~ conn) =>
      -- | Qualified table name
      QualifiedName ->
      m [String]
    analyzeTable ::
      (PersistBackend m, Conn m ~ conn) =>
      -- | Qualified table name
      QualifiedName ->
      m (Maybe TableInfo)
    analyzeTrigger ::
      (PersistBackend m, Conn m ~ conn) =>
      -- | Qualified trigger name
      QualifiedName ->
      m (Maybe String)
    analyzeFunction ::
      (PersistBackend m, Conn m ~ conn) =>
      -- | Qualified function name
      QualifiedName ->
      -- | Argument types, return type, and body
      m (Maybe (Maybe [DbTypePrimitive], Maybe DbTypePrimitive, String))
    getMigrationPack :: (PersistBackend m, Conn m ~ conn) => m (MigrationPack conn)
  