module Hasql.Generate.TH
    ( GenerateConfig
    , fromTable
    , fromType
    , fromView
    , generate
    , withDefaultedCols
    , withDerivations
    , withOverrides
    ) where

----------------------------------------------------------------------------------------------------

import           Control.Applicative                ( (<*>) )
import           Control.Monad                      ( mapM, mapM_, return )
import           Control.Monad.Fail                 ( MonadFail (fail) )

import           Data.Bool
    ( Bool (..)
    , not
    , otherwise
    )
import           Data.Char                          ( toLower, toUpper )
import           Data.Eq                            ( (==) )
import           Data.Foldable                      ( foldl )
import           Data.Function                      ( flip, ($), (.) )
import           Data.Functor                       ( (<$>) )
import qualified Data.Functor.Contravariant         as Contravariant
import           Data.Int                           ( Int )
import           Data.List
    ( break
    , concatMap
    , elem
    , filter
    , intercalate
    , length
    , lookup
    , map
    , notElem
    , null
    , partition
    , zip
    , zipWith
    , (!!)
    )
import           Data.Maybe                         ( Maybe (..) )
import           Data.Semigroup                     ( (<>) )
import           Data.String                        ( String )
import qualified Data.String
import qualified Data.Text
import qualified Data.Tuple

import qualified Hasql.Decoders
import qualified Hasql.Encoders
import qualified Hasql.Generate.Class
import qualified Hasql.Generate.Codec
import qualified Hasql.Generate.Column
import           Hasql.Generate.Config              ( Config (..) )
import           Hasql.Generate.Connection
    ( toConnString
    , withCompileTimeConnection
    )
import           Hasql.Generate.Internal.Introspect
    ( ColumnInfo (..)
    , introspectColumns
    , introspectEnumLabels
    , introspectPrimaryKey
    )
import qualified Hasql.Session
import qualified Hasql.Statement

import           Language.Haskell.TH

import           Prelude
    ( enumFromTo
    , mconcat
    , show
    , (+)
    , (-)
    )

----------------------------------------------------------------------------------------------------

-- | Whether the target relation is a table (full CRUD), a view (read-only), or a type (enum).
data RelationKind = Table | View | Type

