Generate.hs

module Valiant.CLI.Command.Generate
  ( runGenerate
  , GenerateOpts (..)
  ) where

import Data.Char (toUpper)
import Data.List (intercalate, nub, sort)
import Data.List qualified as List
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Valiant.CLI.Cache (CacheColumn (..), CacheEntry (..), CacheParam (..), StatementType (..), findCacheFile)
import Valiant.CLI.Config (AppEnv (..))
import Valiant.CLI.Discover (SqlFile (..), discoverSqlFiles)
import Valiant.CLI.Error (ValiantCliError (..), dieWithError)
import Valiant.CLI.Output
import Valiant.CLI.SqlMetadata (SqlMetadata (..), parseSqlMetadata)
import System.Directory (createDirectoryIfMissing, renameFile)
import System.Exit (ExitCode (..))
import System.FilePath (takeBaseName, takeDirectory, (</>))

data GenerateOpts = GenerateOpts
  { genModulePrefix :: String
  , genOutputDir :: FilePath
  }
  deriving stock (Show)

runGenerate :: AppEnv -> GenerateOpts -> IO ExitCode
runGenerate env opts = do
  files <- discoverSqlFiles (appSqlDir env)
  case files of
    [] -> dieWithError (ErrNoSqlFiles (appSqlDir env))
    _ -> pure ()

  -- Group SQL files by their parent directory → one module per directory
  let grouped = groupByDirectory files
  results <- mapM (generateModule env opts) (Map.toAscList grouped)
  let okCount = length (filter id results)
      total = length results
  if okCount == total
    then pure ExitSuccess
    else pure (ExitFailure 1)

-- | Group SQL files by their first directory component.
-- @users/find_by_id.sql@ → key @"users"@, @posts/insert.sql@ → key @"posts"@
-- Files at the root go under key @""@.
groupByDirectory :: [SqlFile] -> Map String [SqlFile]
groupByDirectory = List.foldl' go Map.empty
  where
    go acc sf =
      let dir = takeDirectory (sqlRelPath sf)
          key = if dir == "." then "" else dir
       in Map.insertWith (++) key [sf] acc

-- | Generate a single Haskell module for a group of SQL files.
generateModule :: AppEnv -> GenerateOpts -> (String, [SqlFile]) -> IO Bool
generateModule env opts (dirKey, files) = do
  -- Load cache entries for all files
  entries <- mapM (loadEntry env) files
  let pairs = [(sf, e) | (sf, Just e) <- zip files entries]
  if null pairs
    then do
      printLn $ "  Skipping " <> T.pack dirKey <> " (no cache entries)"
      pure False
    else do
      let moduleName = buildModuleName (genModulePrefix opts) dirKey
          modulePath = genOutputDir opts </> moduleNameToPath moduleName
          content = renderModule env moduleName pairs
      createDirectoryIfMissing True (takeDirectory modulePath)
      let tmpPath = modulePath <> ".tmp"
      writeFile tmpPath content
      renameFile tmpPath modulePath
      printLn $
        "  Generated "
          <> T.pack modulePath
          <> " ("
          <> T.pack (show (length pairs))
          <> " queries)"
      pure True

loadEntry :: AppEnv -> SqlFile -> IO (Maybe CacheEntry)
loadEntry env sf = findCacheFile (appCacheDir env) (sqlRelPath sf) (sqlHash sf)

-- Module name construction ------------------------------------------------

buildModuleName :: String -> String -> String
buildModuleName prefix dirKey
  | null dirKey = prefix
  | otherwise = prefix <> "." <> toPascalCase dirKey

toPascalCase :: String -> String
toPascalCase = concatMap capitalize . splitOn '/'
  where
    capitalize [] = []
    capitalize (c : cs) = toUpper c : cs

    splitOn _ [] = []
    splitOn sep s =
      let (w, rest) = break (== sep) s
       in w : case rest of
            [] -> []
            (_ : rs) -> splitOn sep rs

moduleNameToPath :: String -> FilePath
moduleNameToPath = (<> ".hs") . map (\c -> if c == '.' then '/' else c)

-- Rendering ---------------------------------------------------------------

renderModule :: AppEnv -> String -> [(SqlFile, CacheEntry)] -> String
renderModule env moduleName pairs =
  unlines $
    [ "-- AUTO-GENERATED by valiant generate. Edit as needed."
    , "-- Re-run: valiant generate --module-prefix " <> takeWhile (/= '.') moduleName <> " --output-dir ..."
    , "{-# OPTIONS_GHC -fplugin=Valiant.Plugin"
    , "                -fplugin-opt=Valiant.Plugin:sql-dir=" <> appSqlDir env
    , "                -fplugin-opt=Valiant.Plugin:cache-dir=" <> appCacheDir env <> " #-}"
    , "module " <> moduleName <> " where"
    , ""
    , "import Valiant"
    ]
      <> renderImports pairs
      <> [""]
      <> concatMap renderBinding pairs

renderImports :: [(SqlFile, CacheEntry)] -> [String]
renderImports pairs =
  let -- Group by module and collect types
      grouped = Map.toAscList $ List.foldl' groupImport Map.empty (concatMap collectImports pairs)
   in map renderImportLine grouped

collectImports :: (SqlFile, CacheEntry) -> [(String, String)]
collectImports (_, entry) =
  [(T.unpack (cpHaskellModule p), baseType (T.unpack (cpHaskellType p))) | p <- ceParams entry]
    <> [(T.unpack (ccHaskellModule c), baseType (T.unpack (ccHaskellType c))) | c <- ceColumns entry]

