{-# LANGUAGE DuplicateRecordFields #-}

-- | Make a plan for the data loader to execute (.Execute).
--
-- It will produce a graph of actions, to be executed by .Execute.
module Hasura.Backends.MySQL.DataLoader.Plan
  ( Ref,
    PlannedAction (..),
    Action (..),
    Select (..),
    Join (..),
    Relationship (..),
    FieldName (..),
    HeadAndTail (..),
    toFieldName,
    runPlan,
    planSelectHeadAndTail,
    actionsForest,
    selectQuery,
  )
where

import Data.Aeson
import Data.Bifunctor
import Data.Graph
import Data.Sequence qualified as Seq
import Data.String
import Hasura.Backends.MySQL.Types qualified as MySQL
import Hasura.Prelude hiding (head, second, tail, tell)

--------------------------------------------------------------------------------
-- Types

-- | A reference to a result of loading a recordset from the database.
data Ref = Ref
  { -- | This index will be generated by the planner.
    idx :: !Int,
    -- | A display name. The idx gives us uniqueness.
    text :: !Text
  }
  deriving (Show, Eq, Generic, Ord)

instance Hashable Ref

-- | A almost-the-same version of Select from Types.Internal, except
-- with some fields used for planning and executing.
data Select = Select
  { selectAggUnwrap :: !(Maybe Text),
    selectFrom :: !(Maybe MySQL.From),
    selectGroupBy :: ![MySQL.FieldName],
    selectHaskellJoins :: ![MySQL.Join],
    selectOrderBy :: !(Maybe (NonEmpty MySQL.OrderBy)),
    selectProjections :: ![MySQL.Projection],
    selectRelationship :: !(Maybe Relationship),
    selectWhere :: !MySQL.Where,
    selectWantedFields :: !(Maybe [Text])
  }
  deriving (Show)

-- | An join action.
data Join = Join
  { -- | Join this side...
    leftRecordSet :: Ref,
    -- | with this side.
    rightRecordSet :: Ref,
    -- | Join only the top N results. It's important that we do this
    -- IN HASKELL, therefore this is not part of the generated SQL.
    joinRhsTop :: !MySQL.Top,
    -- | Offset applied to the right-hand-side table.
    joinRhsOffset :: !(Maybe Int),
    -- | Type of relational join to do.
    joinType :: !MySQL.JoinType,
    -- | Field name to return the join result as; e.g. "albums" for an
    -- artist with an array relation of albums.
    joinFieldName :: !Text,
    -- | The SQL queries may achieve the data using joining fields,
    -- but those fields aren't supposed to be returned back to the
    -- user. To avoid that, we explicitly specify which fields are
    -- wanted from this join. E.g. "title" and "year", but not
    -- artist_id which was used to Haskell-join the row with an
    -- album_artist_id, or whatever.
    wantedFields :: !(Maybe [Text])
  }
  deriving (Show)

-- | An action that the executor will perform. Either pull data from
-- the database directly via a select, or join two other actions'
-- record sets together.
data Action
  = SelectAction Select
  | JoinAction Join
  deriving (Show)

-- | An action planned, with a unique reference. I.e. the @action@
-- performed yields a result stored at reference @ref@.
data PlannedAction = PlannedAction
  { ref :: Ref,
    action :: Action
  }
  deriving (Show)

-- | A relationship lets the executor insert on-the-fly WHERE
-- fkey1=fkey2 for relationships. These can only be inserted
-- on-the-fly and aren't known at the time of planning, because the
-- keys come from the left-hand-side table for a join.
data Relationship = Relationship
  { leftRecordSet :: Ref,
    rightTable :: MySQL.EntityAlias,
    joinType :: MySQL.JoinType
  }
  deriving (Show)

-- | Just a wrapper to clarify some types. It's different from the
-- MySQL.FieldName because it doesn't care about schemas: schemas
-- aren't returned in recordsets from the database.
newtype FieldName
  = FieldName Text
  deriving (Show, Ord, Eq, Hashable, FromJSON, ToJSONKey, IsString)

-- | The reason for this is subtle. Read this documentation. For each
-- join on a select (see above, there is a list), we split that out
-- into three jobs:
--
-- 1. One job for the left hand side (i.e. the select).
-- 2. One job for the right hand side (i.e. the join).
-- 3. One job to join them (And in the darkness bind them...)
--
-- This is performed as a fold, like: @foldM planJoin head joins@. A
-- nice linked-list or tree-like structure arises. The planner code
-- produces a graph out of this; so it's possible that some
-- parallelism can be achieved by running multiple jobs at once.
--
-- The "head" is the first, original select. The "tail" is the
-- (indirectly) linked list of joins. That list may also be empty. In
-- that case, the tail is simply the same as the head.
--
-- If the tail is different to the head, then we choose the tail, as
-- it represents the joined up version of both. If they're the same,
-- we take whichever.
data HeadAndTail = HeadAndTail
  { head :: Ref,
    tail :: Ref
  }

-- | We're simply accumulating a set of actions with this. The counter
-- lets us generate unique refs.
data PlanState = PlanState
  { actions :: !(Seq PlannedAction),
    counter :: !Int
  }

-- | Simple monad to collect actions.
newtype Plan a = Plan
  { unPlan :: State PlanState a
  }
  deriving (Functor, Applicative, Monad, MonadState PlanState)

--------------------------------------------------------------------------------
-- Conversions

-- | Note that we're intentionally discarding the table qualification.
toFieldName :: MySQL.FieldName -> FieldName
toFieldName (MySQL.FieldName {fName = t}) = FieldName t

joinAliasName :: MySQL.EntityAlias -> Text
joinAliasName (MySQL.EntityAlias {entityAliasText}) = entityAliasText

-- | Used for display purposes, not semantic content.
selectFromName :: Maybe MySQL.From -> Text
selectFromName = \case
  Nothing -> ""
  Just x -> case x of
    MySQL.FromQualifiedTable (MySQL.Aliased {aliasedThing = MySQL.TableName {name}}) -> name
    -- This constructor will be removed in the subsequent PR.
    MySQL.FromOpenJson {} -> ""

--------------------------------------------------------------------------------
-- Run planner

runPlan :: Plan r -> (r, [PlannedAction])
runPlan =
  second (toList . actions)
    . flip runState (PlanState {actions = mempty, counter = 0})
    . unPlan

--------------------------------------------------------------------------------
-- Planners

-- | See the documentation for 'HeadAndTail'.
planSelectHeadAndTail :: Maybe Relationship -> Maybe Text -> MySQL.Select -> Plan HeadAndTail
planSelectHeadAndTail relationship joinExtractPath select0 = do
  ref <- generate (selectFromName (MySQL.selectFrom select0))
  let select = fromSelect relationship joinExtractPath select0
      action = SelectAction select
  tell PlannedAction {ref, action}
  joinsFinalRef <- foldM planJoin ref (selectHaskellJoins select)
  pure
    ( let head = ref
          tail = case selectHaskellJoins select of
            [] -> ref
            _ -> joinsFinalRef
       in HeadAndTail {head, tail}
    )

-- | Given a left-hand-side table and a join spec, produce a single
-- reference that refers to the composition of the two.
planJoin :: Ref -> MySQL.Join -> Plan Ref
planJoin leftRecordSet join' = do
  ref <- generate (joinAliasName (MySQL.joinRightTable join'))
  rightRecordSet <-
    fmap
      (\HeadAndTail {..} -> tail)
      ( planSelectHeadAndTail
          ( Just
              ( Relationship
                  { leftRecordSet,
                    joinType = MySQL.joinType join',
                    rightTable = MySQL.joinRightTable join'
                  }
              )
          )
          Nothing
          (MySQL.joinSelect join')
      )
  let action =
        JoinAction
          Join
            { leftRecordSet,
              rightRecordSet,
              wantedFields = MySQL.selectFinalWantedFields (MySQL.joinSelect join'),
              joinRhsTop = MySQL.joinTop join',
              joinRhsOffset = MySQL.joinOffset join',
              joinFieldName = MySQL.joinFieldName join',
              joinType = MySQL.joinType join',
              ..
            }
  tell PlannedAction {ref, action}
  pure ref

--------------------------------------------------------------------------------
-- Monad helpers

-- | Write the planned action to the state, like a writer's @tell@.
tell :: PlannedAction -> Plan ()
tell action = modify' (\s -> s {actions = actions s Seq.:|> action})

-- | Generate a unique reference with a label for debugging.
generate :: Text -> Plan Ref
generate text = do
  idx <- gets counter
  modify' (\s -> s {counter = counter s + 1})
  pure (Ref {idx, text})

--------------------------------------------------------------------------------
-- Graphing the plan to a forest

-- | Graph the set of planned actions ready for execution in the correct order.
actionsForest :: (Graph -> Graph) -> [PlannedAction] -> Forest PlannedAction
actionsForest transform actions =
  let (graph, vertex2Node, _key2Vertex) =
        graphFromEdges
          ( map
              ( \PlannedAction {ref, action} ->
                  ( action,
                    ref,
                    map
                      (\PlannedAction {ref = r} -> r)
                      (filter (elem ref . plannedActionRefs) actions)
                  )
              )
              actions
          )
   in fmap
        ( fmap
            ((\(action, ref, _refs) -> PlannedAction {ref, action}) . vertex2Node)
        )
        (dff (transform graph))
  where
    plannedActionRefs PlannedAction {action} =
      case action of
        SelectAction Select {selectRelationship} ->
          case selectRelationship of
            Just Relationship {leftRecordSet} -> [leftRecordSet]
            Nothing -> mempty
        JoinAction Join {leftRecordSet, rightRecordSet} ->
          [leftRecordSet, rightRecordSet]

--------------------------------------------------------------------------------
-- Build a query

-- | Used by the executor to produce a plain old select that can be
-- sent to the MySQL server.
selectQuery :: Select -> MySQL.Select
selectQuery Select {..} =
  MySQL.Select
    { selectJoins = selectHaskellJoins,
      selectProjections = selectProjections,
      selectTop = MySQL.NoTop,
      selectOffset = Nothing,
      selectFor = MySQL.NoFor,
      ..
    }

-- | From a plain select, and possibly a parent/left-hand-side
-- relationship, produce a select that is useful for execution.
fromSelect :: Maybe Relationship -> Maybe Text -> MySQL.Select -> Select
fromSelect selectRelationship selectAggUnwrap select@MySQL.Select {..} =
  Select
    { selectHaskellJoins = selectJoins,
      selectWantedFields = MySQL.selectFinalWantedFields select,
      selectGroupBy = [],
      selectProjections = toList selectProjections,
      ..
    }