{- | Configuration for a table, view, or type code-generation request.

    Use 'fromTable', 'fromView', or 'fromType' to create a default config,
    then chain modifier functions with @(&)@ to customise it:

    @
    fromTable \"public\" \"users\"
        & withDerivations [''Show, ''Eq, ''Generic]
        & withOverrides [(\"timestamptz\", ''UTCTime)]
        & withDefaultedCols [\"id\"]
    @
-}
data GenerateConfig
    = GenerateConfig
      { gcSchema        :: String
      , gcTable         :: String
      , gcKind          :: RelationKind
      , gcDerivations   :: [Name]
      , gcOverrides     :: [(String, Name)]
      , gcDefaultedCols :: [String]
      }

{- | Create a 'GenerateConfig' for the given schema and table with no derivations
    and no type overrides. Generates full CRUD code.
-}
fromTable :: String -> String -> GenerateConfig
fromTable schema table =
  GenerateConfig
    { gcSchema = schema
    , gcTable = table
    , gcKind = Table
    , gcDerivations = []
    , gcOverrides = []
    , gcDefaultedCols = []
    }

{- | Create a 'GenerateConfig' for the given schema and view with no derivations
    and no type overrides. Generates read-only code: a record type, a decoder,
    a SELECT statement, and a 'HasView' instance.
-}
fromView :: String -> String -> GenerateConfig
fromView schema view =
  GenerateConfig
    { gcSchema = schema
    , gcTable = view
    , gcKind = View
    , gcDerivations = []
    , gcOverrides = []
    , gcDefaultedCols = []
    }

{- | Create a 'GenerateConfig' for the given schema and enum type with no
    derivations. Generates a Haskell sum type, a 'PgCodec' instance,
    and a 'PgColumn' instance. Overrides are ignored for types.

    The splice must appear before any @fromTable@ whose table has a column of
    this enum type, since TH processes splices top to bottom.

    __Orphan instance warning:__ The generated @PgColumn@ instance will trigger
    GHC's @-Worphans@ warning because the functional dependency's determining
    types are @Symbol@ literals, which GHC never considers local to any user
    module. This is expected and harmless. Add the following pragma to any
    module that uses @fromType@:

    @
    \{\-\# OPTIONS_GHC -Wno-orphans \#\-\}
    @
-}
fromType :: String -> String -> GenerateConfig
fromType schema typeName =
  GenerateConfig
    { gcSchema = schema
    , gcTable = typeName
    , gcKind = Type
    , gcDerivations = []
    , gcOverrides = []
    , gcDefaultedCols = []
    }

-- | Append derivation class 'Name's to the config.
withDerivations :: [Name] -> GenerateConfig -> GenerateConfig
withDerivations names cfg = cfg {gcDerivations = gcDerivations cfg <> names}

{- | Append per-table PG type overrides. Each pair maps a PostgreSQL type name
    (e.g. @\"timestamptz\"@) to a Haskell type 'Name' (e.g. @''UTCTime@).
    The override type must still have a 'PgCodec' instance.
-}
withOverrides :: [(String, Name)] -> GenerateConfig -> GenerateConfig
withOverrides ovs cfg = cfg {gcOverrides = gcOverrides cfg <> ovs}

{- | Exclude the named columns from INSERT statements. Each column must exist
    in the table and have a database default (@atthasdef = true@ in
    @pg_catalog@); a compile-time error is raised otherwise.

    @
    fromTable \"public\" \"users\" & withDefaultedCols [\"id\", \"created_at\"]
    @
-}
withDefaultedCols :: [String] -> GenerateConfig -> GenerateConfig
withDefaultedCols cols cfg = cfg {gcDefaultedCols = gcDefaultedCols cfg <> cols}

{- | Terminal step that introspects a PostgreSQL database at compile time and
    generates the provided data construct:

    Usage:
    @
    \$(generate def (fromTable \"public\" \"users\" & withDerivations [''Show, ''Eq]))
    @
-}
generate :: Config -> GenerateConfig -> Q [Dec]
generate config genConfig =
  case kind of
    Type -> do
      labels <- runIO $ withCompileTimeConnection connStr $ \conn ->
        introspectEnumLabels conn schema table
      if null labels
        then fail ("hasql-generate: no enum labels found for " <> schema <> "." <> table)
        else generateTypeDecs schema table (pascalCase table) derivNames labels
    Table -> do
      (typeName, resolvedCols, pkCols) <- getPgData
      let pkInfo = buildPkInfo resolvedCols pkCols
      generateAllDecs config defaultedCols schema table typeName derivNames resolvedCols pkInfo
    View -> do
      (typeName, resolvedCols, _) <- getPgData
      generateViewDecs schema table typeName derivNames resolvedCols
  where
    connStr = toConnString (connection config)
    schema = gcSchema genConfig
    table = gcTable genConfig
    kind = gcKind genConfig
    derivNames = gcDerivations genConfig
    defaultedCols = gcDefaultedCols genConfig
    mergedOverrides = gcOverrides genConfig <> globalOverrides config
    getPgData = do
      (columns, pkCols') <- runIO $ withCompileTimeConnection connStr $ \conn -> do
        cols <- introspectColumns conn schema table
        pks <- case kind of
          Table -> introspectPrimaryKey conn schema table
          _     -> return []
        return (cols, pks)
      if null columns
        then fail ("hasql-generate: no columns found for " <> schema <> "." <> table)
        else do
          resolvedCols <- mapM (resolveColumnWithOverrides mergedOverrides) columns
          let typName = pascalCase table
              resolvedCols' =
                if allowDuplicateRecordFields config
                  then resolvedCols
                  else map (prefixFieldName typName) resolvedCols
          return (typName, resolvedCols', pkCols')

----------------------------------------------------------------------------------------------------

data ResolvedColumn
    = ResolvedColumn
      { rcColName    :: String
      , rcFieldName  :: String
      , rcType       :: Type
      , rcNotNull    :: Bool
      , rcHasDefault :: Bool
      , rcPgCast     :: Maybe String
      }

data PkInfo
    = NoPrimaryKey
    | SinglePk ResolvedColumn
    | CompositePk [ResolvedColumn]

----------------------------------------------------------------------------------------------------

{-  Resolve a column's Haskell type. If the column's PG type name appears in
    the overrides list, use the override 'Name' directly (as @ConT@). Otherwise
    fall back to 'PgColumn' instance resolution via 'reifyInstances'.
-}
resolveColumnWithOverrides :: [(String, Name)] -> ColumnInfo -> Q ResolvedColumn
resolveColumnWithOverrides overrides col = do
  hsType <- case lookup (colPgType col) overrides of
    Just overrideName -> return (ConT overrideName)
    Nothing -> do
      a <- newName "a"
      insts <-
        reifyInstances
          ''Hasql.Generate.Column.PgColumn
          [ LitT (StrTyLit (colPgSchema col))
          , LitT (StrTyLit (colPgType col))
          , VarT a
          ]
      case insts of
        [InstanceD _ _ (AppT (AppT (AppT _ _) _) ty) _] -> return ty
        [] ->
          fail
            ( "hasql-generate: no PgColumn instance for pg type '"
                <> colPgSchema col
                <> "."
                <> colPgType col
                <> "' (column '"
                <> colName col
                <> "')"
            )
        _ ->
          fail
            ( "hasql-generate: multiple PgColumn instances for pg type '"
                <> colPgSchema col
                <> "."
                <> colPgType col
                <> "' (column '"
                <> colName col
                <> "') — expected exactly one"
            )
  let pgCast =
        if colIsEnum col
          then Just (quoteIdent (colPgSchema col) <> "." <> quoteIdent (colPgType col))
          else Nothing
  return
    ResolvedColumn
      { rcColName = colName col
      , rcFieldName = sanitizeField (camelCase (colName col))
      , rcType = hsType
      , rcNotNull = colNotNull col
      , rcHasDefault = colHasDefault col
      , rcPgCast = pgCast
      }

----------------------------------------------------------------------------------------------------

buildPkInfo :: [ResolvedColumn] -> [String] -> PkInfo
buildPkInfo _ [] = NoPrimaryKey
buildPkInfo resolvedCols [pkName] =
  case filter (\rc -> rcColName rc == pkName) resolvedCols of
    [rc] -> SinglePk rc
    _    -> NoPrimaryKey
buildPkInfo resolvedCols pkNames =
  let lookupPk pkn = filter (\rc -> rcColName rc == pkn) resolvedCols
      pkCols = concatMap lookupPk pkNames
   in if length pkCols == length pkNames
        then CompositePk pkCols
        else NoPrimaryKey

----------------------------------------------------------------------------------------------------

generateAllDecs
  :: Config -> [String] -> String -> String -> String -> [Name] -> [ResolvedColumn] -> PkInfo -> Q [Dec]
generateAllDecs config defaultedCols schema table typName derivNames resolvedCols pkInfo = do
  let useNtPk = newtypePrimaryKeys config
      pkTypName = typName <> "Pk"
      allowDupFields = allowDuplicateRecordFields config

  -- Generate PK wrapper type declarations (before main record so newtype is in scope)
  pkTypeDecs <-
    if useNtPk
      then case pkInfo of
        SinglePk rc -> do
          ntDec <- genPkNewtype pkTypName derivNames rc
          codecDec <- genPkNewtypeCodec pkTypName
          return (ntDec <> codecDec)
        CompositePk rcs -> do
          let pkRcs = map (pkRecordFieldName pkTypName allowDupFields) rcs
          recDec <- genPkRecord pkTypName derivNames pkRcs
          encDec <- genPkRecordEncoder pkTypName pkRcs
          return (recDec <> encDec)
        NoPrimaryKey -> return []
      else return []

  -- For single PK with newtypes, rewrite the PK column's type to the newtype
  let (resolvedCols', pkInfo') =
        if useNtPk
          then case pkInfo of
            SinglePk rc ->
              let rc' = rc {rcType = ConT (mkName pkTypName)}
                  cols' = map (\c -> if rcColName c == rcColName rc then rc' else c) resolvedCols
               in (cols', SinglePk rc')
            _ -> (resolvedCols, pkInfo)
          else (resolvedCols, pkInfo)

  -- Compute pkTypOverride for composite PK with newtypes
  let pkTypOverride = case (useNtPk, pkInfo) of
        (True, CompositePk _) -> Just pkTypName
        _                     -> Nothing

  -- Generate main record, decoder, encoder, insert (single + batch)
  dataDec <- genDataType typName derivNames resolvedCols'
  decoderDec <- genDecoder typName resolvedCols'
  encoderDec <- genEncoder typName resolvedCols'
  insertCols <- computeInsertCols defaultedCols resolvedCols'
  insDec <- genInsert schema table typName resolvedCols' insertCols
  insMany <- genInsertMany schema table typName resolvedCols' insertCols
  hasInsInst <- genHasInsertInstance typName

  -- Generate PK-dependent statements and instances (single + batch)
  pkDecs <- case pkInfo' of
    NoPrimaryKey -> return []
    _ -> do
      sel <- genSelectByPk schema table typName resolvedCols' pkInfo' pkTypOverride
      selMany <- genSelectMany schema table typName resolvedCols' pkInfo' pkTypOverride
      upd <- genUpdate schema table typName resolvedCols' pkInfo'
      updMany <- genUpdateMany schema table typName resolvedCols' pkInfo'
      del <- genDeleteByPk schema table typName pkInfo' pkTypOverride
      delMany <- genDeleteMany schema table typName pkInfo' pkTypOverride
      hasSelInst <- genHasSelectInstance typName pkInfo' pkTypOverride
      hasUpdInst <- genHasUpdateInstance typName
      hasDelInst <- genHasDeleteInstance typName pkInfo' pkTypOverride
      return (sel <> selMany <> upd <> updMany <> del <> delMany <> hasSelInst <> hasUpdInst <> hasDelInst)

  -- Generate HasPrimaryKey instance (for ALL tables with PKs)
  hasPkInst <- case pkInfo of
    NoPrimaryKey -> return []
    _ -> genHasPrimaryKeyInstance typName pkInfo pkInfo' useNtPk pkTypOverride allowDupFields

  return (pkTypeDecs <> dataDec <> decoderDec <> encoderDec <> insDec <> insMany <> hasInsInst <> pkDecs <> hasPkInst)

{-  Generate read-only declarations for a view: data type, decoder, SELECT
    statement, and 'HasView' instance.
-}
generateViewDecs
  :: String -> String -> String -> [Name] -> [ResolvedColumn] -> Q [Dec]
generateViewDecs schema table typName derivNames resolvedCols = do
  dataDec <- genDataType typName derivNames resolvedCols
  decoderDec <- genDecoder typName resolvedCols
  selectDec <- genSelectView schema table typName resolvedCols
  hasViewInst <- genHasViewInstance typName
  return (dataDec <> decoderDec <> selectDec <> hasViewInst)

{-  Generate declarations for a PostgreSQL enum type: a Haskell sum type,
    a 'PgCodec' instance (using @Decoders.enum@ / @Encoders.enum@), and
    a 'PgColumn' instance mapping the schema-qualified PG type name to the
    generated Haskell type.
-}
generateTypeDecs
  :: String -> String -> String -> [Name] -> [String] -> Q [Dec]
generateTypeDecs schema pgTypeName hsTypeName derivNames labels = do
  sumType <- genSumType hsTypeName derivNames labels
  codecInst <- genPgCodecInstance hsTypeName labels
  columnInst <- genPgColumnInstance schema pgTypeName hsTypeName
  enumInst <- genHasEnumInstance hsTypeName labels
  return (sumType <> codecInst <> columnInst <> enumInst)

genSumType :: String -> [Name] -> [String] -> Q [Dec]
genSumType hsTypeName derivNames labels = do
  let tName = mkName hsTypeName
      cons = map (\lbl -> NormalC (mkName (pascalCase lbl)) []) labels
      derivClauses = mkDerivClauses derivNames
  return [DataD [] tName [] Nothing cons derivClauses]

genPgCodecInstance :: String -> [String] -> Q [Dec]
genPgCodecInstance hsTypeName labels = do
  let tName = mkName hsTypeName
      -- pgDecode: Decoders.enum (\t -> case Data.Text.unpack t of "lbl" -> Just Con; ... _ -> Nothing)
      t = mkName "t"
      decoderMatches =
        map
          ( \lbl ->
              Match
                (LitP (StringL lbl))
                (NormalB (AppE (ConE 'Just) (ConE (mkName (pascalCase lbl)))))
                []
          )
          labels
          <> [Match WildP (NormalB (ConE 'Nothing)) []]
      decoderLam =
        LamE
          [VarP t]
          (CaseE (AppE (VarE 'Data.Text.unpack) (VarE t)) decoderMatches)
      decoderBody = AppE (VarE 'Hasql.Decoders.enum) decoderLam
      decoderDec = FunD 'Hasql.Generate.Codec.pgDecode [Clause [] (NormalB decoderBody) []]
      -- pgEncode: Encoders.enum (\x -> case x of Con -> "lbl"; ...)
      x = mkName "x"
      encoderMatches =
        map
          ( \lbl ->
              Match
                (ConP (mkName (pascalCase lbl)) [] [])
                (NormalB (textLit lbl))
                []
          )
          labels
      encoderLam = LamE [VarP x] (CaseE (VarE x) encoderMatches)
      encoderBody = AppE (VarE 'Hasql.Encoders.enum) encoderLam
      encoderDec = FunD 'Hasql.Generate.Codec.pgEncode [Clause [] (NormalB encoderBody) []]
      instanceType = AppT (ConT ''Hasql.Generate.Codec.PgCodec) (ConT tName)
  return [InstanceD Nothing [] instanceType [decoderDec, encoderDec]]

genPgColumnInstance :: String -> String -> String -> Q [Dec]
genPgColumnInstance schema pgTypeName hsTypeName = do
  let tName = mkName hsTypeName
      instanceType =
        AppT
          ( AppT
              (AppT (ConT ''Hasql.Generate.Column.PgColumn) (LitT (StrTyLit schema)))
              (LitT (StrTyLit pgTypeName))
          )
          (ConT tName)
  return [InstanceD Nothing [] instanceType []]

{-  Generate a @HasEnum@ instance for the given sum type.

    @allValues@ is a list of all constructors. @toText@ is a multi-clause
    function mapping each constructor to its PostgreSQL label. @fromText@
    unpacks the 'Text' to a 'String' and matches on string literals.
-}
genHasEnumInstance :: String -> [String] -> Q [Dec]
genHasEnumInstance hsTypeName labels = do
  let tName = mkName hsTypeName
      -- allValues = [Con1, Con2, ...]
      allValuesDec =
        FunD
          'Hasql.Generate.Class.allValues
          [Clause [] (NormalB (ListE (map (ConE . mkName . pascalCase) labels))) []]
      -- toText Con1 = "label1"; toText Con2 = "label2"; ...
      toTextClauses =
        map
          ( \lbl ->
              Clause
                [ConP (mkName (pascalCase lbl)) [] []]
                (NormalB (textLit lbl))
                []
          )
          labels
      toTextDec = FunD 'Hasql.Generate.Class.toText toTextClauses
      -- fromText t = case Data.Text.unpack t of "label1" -> Just Con1; ... _ -> Nothing
      t = mkName "t"
      fromTextMatches =
        map
          ( \lbl ->
              Match
                (LitP (StringL lbl))
                (NormalB (AppE (ConE 'Just) (ConE (mkName (pascalCase lbl)))))
                []
          )
          labels
          <> [Match WildP (NormalB (ConE 'Nothing)) []]
      fromTextBody =
        LamE
          [VarP t]
          (CaseE (AppE (VarE 'Data.Text.unpack) (VarE t)) fromTextMatches)
      fromTextDec =
        FunD
          'Hasql.Generate.Class.fromText
          [Clause [] (NormalB fromTextBody) []]
      instanceType = AppT (ConT ''Hasql.Generate.Class.HasEnum) (ConT tName)
  return [InstanceD Nothing [] instanceType [allValuesDec, toTextDec, fromTextDec]]

----------------------------------------------------------------------------------------------------
-- Data type generation
----------------------------------------------------------------------------------------------------

genDataType :: String -> [Name] -> [ResolvedColumn] -> Q [Dec]
genDataType typName derivNames cols = do
  let tName = mkName typName
      fields = map mkField cols
      con = RecC tName fields
      derivClauses = mkDerivClauses derivNames
  return [DataD [] tName [] Nothing [con] derivClauses]
  where
    mkField :: ResolvedColumn -> VarBangType
    mkField resCol =
      let fName = mkName (rcFieldName resCol)
          fieldBang = Bang NoSourceUnpackedness SourceStrict
          typ = fieldType resCol
       in (fName, fieldBang, typ)

fieldType :: ResolvedColumn -> Type
fieldType resCol
  | rcNotNull resCol = rcType resCol
  | otherwise = AppT (ConT ''Maybe) (rcType resCol)

{-  Partition deriving class 'Name's into stock and anyclass 'DerivClause's.
    Known stock-derivable classes get @deriving stock@; everything else gets
    @deriving anyclass@. Empty groups are omitted.
-}
mkDerivClauses :: [Name] -> [DerivClause]
mkDerivClauses names =
  let (stock, anyclass) = partition isStockDerivable names
   in clauseFor StockStrategy stock <> clauseFor AnyclassStrategy anyclass
  where
    clauseFor _ []      = []
    clauseFor strat nms = [DerivClause (Just strat) (map ConT nms)]

isStockDerivable :: Name -> Bool
isStockDerivable nm =
  nameBase nm
    `elem` [ "Show"
           , "Read"
           , "Eq"
           , "Ord"
           , "Bounded"
           , "Enum"
           , "Ix"
           , "Generic"
           , "Generic1"
           , "Data"
           , "Typeable"
           , "Lift"
           , "Functor"
           , "Foldable"
           , "Traversable"
           ]

----------------------------------------------------------------------------------------------------
-- Decoder generation
----------------------------------------------------------------------------------------------------

genDecoder :: String -> [ResolvedColumn] -> Q [Dec]
genDecoder typName cols = do
  let decName = mkName (camelCase typName <> "Decoder")
      tName = mkName typName
      sigTy = AppT (ConT ''Hasql.Decoders.Row) (ConT tName)
  sig <- sigD decName (return sigTy)
  body <- genDecoderBody tName cols
  let dec = ValD (VarP decName) (NormalB body) []
  return [sig, dec]

genDecoderBody :: Name -> [ResolvedColumn] -> Q Exp
genDecoderBody tName cols =
  case cols of
    [] -> fail "hasql-generate: cannot generate decoder for table with no columns"
    (x : xs) -> do
      let first = applyFmap (ConE tName) (columnDecodeExp x)
      foldl applyAp (return first) (map (return . columnDecodeExp) xs)

columnDecodeExp :: ResolvedColumn -> Exp
columnDecodeExp resCol =
  let nullability =
        if rcNotNull resCol
          then VarE 'Hasql.Decoders.nonNullable
          else VarE 'Hasql.Decoders.nullable
      codec = VarE 'Hasql.Generate.Codec.pgDecode
   in AppE (VarE 'Hasql.Decoders.column) (AppE nullability codec)

applyFmap :: Exp -> Exp -> Exp
applyFmap f x = InfixE (Just f) (VarE '(<$>)) (Just x)

applyAp :: Q Exp -> Q Exp -> Q Exp
applyAp qf qx = do
  f <- qf
  InfixE (Just f) (VarE '(<*>)) . Just <$> qx

----------------------------------------------------------------------------------------------------
-- Encoder generation
----------------------------------------------------------------------------------------------------

genEncoder :: String -> [ResolvedColumn] -> Q [Dec]
genEncoder typName cols = do
  let encName = mkName (camelCase typName <> "Encoder")
      tName = mkName typName
      sigTy = AppT (ConT ''Hasql.Encoders.Params) (ConT tName)
  sig <- sigD encName (return sigTy)
  let body = AppE (VarE 'mconcat) (ListE (map (columnEncodeExp tName) cols))
      dec = ValD (VarP encName) (NormalB body) []
  return [sig, dec]

{-  Generate a single encoder term:
    @contramap fieldSelector (Hasql.Encoders.param (Hasql.Encoders.nonNullable pgEncode))@

    Uses a record pattern match lambda to extract the field, which works
    regardless of @DuplicateRecordFields@ or @NoFieldSelectors@ settings.
-}
columnEncodeExp :: Name -> ResolvedColumn -> Exp
columnEncodeExp tName rc =
  let nullability =
        if rcNotNull rc
          then VarE 'Hasql.Encoders.nonNullable
          else VarE 'Hasql.Encoders.nullable
      codec = VarE 'Hasql.Generate.Codec.pgEncode
      paramEnc = AppE (VarE 'Hasql.Encoders.param) (AppE nullability codec)
      x = mkName "x"
      selector =
        LamE
          [RecP tName [(mkName (rcFieldName rc), VarP x)]]
          (VarE x)
   in AppE (AppE (VarE 'Contravariant.contramap) selector) paramEnc

{-  Generate a single array encoder term for batch operations:

    @contramap (map (\\TypeName{field = x} -> x))
      (Encoders.param (Encoders.nonNullable
        (Encoders.foldableArray (nonNullable\/nullable pgEncode))))@

    The outer @param@ is always @nonNullable@ (the array itself is always present).
    The inner @foldableArray@ element nullability follows 'rcNotNull'.
    The selector uses @map@ over a @RecP@ lambda pattern.
-}
columnArrayEncodeExp :: Name -> ResolvedColumn -> Exp
columnArrayEncodeExp tName rc =
  let elementNullability =
        if rcNotNull rc
          then VarE 'Hasql.Encoders.nonNullable
          else VarE 'Hasql.Encoders.nullable
      codec = VarE 'Hasql.Generate.Codec.pgEncode
      arrayEnc =
        AppE
          (VarE 'Hasql.Encoders.foldableArray)
          (AppE elementNullability codec)
      paramEnc =
        AppE
          (VarE 'Hasql.Encoders.param)
          (AppE (VarE 'Hasql.Encoders.nonNullable) arrayEnc)
      x = mkName "x"
      fieldExtractor =
        LamE
          [RecP tName [(mkName (rcFieldName rc), VarP x)]]
          (VarE x)
      mapSelector = AppE (VarE 'map) fieldExtractor
   in AppE (AppE (VarE 'Contravariant.contramap) mapSelector) paramEnc

{-  Generate an array encoder for a single PK value list:

    @Encoders.param (Encoders.nonNullable
      (Encoders.foldableArray (nonNullable\/nullable pgEncode)))@
-}
singlePkArrayEncoder :: ResolvedColumn -> Exp
singlePkArrayEncoder resCol =
  let elementNullability =
        if rcNotNull resCol
          then VarE 'Hasql.Encoders.nonNullable
          else VarE 'Hasql.Encoders.nullable
      codec = VarE 'Hasql.Generate.Codec.pgEncode
      arrayEnc =
        AppE
          (VarE 'Hasql.Encoders.foldableArray)
          (AppE elementNullability codec)
   in AppE
        (VarE 'Hasql.Encoders.param)
        (AppE (VarE 'Hasql.Encoders.nonNullable) arrayEnc)

{-  Generate an array encoder for batch PK operations over @[PkType]@ input.

    * @NoPrimaryKey@: @noParams@
    * @SinglePk@: 'singlePkArrayEncoder'
    * @CompositePk@ without newtypes: @mconcat@ of tuple-field array encoders
    * @CompositePk@ with newtypes: @mconcat@ of RecP-based array encoders
      over the PK record type
-}
pkArrayEncoder :: Maybe String -> PkInfo -> Q Exp
pkArrayEncoder (Just pkTypName) (CompositePk rcs) = do
  let tName = mkName pkTypName
  return (AppE (VarE 'mconcat) (ListE (map (columnArrayEncodeExp tName) rcs)))
pkArrayEncoder _ NoPrimaryKey = return (VarE 'Hasql.Encoders.noParams)
pkArrayEncoder _ (SinglePk rc) = return (singlePkArrayEncoder rc)
pkArrayEncoder _ (CompositePk rcs) = do
  let n = length rcs
      encoders =
        zipWith
          (flip (tupleFieldArrayEncoder n))
          rcs
          (enumFromTo 0 (n - 1))
  return (AppE (VarE 'mconcat) (ListE encoders))

{-  Generate an array encoder for one field of a tuple PK:

    @contramap (map (\\(_, x, _) -> x))
      (Encoders.param (Encoders.nonNullable (Encoders.foldableArray ...)))@
-}
tupleFieldArrayEncoder :: Int -> Int -> ResolvedColumn -> Exp
tupleFieldArrayEncoder tupleSize idx resCol =
  let elementNullability =
        if rcNotNull resCol
          then VarE 'Hasql.Encoders.nonNullable
          else VarE 'Hasql.Encoders.nullable
      codec = VarE 'Hasql.Generate.Codec.pgEncode
      arrayEnc =
        AppE
          (VarE 'Hasql.Encoders.foldableArray)
          (AppE elementNullability codec)
      paramEnc =
        AppE
          (VarE 'Hasql.Encoders.param)
          (AppE (VarE 'Hasql.Encoders.nonNullable) arrayEnc)
      accessor = tupleAccessor tupleSize idx
      mapSelector = AppE (VarE 'map) accessor
   in AppE (AppE (VarE 'Contravariant.contramap) mapSelector) paramEnc

----------------------------------------------------------------------------------------------------
-- CRUD statement generation
----------------------------------------------------------------------------------------------------

{-  Assemble a typed top-level @Statement@ binding from its parts: name, type
    signature, SQL literal, encoder expression, and decoder expression.
    Every statement generator delegates to this for the final assembly.
-}
genStatement :: Name -> Type -> Exp -> Exp -> Exp -> Q [Dec]
genStatement stmtName sigTy sql enc dec = do
  sig <- sigD stmtName (return sigTy)
  let body = applyStatement sql enc dec
      valDec = ValD (VarP stmtName) (NormalB body) []
  return [sig, valDec]

genSelectByPk
  :: String -> String -> String -> [ResolvedColumn] -> PkInfo -> Maybe String -> Q [Dec]
genSelectByPk schema table typName cols pkInfo pkTypOverride = do
  let stmtName = mkName ("select" <> typName)
      decName = mkName (camelCase typName <> "Decoder")
      pkType = pkParamType pkTypOverride pkInfo
      sql = selectByPkSql schema table cols pkInfo
      sigTy =
        AppT
          (AppT (ConT ''Hasql.Statement.Statement) pkType)
          (AppT (ConT ''Maybe) (ConT (mkName typName)))
  pkEnc <- pkEncoder pkTypOverride pkInfo
  genStatement stmtName sigTy sql pkEnc (AppE (VarE 'Hasql.Decoders.rowMaybe) (VarE decName))

{-  Compute the columns to include in INSERT. Each column name in the
    @defaultedNames@ list is validated to exist in the table and to have a
    database default (@atthasdef@); a compile-time error is raised otherwise.
    Validated columns are then excluded from the insert column list.
-}
computeInsertCols :: [String] -> [ResolvedColumn] -> Q [ResolvedColumn]
computeInsertCols [] allCols = return allCols
computeInsertCols defaultedNames allCols = do
  mapM_ (validateDefaulted allCols) defaultedNames
  return (filter (\rc -> rcColName rc `notElem` defaultedNames) allCols)

validateDefaulted :: [ResolvedColumn] -> String -> Q ()
validateDefaulted allCols name =
  case filter (\rc -> rcColName rc == name) allCols of
    [] ->
      fail
        ( "hasql-generate: withDefaultedCols: column '"
            <> name
            <> "' not found"
        )
    (rc : _) ->
      if rcHasDefault rc
        then return ()
        else
          fail
            ( "hasql-generate: withDefaultedCols: column '"
                <> name
                <> "' has no database default"
            )

{-  Generate the INSERT statement. @allCols@ is the full record column list
    (used for RETURNING and the decoder), @insertCols@ is the subset actually
    included in the INSERT column list and VALUES params. When PK columns have
    defaults, @insertCols@ excludes them.
-}
genInsert
  :: String -> String -> String -> [ResolvedColumn] -> [ResolvedColumn] -> Q [Dec]
genInsert schema table typName allCols insertCols =
  let stmtName = mkName ("insert" <> typName)
      tName = mkName typName
      decName = mkName (camelCase typName <> "Decoder")
      sql = insertSql schema table allCols insertCols
      sigTy = AppT (AppT (ConT ''Hasql.Statement.Statement) (ConT tName)) (ConT tName)
      enc = insertEncoder tName insertCols
   in genStatement stmtName sigTy sql enc (AppE (VarE 'Hasql.Decoders.singleRow) (VarE decName))

{-  Generate an inline encoder for the INSERT that only encodes the insert
    columns (which may be a subset of all columns when PK columns are excluded).
-}
insertEncoder :: Name -> [ResolvedColumn] -> Exp
insertEncoder tName cols =
  AppE (VarE 'mconcat) (ListE (map (columnEncodeExp tName) cols))

{-  The UPDATE uses the full record encoder, so SQL parameter positions match
    the record field order. SET covers non-PK columns and WHERE covers PK
    columns, each referencing the correct @$N@ from the encoder.
-}
genUpdate
  :: String -> String -> String -> [ResolvedColumn] -> PkInfo -> Q [Dec]
genUpdate schema table typName cols pkInfo =
  let stmtName = mkName ("update" <> typName)
      tName = mkName typName
      encName = mkName (camelCase typName <> "Encoder")
      decName = mkName (camelCase typName <> "Decoder")
      sql = updateSql schema table cols pkInfo
      sigTy =
        AppT
          (AppT (ConT ''Hasql.Statement.Statement) (ConT tName))
          (AppT (ConT ''Maybe) (ConT tName))
   in genStatement stmtName sigTy sql (VarE encName) (AppE (VarE 'Hasql.Decoders.rowMaybe) (VarE decName))

genDeleteByPk :: String -> String -> String -> PkInfo -> Maybe String -> Q [Dec]
genDeleteByPk schema table _typName pkInfo pkTypOverride = do
  let stmtName = mkName ("delete" <> pascalCase table)
      pkType = pkParamType pkTypOverride pkInfo
      sql = deleteSql schema table pkInfo
      sigTy = AppT (AppT (ConT ''Hasql.Statement.Statement) pkType) (TupleT 0)
  pkEnc <- pkEncoder pkTypOverride pkInfo
  genStatement stmtName sigTy sql pkEnc (VarE 'Hasql.Decoders.noResult)

{-  Generate @selectMany\<TypeName\> :: Statement [PkType] [TypeName]@.

    Uses 'selectManySql' and 'pkArrayEncoder' for the input encoding.
-}
genSelectMany
  :: String -> String -> String -> [ResolvedColumn] -> PkInfo -> Maybe String -> Q [Dec]
genSelectMany schema table typName cols pkInfo pkTypOverride = do
  let stmtName = mkName ("selectMany" <> typName)
      tName = mkName typName
      decName = mkName (camelCase typName <> "Decoder")
      pkType = pkParamType pkTypOverride pkInfo
      sql = selectManySql schema table cols pkInfo
      sigTy =
        AppT
          (AppT (ConT ''Hasql.Statement.Statement) (AppT ListT pkType))
          (AppT ListT (ConT tName))
  pkEnc <- pkArrayEncoder pkTypOverride pkInfo
  genStatement stmtName sigTy sql pkEnc (AppE (VarE 'Hasql.Decoders.rowList) (VarE decName))

{-  Generate @deleteMany\<TypeName\> :: Statement [PkType] ()@.

    Uses 'deleteManySql' and 'pkArrayEncoder'.
-}
genDeleteMany
  :: String -> String -> String -> PkInfo -> Maybe String -> Q [Dec]
genDeleteMany schema table _typName pkInfo pkTypOverride = do
  let stmtName = mkName ("deleteMany" <> pascalCase table)
      pkType = pkParamType pkTypOverride pkInfo
      sql = deleteManySql schema table pkInfo
      sigTy = AppT (AppT (ConT ''Hasql.Statement.Statement) (AppT ListT pkType)) (TupleT 0)
  pkEnc <- pkArrayEncoder pkTypOverride pkInfo
  genStatement stmtName sigTy sql pkEnc (VarE 'Hasql.Decoders.noResult)

{-  Generate @insertMany\<TypeName\> :: Statement [TypeName] [TypeName]@.

    Uses 'insertManySql' and @mconcat@ of 'columnArrayEncodeExp' for
    each insert column. Respects 'computeInsertCols' (excludes defaulted
    PK columns).
-}
genInsertMany
  :: String -> String -> String -> [ResolvedColumn] -> [ResolvedColumn] -> Q [Dec]
genInsertMany schema table typName allCols insertCols =
  let stmtName = mkName ("insertMany" <> typName)
      tName = mkName typName
      decName = mkName (camelCase typName <> "Decoder")
      sql = insertManySql schema table allCols insertCols
      sigTy =
        AppT
          (AppT (ConT ''Hasql.Statement.Statement) (AppT ListT (ConT tName)))
          (AppT ListT (ConT tName))
      enc = AppE (VarE 'mconcat) (ListE (map (columnArrayEncodeExp tName) insertCols))
   in genStatement stmtName sigTy sql enc (AppE (VarE 'Hasql.Decoders.rowList) (VarE decName))

{-  Generate @updateMany\<TypeName\> :: Statement [TypeName] [TypeName]@.

    Uses 'updateManySql' and @mconcat@ of 'columnArrayEncodeExp' for ALL
    columns (both PK and non-PK appear in the unnest subquery). Only generated
    when the table has a primary key.
-}
genUpdateMany
  :: String -> String -> String -> [ResolvedColumn] -> PkInfo -> Q [Dec]
genUpdateMany schema table typName cols pkInfo =
  let stmtName = mkName ("updateMany" <> typName)
      tName = mkName typName
      decName = mkName (camelCase typName <> "Decoder")
      sql = updateManySql schema table cols pkInfo
      sigTy =
        AppT
          (AppT (ConT ''Hasql.Statement.Statement) (AppT ListT (ConT tName)))
          (AppT ListT (ConT tName))
      enc = AppE (VarE 'mconcat) (ListE (map (columnArrayEncodeExp tName) cols))
   in genStatement stmtName sigTy sql enc (AppE (VarE 'Hasql.Decoders.rowList) (VarE decName))

----------------------------------------------------------------------------------------------------
-- Typeclass instance generation
----------------------------------------------------------------------------------------------------

-- | Build @\\param -> Hasql.Session.statement param stmtName@.
sessionLambda :: String -> Name -> Exp
sessionLambda paramName stmtName =
  LamE
    [VarP (mkName paramName)]
    ( AppE
        (AppE (VarE 'Hasql.Session.statement) (VarE (mkName paramName)))
        (VarE stmtName)
    )

{-  Generate an instance with two session-wrapper methods (single + batch) and
    no associated types. Used for 'HasInsert' and 'HasUpdate'.
-}
genSimpleCrudInstance :: Name -> Name -> Name -> String -> String -> Q [Dec]
genSimpleCrudInstance className singleMethod batchMethod prefix typName = do
  let tName = mkName typName
      singleDec = FunD singleMethod [Clause [] (NormalB (sessionLambda "x" (mkName (prefix <> typName)))) []]
      batchDec = FunD batchMethod [Clause [] (NormalB (sessionLambda "xs" (mkName (prefix <> "Many" <> typName)))) []]
  return [InstanceD Nothing [] (AppT (ConT className) (ConT tName)) [singleDec, batchDec]]

{-  Generate an instance with two session-wrapper methods (single + batch) and
    one associated type for the key. Used for 'HasSelect' and 'HasDelete'.
-}
genKeyCrudInstance :: Name -> Name -> Name -> Name -> String -> String -> PkInfo -> Maybe String -> Q [Dec]
genKeyCrudInstance className singleMethod batchMethod assocType prefix typName pkInfo pkTypOverride = do
  let tName = mkName typName
      keyType = pkParamType pkTypOverride pkInfo
      singleDec = FunD singleMethod [Clause [] (NormalB (sessionLambda "k" (mkName (prefix <> typName)))) []]
      batchDec = FunD batchMethod [Clause [] (NormalB (sessionLambda "ks" (mkName (prefix <> "Many" <> typName)))) []]
      keyTySynInst = TySynInstD (TySynEqn Nothing (AppT (ConT assocType) (ConT tName)) keyType)
  return [InstanceD Nothing [] (AppT (ConT className) (ConT tName)) [keyTySynInst, singleDec, batchDec]]

genHasInsertInstance :: String -> Q [Dec]
genHasInsertInstance =
  genSimpleCrudInstance ''Hasql.Generate.Class.HasInsert 'Hasql.Generate.Class.insert 'Hasql.Generate.Class.insertMany "insert"

genHasSelectInstance :: String -> PkInfo -> Maybe String -> Q [Dec]
genHasSelectInstance =
  genKeyCrudInstance ''Hasql.Generate.Class.HasSelect 'Hasql.Generate.Class.select 'Hasql.Generate.Class.selectMany ''Hasql.Generate.Class.SelectKey "select"

genHasUpdateInstance :: String -> Q [Dec]
genHasUpdateInstance =
  genSimpleCrudInstance ''Hasql.Generate.Class.HasUpdate 'Hasql.Generate.Class.update 'Hasql.Generate.Class.updateMany "update"

genHasDeleteInstance :: String -> PkInfo -> Maybe String -> Q [Dec]
genHasDeleteInstance =
  genKeyCrudInstance ''Hasql.Generate.Class.HasDelete 'Hasql.Generate.Class.delete 'Hasql.Generate.Class.deleteMany ''Hasql.Generate.Class.DeleteKey "delete"

----------------------------------------------------------------------------------------------------
-- View-specific generators
----------------------------------------------------------------------------------------------------

{-  Generate a @select<TypeName> :: Statement () [TypeName]@ that selects
    all rows from the view with no parameters.
-}
genSelectView
  :: String -> String -> String -> [ResolvedColumn] -> Q [Dec]
genSelectView schema table typName cols =
  let stmtName = mkName ("select" <> typName)
      tName = mkName typName
      decName = mkName (camelCase typName <> "Decoder")
      sql = selectViewSql schema table cols
      sigTy =
        AppT
          (AppT (ConT ''Hasql.Statement.Statement) (TupleT 0))
          (AppT ListT (ConT tName))
   in genStatement stmtName sigTy sql (VarE 'Hasql.Encoders.noParams) (AppE (VarE 'Hasql.Decoders.rowList) (VarE decName))

selectViewSql :: String -> String -> [ResolvedColumn] -> Exp
selectViewSql schema table cols =
  sqlLit
    ( "SELECT "
        <> columnList cols
        <> " FROM "
        <> qualifiedName schema table
    )

-- | Generate a @HasView@ instance for a view.
genHasViewInstance :: String -> Q [Dec]
genHasViewInstance typName = do
  let tName = mkName typName
      stmtName = mkName ("select" <> typName)
      body =
        AppE
          (AppE (VarE 'Hasql.Session.statement) (TupE []))
          (VarE stmtName)
      selectViewMethod = FunD 'Hasql.Generate.Class.selectView [Clause [] (NormalB body) []]
  return [InstanceD Nothing [] (AppT (ConT ''Hasql.Generate.Class.HasView) (ConT tName)) [selectViewMethod]]

----------------------------------------------------------------------------------------------------
-- SQL generation helpers
----------------------------------------------------------------------------------------------------

selectByPkSql :: String -> String -> [ResolvedColumn] -> PkInfo -> Exp
selectByPkSql schema table cols pkInfo =
  let pkWhere = pkWhereClause pkInfo 1
   in sqlLit
        ( "SELECT "
            <> columnList cols
            <> " FROM "
            <> qualifiedName schema table
            <> " WHERE "
            <> pkWhere
        )

insertSql :: String -> String -> [ResolvedColumn] -> [ResolvedColumn] -> Exp
insertSql schema table allCols insertCols =
  let insertColNames = columnList insertCols
      returnColNames = columnList allCols
      params = typedParamList 1 insertCols
   in sqlLit
        ( "INSERT INTO "
            <> qualifiedName schema table
            <> " ("
            <> insertColNames
            <> ")"
            <> " VALUES ("
            <> params
            <> ")"
            <> " RETURNING "
            <> returnColNames
        )

updateSql :: String -> String -> [ResolvedColumn] -> PkInfo -> Exp
updateSql schema table cols pkInfo =
  let pkNames = pkColumnNames pkInfo
      indexed = zip cols (enumFromTo 1 (length cols))
      nonPkSet = filter (\(rc, _) -> rcColName rc `notElem` pkNames) indexed
      pkWhere = filter (\(rc, _) -> rcColName rc `elem` pkNames) indexed
      setClauses =
        intercalate
          ", "
          (map (\(rc, i) -> quoteIdent (rcColName rc) <> " = " <> paramRef i rc) nonPkSet)
      whereClauses =
        intercalate
          " AND "
          (map (\(rc, i) -> quoteIdent (rcColName rc) <> " = " <> paramRef i rc) pkWhere)
   in sqlLit
        ( "UPDATE "
            <> qualifiedName schema table
            <> " SET "
            <> setClauses
            <> " WHERE "
            <> whereClauses
            <> " RETURNING "
            <> columnList cols
        )

deleteSql :: String -> String -> PkInfo -> Exp
deleteSql schema table pkInfo =
  let pkWhere = pkWhereClause pkInfo 1
   in sqlLit
        ( "DELETE FROM "
            <> qualifiedName schema table
            <> " WHERE "
            <> pkWhere
        )

----------------------------------------------------------------------------------------------------
-- Batch SQL generation helpers
----------------------------------------------------------------------------------------------------

{-  Generate SQL for batch select: @SELECT cols FROM s.t WHERE pk = ANY($1)@
    for single PK, or @WHERE (c1, c2) IN (SELECT unnest($1), unnest($2))@
    for composite PK.
-}
selectManySql :: String -> String -> [ResolvedColumn] -> PkInfo -> Exp
selectManySql schema table cols pkInfo =
  sqlLit
    ( "SELECT "
        <> columnList cols
        <> " FROM "
        <> qualifiedName schema table
        <> " WHERE "
        <> batchPkWhereClause pkInfo
    )

{-  Generate SQL for batch delete: same WHERE clause logic as 'selectManySql'.
-}
deleteManySql :: String -> String -> PkInfo -> Exp
deleteManySql schema table pkInfo =
  sqlLit
    ( "DELETE FROM "
        <> qualifiedName schema table
        <> " WHERE "
        <> batchPkWhereClause pkInfo
    )

{-  Generate SQL for batch insert using @unnest@-based array parameters:

    @INSERT INTO s.t (cols) SELECT * FROM unnest($1, $2, ...) RETURNING allCols@
-}
insertManySql :: String -> String -> [ResolvedColumn] -> [ResolvedColumn] -> Exp
insertManySql schema table allCols insertCols =
  let insertColNames = columnList insertCols
      returnColNames = columnList allCols
      params = typedArrayParamList 1 insertCols
   in sqlLit
        ( "INSERT INTO "
            <> qualifiedName schema table
            <> " ("
            <> insertColNames
            <> ") SELECT * FROM unnest("
            <> params
            <> ") RETURNING "
            <> returnColNames
        )

{-  Generate SQL for batch update using an @unnest@ subquery join:

    @UPDATE s.t SET c1 = d.c1, c2 = d.c2
    FROM (SELECT unnest($1) AS pk, unnest($2) AS c1, unnest($3) AS c2) d
    WHERE s.t.pk = d.pk
    RETURNING s.t.*@

    All columns (PK and non-PK) appear in the unnest subquery. SET covers
    non-PK columns. WHERE joins on PK columns.
-}
updateManySql :: String -> String -> [ResolvedColumn] -> PkInfo -> Exp
updateManySql schema table cols pkInfo =
  let pkNames = pkColumnNames pkInfo
      nonPkCols = filter (\rc -> rcColName rc `notElem` pkNames) cols
      setClauses =
        intercalate
          ", "
          (map (\rc -> quoteIdent (rcColName rc) <> " = d." <> quoteIdent (rcColName rc)) nonPkCols)
      -- All columns in the unnest subquery
      unnestParams =
        intercalate
          ", "
          ( zipWith
              (\rc i -> "unnest(" <> arrayParamRef i rc <> ") AS " <> quoteIdent (rcColName rc))
              cols
              (enumFromTo 1 (length cols))
          )
      -- WHERE join on PK columns
      pkJoin =
        intercalate
          " AND "
          ( map
              (\pkn -> qualifiedName schema table <> "." <> quoteIdent pkn <> " = d." <> quoteIdent pkn)
              pkNames
          )
   in sqlLit
        ( "UPDATE "
            <> qualifiedName schema table
            <> " SET "
            <> setClauses
            <> " FROM (SELECT "
            <> unnestParams
            <> ") d WHERE "
            <> pkJoin
            <> " RETURNING "
            <> qualifiedName schema table
            <> ".*"
        )

{-  Generate the WHERE clause for batch PK operations.

    Single PK: @pk = ANY($1)@ — uses @arrayParamRef@ for enum cast support.
    Composite PK: @(c1, c2) IN (SELECT unnest($1), unnest($2))@.
-}
batchPkWhereClause :: PkInfo -> String
batchPkWhereClause NoPrimaryKey = "true"
batchPkWhereClause (SinglePk rc) =
  quoteIdent (rcColName rc) <> " = ANY(" <> arrayParamRef 1 rc <> ")"
batchPkWhereClause (CompositePk rcs) =
  let pkTuple = "(" <> intercalate ", " (map (quoteIdent . rcColName) rcs) <> ")"
      unnests =
        intercalate
          ", "
          ( zipWith
              (\rc i -> "unnest(" <> arrayParamRef i rc <> ")")
              rcs
              (enumFromTo 1 (length rcs))
          )
   in pkTuple <> " IN (SELECT " <> unnests <> ")"

{-  Generate a comma-separated array parameter list with per-column type casts,
    using 'arrayParamRef' so enum types get @[]@ appended.
-}
typedArrayParamList :: Int -> [ResolvedColumn] -> String
typedArrayParamList start cols =
  intercalate ", " (zipWith arrayParamRef (enumFromTo start (start + length cols - 1)) cols)

textLit :: String -> Exp
textLit s = AppE (VarE 'Data.String.fromString) (LitE (StringL s))

sqlLit :: String -> Exp
sqlLit = textLit

applyStatement :: Exp -> Exp -> Exp -> Exp
applyStatement sql enc dec =
  foldl AppE (ConE 'Hasql.Statement.Statement) [sql, enc, dec, ConE 'True]

columnList :: [ResolvedColumn] -> String
columnList cols = intercalate ", " (map (quoteIdent . rcColName) cols)

{-  Format a single parameter reference, appending a @::type@ cast when the
    column has 'rcPgCast' set (e.g. for enum types).

    @paramRef 3 rc@  →  @\"$3\"@ or @\"$3::hg_test.user_role\"@
-}
paramRef :: Int -> ResolvedColumn -> String
paramRef i rc = case rcPgCast rc of
  Nothing   -> "$" <> show i
  Just cast -> "$" <> show i <> "::" <> cast

{-  Format an array parameter reference for batch operations. Appends @[]@ to
    the cast suffix for enum types so that PostgreSQL receives the correct
    array type annotation.

    @arrayParamRef 1 rcEnum@  produces @\"$1::hg_test.user_role[]\"@
    @arrayParamRef 2 rcPlain@ produces @\"$2\"@
-}
arrayParamRef :: Int -> ResolvedColumn -> String
arrayParamRef i rc = case rcPgCast rc of
  Nothing   -> "$" <> show i
  Just cast -> "$" <> show i <> "::" <> cast <> "[]"

{-  Generate a comma-separated parameter list with per-column type casts.

    @typedParamList 1 [rcPlain, rcEnum]@  →  @\"$1, $2::hg_test.user_role\"@
-}
typedParamList :: Int -> [ResolvedColumn] -> String
typedParamList start cols =
  intercalate ", " (zipWith paramRef (enumFromTo start (start + length cols - 1)) cols)

pkWhereClause :: PkInfo -> Int -> String
pkWhereClause NoPrimaryKey _ = "true"
pkWhereClause (SinglePk rc) startIdx =
  quoteIdent (rcColName rc) <> " = " <> paramRef startIdx rc
pkWhereClause (CompositePk rcs) startIdx =
  intercalate
    " AND "
    ( zipWith
        (\rc i -> quoteIdent (rcColName rc) <> " = " <> paramRef i rc)
        rcs
        (enumFromTo startIdx (startIdx + length rcs - 1))
    )

pkColumnNames :: PkInfo -> [String]
pkColumnNames NoPrimaryKey      = []
pkColumnNames (SinglePk rc)     = [rcColName rc]
pkColumnNames (CompositePk rcs) = map rcColName rcs

----------------------------------------------------------------------------------------------------
-- PK type and encoder helpers
----------------------------------------------------------------------------------------------------

pkParamType :: Maybe String -> PkInfo -> Type
pkParamType (Just pkTypName) (CompositePk _) = ConT (mkName pkTypName)
pkParamType _ NoPrimaryKey = TupleT 0
pkParamType _ (SinglePk rc) = fieldType rc
pkParamType _ (CompositePk rcs) =
  foldl AppT (TupleT (length rcs)) (map fieldType rcs)

pkEncoder :: Maybe String -> PkInfo -> Q Exp
pkEncoder (Just pkTypName) (CompositePk _) =
  return (VarE (mkName (camelCase pkTypName <> "Encoder")))
pkEncoder _ NoPrimaryKey = return (VarE 'Hasql.Encoders.noParams)
pkEncoder _ (SinglePk rc) = return (singleParamEncoder rc)
pkEncoder _ (CompositePk rcs) = return (compositeEncoder rcs)

singleParamEncoder :: ResolvedColumn -> Exp
singleParamEncoder rc =
  let nullability =
        if rcNotNull rc
          then VarE 'Hasql.Encoders.nonNullable
          else VarE 'Hasql.Encoders.nullable
      codec = VarE 'Hasql.Generate.Codec.pgEncode
   in AppE (VarE 'Hasql.Encoders.param) (AppE nullability codec)

compositeEncoder :: [ResolvedColumn] -> Exp
compositeEncoder rcs =
  let n = length rcs
      encoders =
        zipWith
          (flip (tupleFieldEncoder n))
          rcs
          (enumFromTo 0 (n - 1))
   in AppE (VarE 'mconcat) (ListE encoders)

tupleFieldEncoder :: Int -> Int -> ResolvedColumn -> Exp
tupleFieldEncoder tupleSize idx rc =
  let nullability =
        if rcNotNull rc
          then VarE 'Hasql.Encoders.nonNullable
          else VarE 'Hasql.Encoders.nullable
      codec = VarE 'Hasql.Generate.Codec.pgEncode
      paramEnc = AppE (VarE 'Hasql.Encoders.param) (AppE nullability codec)
      accessor = tupleAccessor tupleSize idx
   in AppE (AppE (VarE 'Contravariant.contramap) accessor) paramEnc

{-  Build a lambda extracting the Nth element from a tuple of known size.
    For pairs, uses @fst@ and @snd@. For larger tuples, generates a pattern match:
    @\\(_, _, x, _) -> x@
-}
tupleAccessor :: Int -> Int -> Exp
tupleAccessor 2 0 = VarE 'Data.Tuple.fst
tupleAccessor 2 1 = VarE 'Data.Tuple.snd
tupleAccessor size idx =
  let names = map (\i -> mkName ("t" <> show i)) (enumFromTo 0 (size - 1))
      pats = map VarP names
      target = names !! idx
   in LamE [TupP pats] (VarE target)

----------------------------------------------------------------------------------------------------
-- Newtype PK generation
----------------------------------------------------------------------------------------------------

{-  Generate a newtype wrapper for a single-column primary key.

    @genPkNewtype \"UsersPk\" [''Show, ''Eq] rc@  where @rc@ has type @UUID@

    produces:

    @newtype UsersPk = UsersPk { getUsersPk :: !UUID } deriving stock (Show, Eq)@
-}
genPkNewtype :: String -> [Name] -> ResolvedColumn -> Q [Dec]
genPkNewtype pkTypName derivNames rc = do
  let tName = mkName pkTypName
      unwrapperName = mkName ("get" <> pkTypName)
      fieldBang = Bang NoSourceUnpackedness NoSourceStrictness
      field = (unwrapperName, fieldBang, rcType rc)
      con = RecC tName [field]
      derivClauses = mkDerivClauses derivNames
  return [NewtypeD [] tName [] Nothing con derivClauses]

{-  Generate a @PgCodec@ instance for a single-column PK newtype.

    @genPkNewtypeCodec \"UsersPk\"@  produces:

    @
    instance PgCodec UsersPk where
      pgDecode = UsersPk \<$\> pgDecode
      pgEncode = contramap getUsersPk pgEncode
    @
-}
genPkNewtypeCodec :: String -> Q [Dec]
genPkNewtypeCodec pkTypName = do
  let tName = mkName pkTypName
      conName = mkName pkTypName
      unwrapperName = mkName ("get" <> pkTypName)
      -- pgDecode = Con <$> pgDecode
      decoderBody =
        InfixE
          (Just (ConE conName))
          (VarE '(<$>))
          (Just (VarE 'Hasql.Generate.Codec.pgDecode))
      decoderDec = FunD 'Hasql.Generate.Codec.pgDecode [Clause [] (NormalB decoderBody) []]
      -- pgEncode = contramap unwrapper pgEncode
      encoderBody =
        AppE
          (AppE (VarE 'Contravariant.contramap) (VarE unwrapperName))
          (VarE 'Hasql.Generate.Codec.pgEncode)
      encoderDec = FunD 'Hasql.Generate.Codec.pgEncode [Clause [] (NormalB encoderBody) []]
      instanceType = AppT (ConT ''Hasql.Generate.Codec.PgCodec) (ConT tName)
  return [InstanceD Nothing [] instanceType [decoderDec, encoderDec]]

{-  Generate a data record for a composite primary key.

    @genPkRecord \"CompositePkPk\" [''Show, ''Eq] pkRcs@  produces a record with
    fields matching the PK columns.
-}
genPkRecord :: String -> [Name] -> [ResolvedColumn] -> Q [Dec]
genPkRecord pkTypName derivNames pkRcs = do
  let tName = mkName pkTypName
      fields = map mkPkField pkRcs
      con = RecC tName fields
      derivClauses = mkDerivClauses derivNames
  return [DataD [] tName [] Nothing [con] derivClauses]
  where
    mkPkField :: ResolvedColumn -> VarBangType
    mkPkField rc =
      let fName = mkName (rcFieldName rc)
          fieldBang = Bang NoSourceUnpackedness SourceStrict
       in (fName, fieldBang, fieldType rc)

{-  Generate an encoder for a composite PK record type.

    @genPkRecordEncoder \"CompositePkPk\" pkRcs@  produces:

    @
    compositePkPkEncoder :: Encoders.Params CompositePkPk
    compositePkPkEncoder = mconcat [...]
    @
-}
genPkRecordEncoder :: String -> [ResolvedColumn] -> Q [Dec]
genPkRecordEncoder pkTypName pkRcs = do
  let encName = mkName (camelCase pkTypName <> "Encoder")
      tName = mkName pkTypName
      sigTy = AppT (ConT ''Hasql.Encoders.Params) (ConT tName)
  sig <- sigD encName (return sigTy)
  let body = AppE (VarE 'mconcat) (ListE (map (columnEncodeExp tName) pkRcs))
      dec = ValD (VarP encName) (NormalB body) []
  return [sig, dec]

{-  Compute the field name for a PK record column. Uses @rcColName@ (the raw PG
    column name) rather than the main record's potentially-prefixed @rcFieldName@.
-}
pkRecordFieldName :: String -> Bool -> ResolvedColumn -> ResolvedColumn
pkRecordFieldName pkTypName allowDupFields rc =
  if allowDupFields
    then rc {rcFieldName = sanitizeField (camelCase (rcColName rc))}
    else rc {rcFieldName = sanitizeField (camelCase pkTypName <> pascalCase (rcColName rc))}

----------------------------------------------------------------------------------------------------
-- HasPrimaryKey instance generation
----------------------------------------------------------------------------------------------------

{-  Generate a @HasPrimaryKey@ instance for a table.

    Handles all four cases: single/composite × newtypes/no-newtypes.
    @rawPk@ is generated explicitly since the type families are non-injective
    and GHC cannot resolve a default @unwrapPk . toPk@ composition.
-}
genHasPrimaryKeyInstance
  :: String -> PkInfo -> PkInfo -> Bool -> Maybe String -> Bool -> Q [Dec]
genHasPrimaryKeyInstance typName origPkInfo pkInfo' useNtPk pkTypOverride allowDupFields = do
  let tName = mkName typName
      pkOfType = pkParamType pkTypOverride pkInfo'
      rawPkOfType = pkParamType Nothing origPkInfo
      pkOfTySynInst =
        TySynInstD
          (TySynEqn Nothing (AppT (ConT ''Hasql.Generate.Class.PkOf) (ConT tName)) pkOfType)
      rawPkOfTySynInst =
        TySynInstD
          (TySynEqn Nothing (AppT (ConT ''Hasql.Generate.Class.RawPkOf) (ConT tName)) rawPkOfType)
  toPkBody <- genToPk typName pkInfo' pkTypOverride allowDupFields
  wrapPkBody <- genWrapPk origPkInfo pkInfo' useNtPk pkTypOverride allowDupFields
  unwrapPkBody <- genUnwrapPk origPkInfo pkInfo' useNtPk pkTypOverride allowDupFields
  rawPkBody <- genRawPk typName origPkInfo pkInfo' useNtPk pkTypOverride allowDupFields
  let toPkDec = FunD 'Hasql.Generate.Class.toPk [Clause [] (NormalB toPkBody) []]
      wrapPkDec = FunD 'Hasql.Generate.Class.wrapPk [Clause [] (NormalB wrapPkBody) []]
      unwrapPkDec = FunD 'Hasql.Generate.Class.unwrapPk [Clause [] (NormalB unwrapPkBody) []]
      rawPkDec = FunD 'Hasql.Generate.Class.rawPk [Clause [] (NormalB rawPkBody) []]
      instanceType = AppT (ConT ''Hasql.Generate.Class.HasPrimaryKey) (ConT tName)
  return
    [ InstanceD
        Nothing
        []
        instanceType
        [pkOfTySynInst, rawPkOfTySynInst, toPkDec, wrapPkDec, unwrapPkDec, rawPkDec]
    ]

{-  Generate the @toPk@ method body: extracts PK field(s) from the main record.

    Single PK: @\\Rec{field = x} -> x@
    Composite, no newtypes: @\\Rec{f1 = x0, f2 = x1} -> (x0, x1)@
    Composite, with newtypes: @\\Rec{f1 = x0, f2 = x1} -> PkRec{pf1 = x0, pf2 = x1}@
-}
genToPk :: String -> PkInfo -> Maybe String -> Bool -> Q Exp
genToPk typName pkInfo' pkTypOverride allowDupFields = do
  let tName = mkName typName
  case pkInfo' of
    NoPrimaryKey -> fail "hasql-generate: genToPk called with NoPrimaryKey"
    SinglePk rc -> do
      let x = mkName "x"
      return (LamE [RecP tName [(mkName (rcFieldName rc), VarP x)]] (VarE x))
    CompositePk rcs -> do
      let xs = map (\i -> mkName ("x" <> show i)) (enumFromTo 0 (length rcs - 1))
          pat = RecP tName (zipWith (\rc x -> (mkName (rcFieldName rc), VarP x)) rcs xs)
      case pkTypOverride of
        Just pkTypName -> do
          let pkCon = mkName pkTypName
              pkRcs = map (pkRecordFieldName pkTypName allowDupFields) rcs
              body = RecConE pkCon (zipWith (\prc x -> (mkName (rcFieldName prc), VarE x)) pkRcs xs)
          return (LamE [pat] body)
        Nothing -> do
          let body = TupE (map (Just . VarE) xs)
          return (LamE [pat] body)

{-  Generate the @wrapPk@ method body: converts raw PK value(s) to the PK type.

    No newtypes: @\\x -> x@ (identity)
    Single + newtypes: @ConE pkTypName@ (newtype constructor)
    Composite + newtypes: @\\(x0, x1) -> PkRec{pf1 = x0, pf2 = x1}@
-}
genWrapPk :: PkInfo -> PkInfo -> Bool -> Maybe String -> Bool -> Q Exp
genWrapPk _origPkInfo _pkInfo' useNtPk pkTypOverride allowDupFields
  | not useNtPk = return identityLam
  | otherwise = case (_origPkInfo, pkTypOverride) of
      (SinglePk _, _) -> do
        let pkTypName = case _pkInfo' of
              SinglePk rc -> case rcType rc of
                ConT n -> nameBase n
                _      -> ""
              _ -> ""
        return (ConE (mkName pkTypName))
      (CompositePk rcs, Just pkTypName) -> do
        let n = length rcs
            xs = map (\i -> mkName ("x" <> show i)) (enumFromTo 0 (n - 1))
            pat = TupP (map VarP xs)
            pkRcs = map (pkRecordFieldName pkTypName allowDupFields) rcs
            body = RecConE (mkName pkTypName) (zipWith (\prc x -> (mkName (rcFieldName prc), VarE x)) pkRcs xs)
        return (LamE [pat] body)
      _ -> return identityLam

{-  Generate the @unwrapPk@ method body: converts PK type to raw value(s).

    No newtypes: @\\x -> x@ (identity)
    Single + newtypes: @VarE getterName@ (newtype field accessor)
    Composite + newtypes: @\\PkRec{pf1 = x0, pf2 = x1} -> (x0, x1)@
-}
genUnwrapPk :: PkInfo -> PkInfo -> Bool -> Maybe String -> Bool -> Q Exp
genUnwrapPk _origPkInfo _pkInfo' useNtPk pkTypOverride allowDupFields
  | not useNtPk = return identityLam
  | otherwise = case (_origPkInfo, pkTypOverride) of
      (SinglePk _, _) -> do
        let pkTypName = case _pkInfo' of
              SinglePk rc -> case rcType rc of
                ConT n -> nameBase n
                _      -> ""
              _ -> ""
        return (VarE (mkName ("get" <> pkTypName)))
      (CompositePk rcs, Just pkTypName) -> do
        let n = length rcs
            xs = map (\i -> mkName ("x" <> show i)) (enumFromTo 0 (n - 1))
            pkRcs = map (pkRecordFieldName pkTypName allowDupFields) rcs
            pat = RecP (mkName pkTypName) (zipWith (\prc x -> (mkName (rcFieldName prc), VarP x)) pkRcs xs)
            body = TupE (map (Just . VarE) xs)
        return (LamE [pat] body)
      _ -> return identityLam

{-  Generate the @rawPk@ method body: extracts raw PK value(s) directly from the
    main record.

    Without newtypes: same as @toPk@ — the PK type is already the raw type.
    With newtypes, single PK: pattern matches through the newtype to extract
    the raw value (e.g. @\\Rec{id = NtSinglePk x} -> x@).
    With newtypes, composite PK: extracts raw fields into a tuple (the record
    fields are already raw types, only the PK wrapper is a record).
-}
genRawPk :: String -> PkInfo -> PkInfo -> Bool -> Maybe String -> Bool -> Q Exp
genRawPk typName origPkInfo pkInfo' useNtPk _pkTypOverride allowDupFields
  | not useNtPk = genToPk typName origPkInfo Nothing allowDupFields
  | otherwise = do
      let tName = mkName typName
      case origPkInfo of
        NoPrimaryKey -> fail "hasql-generate: genRawPk called with NoPrimaryKey"
        SinglePk _rc -> do
          -- The record field is the newtype; pattern match through it
          case pkInfo' of
            SinglePk rc' -> do
              let pkTypName = case rcType rc' of
                    ConT n -> nameBase n
                    _      -> ""
                  x = mkName "x"
                  innerPat = ConP (mkName pkTypName) [] [VarP x]
                  pat = RecP tName [(mkName (rcFieldName rc'), innerPat)]
              return (LamE [pat] (VarE x))
            _ -> fail "hasql-generate: genRawPk single PK mismatch"
        CompositePk rcs -> do
          -- Composite PK fields in the main record are raw types (not rewritten)
          let xs = map (\i -> mkName ("x" <> show i)) (enumFromTo 0 (length rcs - 1))
              -- Use pkInfo' field names since those are what the main record uses
              rcs' = case pkInfo' of
                CompositePk cs -> cs
                _              -> rcs
              pat = RecP tName (zipWith (\rc x -> (mkName (rcFieldName rc), VarP x)) rcs' xs)
              body = TupE (map (Just . VarE) xs)
          return (LamE [pat] body)

identityLam :: Exp
identityLam =
  let x = mkName "x"
   in LamE [VarP x] (VarE x)

----------------------------------------------------------------------------------------------------
-- Utility functions
----------------------------------------------------------------------------------------------------

{-  Wrap a single PostgreSQL identifier in double-quotes, escaping any
    internal double-quote characters per the SQL standard (@\"\"@ → @\"\"\"\"@).
-}
quoteIdent :: String -> String
quoteIdent s = "\"" <> concatMap escDQ s <> "\""
  where
    escDQ '"' = "\"\""
    escDQ c   = [c]

{-  Build a schema-qualified, double-quoted identifier pair.

    @qualifiedName \"public\" \"users\"@  →  @\"\\\"public\\\".\\\"users\\\"\"@
-}
qualifiedName :: String -> String -> String
qualifiedName schema name = quoteIdent schema <> "." <> quoteIdent name

{-  Prefix a resolved column's field name with the type name in camelCase.

    @prefixFieldName \"Users\" rc{rcColName=\"tenant_id\"}@  gives
    @rc{rcFieldName=\"usersTenantId\"}@
-}
prefixFieldName :: String -> ResolvedColumn -> ResolvedColumn
prefixFieldName typName rc =
  rc {rcFieldName = sanitizeField (camelCase typName <> pascalCase (rcColName rc))}

{-  Append an apostrophe to a camelCase identifier if it collides with a
    Haskell reserved keyword. This is idiomatic Haskell (e.g. @type'@, @class'@).
-}
sanitizeField :: String -> String
sanitizeField s
  | s `elem` haskellKeywords = s <> "'"
  | otherwise = s

-- | Haskell reserved keywords that may collide with generated camelCase identifiers.
haskellKeywords :: [String]
haskellKeywords =
  [ "case"
  , "class"
  , "data"
  , "default"
  , "deriving"
  , "do"
  , "else"
  , "forall"
  , "foreign"
  , "hiding"
  , "if"
  , "import"
  , "in"
  , "infix"
  , "infixl"
  , "infixr"
  , "instance"
  , "let"
  , "module"
  , "newtype"
  , "of"
  , "qualified"
  , "then"
  , "type"
  , "where"
  , "mdo"
  , "rec"
  , "proc"
  , "pattern"
  , "role"
  , "family"
  , "stock"
  , "anyclass"
  , "via"
  ]

{-  Convert a snake_case or plain identifier to PascalCase.

    @\"user_emails\"  →  \"UserEmails\"@
    @\"users\"        →  \"Users\"@
-}
pascalCase :: String -> String
pascalCase = concatMap titleWord . splitSnake

{-  Convert a snake_case or PascalCase identifier to camelCase.

    @\"tenant_id\"    →  \"tenantId\"@
    @\"UserEmails\"   →  \"userEmails\"@
    @\"name\"         →  \"name\"@
-}
camelCase :: String -> String
camelCase s = case splitSnake s of
  []       -> []
  (w : ws) -> lowerFirst w <> concatMap titleWord ws

titleWord :: String -> String
titleWord []       = []
titleWord (c : cs) = toUpper c : cs

lowerFirst :: String -> String
lowerFirst []       = []
lowerFirst (c : cs) = toLower c : cs

{-  Split a string on underscores.

    @\"tenant_id\"  →  [\"tenant\", \"id\"]@
    @\"name\"       →  [\"name\"]@
-}
splitSnake :: String -> [String]
splitSnake [] = []
splitSnake s = case break (== '_') s of
  (word, [])       -> [word]
  (word, _ : rest) -> word : splitSnake rest