-- Strip "Maybe " prefix to get the base type for import
baseType :: String -> String
baseType ('M' : 'a' : 'y' : 'b' : 'e' : ' ' : rest) = rest
baseType t = t

groupImport :: Map String [String] -> (String, String) -> Map String [String]
groupImport acc (modName, typeName)
  | modName == "Prelude" = acc
  | otherwise = Map.insertWith (\new old -> nub (old ++ new)) modName [typeName] acc

renderImportLine :: (String, [String]) -> String
renderImportLine (modName, types) =
  "import " <> modName <> " (" <> intercalate ", " (sort (nub types)) <> ")"

renderBinding :: (SqlFile, CacheEntry) -> [String]
renderBinding (sf, entry) =
  let meta = parseSqlMetadata (TE.decodeUtf8 (sqlContent sf))
      bindName = case smName meta of
        Just name -> T.unpack name
        Nothing -> sqlFileToBindingName (sqlRelPath sf)
      paramType = formatParamType (ceParams entry)
      resultType = case smResult meta of
        Just typeName ->
          let wrap = resultWrapper (sqlRelPath sf) (smSingle meta) (ceStatementType entry) (ceColumns entry)
           in wrap (T.unpack typeName)
        Nothing -> formatResultType (sqlRelPath sf) (smSingle meta) (ceStatementType entry) (ceColumns entry)
      queryFn = case smResult meta of
        Just typeName -> "queryFileAs @" <> T.unpack typeName
        Nothing -> "queryFile"
      sqlComment = T.unpack (T.takeWhile (/= '\n') (ceSql entry))
   in [ ""
      , "-- " <> sqlRelPath sf
      , "-- " <> sqlComment
      , bindName <> " :: Statement " <> paramType <> " " <> resultType
      , bindName <> " = " <> queryFn <> " " <> show (sqlRelPath sf)
      ]

-- | Convert @users/find_by_id.sql@ → @findById@
sqlFileToBindingName :: FilePath -> String
sqlFileToBindingName path = snakeToCamel (takeBaseName path)

snakeToCamel :: String -> String
snakeToCamel [] = []
snakeToCamel s =
  let (word, rest) = break (== '_') s
   in word <> case rest of
        [] -> []
        (_ : rs) -> capitalize rs
  where
    capitalize [] = []
    capitalize (c : cs) =
      let (w, rest') = break (== '_') (c : cs)
       in case w of
            (x : xs) -> (toUpper x : xs) <> case rest' of
              [] -> []
              (_ : rs) -> capitalize rs
            [] -> case rest' of
              [] -> []
              (_ : rs) -> capitalize rs

formatParamType :: [CacheParam] -> String
formatParamType [] = "()"
formatParamType [p] = parensIfNeeded (T.unpack (cpHaskellType p))
formatParamType ps = "(" <> intercalate ", " (map (T.unpack . cpHaskellType) ps) <> ")"

formatResultType :: FilePath -> Bool -> StatementType -> [CacheColumn] -> String
formatResultType sqlPath single stmtType cols = case cols of
  [] -> "()"
  _ ->
    let inner = case cols of
          [c] -> parensIfNeeded (T.unpack (ccHaskellType c))
          cs -> "(" <> intercalate ", " (map (T.unpack . ccHaskellType) cs) <> ")"
     in if single
          then inner
          else inferReturnWrapper sqlPath stmtType cols inner

-- | Determine the wrapper for a named result type based on SQL conventions and metadata.
resultWrapper :: FilePath -> Bool -> StatementType -> [CacheColumn] -> String -> String
resultWrapper sqlPath single stmtType cols typeName
  | single = typeName
  | otherwise = inferReturnWrapper sqlPath stmtType cols typeName

-- | Infer the return type wrapper from the SQL file name prefix.
-- The naming convention follows the roadmap:
--
--   * @find_*@ → @Maybe r@ (zero or one row)
--   * @get_*@  → @r@ (exactly one row, no wrapper)
--   * @list_*@ → @[r]@ (multiple rows)
--   * @count_*@ → scalar type directly (no tuple, no wrapper)
--   * @exists_*@ → @Bool@
--   * @insert*@/@update_*@/@delete_*@/@upsert_*@ → @()@ when no columns
--   * default  → no wrapper (raw type)
inferReturnWrapper :: FilePath -> StatementType -> [CacheColumn] -> String -> String
inferReturnWrapper sqlPath _stmtType cols inner
  | "find_" `isPrefixOfBase` base = "(Maybe " <> inner <> ")"
  | "list_" `isPrefixOfBase` base = "[" <> inner <> "]"
  | "count_" `isPrefixOfBase` base = case cols of
      [c] -> parensIfNeeded (T.unpack (ccHaskellType c))
      _ -> inner
  | "exists_" `isPrefixOfBase` base = "Bool"
  | "get_" `isPrefixOfBase` base = inner
  | isCommandPrefix base, null cols = "()"
  | isCommandPrefix base = inner
  | otherwise = inner
  where
    base = takeBaseName sqlPath

    isPrefixOfBase :: String -> String -> Bool
    isPrefixOfBase prefix str = take (length prefix) str == prefix

    isCommandPrefix :: String -> Bool
    isCommandPrefix b =
      "insert" `isPrefixOfBase` b
        || "update_" `isPrefixOfBase` b
        || "delete_" `isPrefixOfBase` b
        || "upsert_" `isPrefixOfBase` b

-- | Wrap in parens if the type contains spaces (e.g. "Maybe Text" → "(Maybe Text)")
parensIfNeeded :: String -> String
parensIfNeeded s
  | ' ' `elem` s = "(" <> s <> ")"
  | otherwise = s