{-# 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.HashSet.InsOrd qualified as OSet
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 :: 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],
    selectSqlOffset :: Maybe Int,
    selectSqlTop :: MySQL.Top
  }
  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 :: MySQL.From -> Text
selectFromName =
  \case
    MySQL.FromQualifiedTable (MySQL.Aliased {aliasedThing = MySQL.TableName {name}}) ->
      name
    MySQL.FromSelect (MySQL.Aliased {aliasedThing = MySQL.Select {selectFrom}}) ->
      selectFromName selectFrom

--------------------------------------------------------------------------------
-- 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 = OSet.fromList selectProjections,
      selectFinalWantedFields = selectWantedFields,
      ..
    }

-- | 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,
      ..
    }