Tweak.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE RecordWildCards #-}
-- |
-- Module      : Warlock.Tweak
-- Description : Generate new data types from existing ones with modifications
-- License     : MIT
--
-- = Overview
--
-- @Warlock.Tweak@ allows you to derive new data types from existing ones by
-- specifying field modifications (selection, renaming, prefix operations).
-- It automatically generates the new type and optionally creates 'Warlock.From'
-- instances between the source and derived types.
--
-- = Quick Example
--
-- @
-- {-# LANGUAGE TemplateHaskell #-}
-- import Warlock.Tweak
--
-- data User = User
--   { userId :: Int
--   , userName :: String
--   , userEmail :: String
--   , userPassword :: ByteString
--   }
--
-- -- Generate DTO without sensitive fields
-- tweakType
--   (DropFrom $ defaultTweakConfig
--     \`withoutFields\` ['userPassword]
--     \`stripPrefix\` "user")
--   "UserDTO"
--   ''User
--
-- -- Generates:
-- -- data UserDTO = UserDTO
-- --   { id :: Int
-- --   , name :: String
-- --   , email :: String
-- --   }
-- -- derive (ByName config) ''User ''UserDTO
-- @
--
-- = Use Cases
--
-- * **DTOs**: Create API response/request types from database entities
-- * **Projections**: Subset of fields for specific use cases
-- * **API Versioning**: Generate V1/V2 types with field evolution
-- * **Type Refinement**: Create more specific types by dropping optional fields
--
-- = Advanced Field Rules
--
-- Warlock.Tweak now supports all Warlock field rules including:
--
-- * 'combineFields' - Combine multiple source fields into one
-- * 'disassembleFields' - Split one source field into multiple
-- * 'mapField' - Apply custom transformations to fields
-- * 'virtualField' - Use HasField for dynamic field access
-- * 'computedField' - Compute field values from expressions
--
-- @
-- -- Combine firstName and lastName into fullName:
-- pick ''Person "PersonDTO" ['fullName, 'age] \`withRules\`
--   [ combineFields 'fullName $ do
--       first <- get 'firstName
--       last <- get 'lastName
--       pure [| $first ++ \" \" ++ $last |]
--   ]
--
-- -- Transform email to lowercase:
-- omit ''User "PublicUser" ['password] \`withRules\`
--   [ mapField 'email (\src _dst -> [| map toLower $(varE src) |])
--   ]
-- @
--
-- = Configuration
--
-- Configuration is composable using helper functions:
--
-- @
-- defaultTweakConfig
--   \`withFields\` ['field1, 'field2]        -- Keep only these
--   \`withRenames\` [('oldName, 'newName)]   -- Explicit renames
--   \`stripPrefix\` "user"                   -- Remove prefix from all fields
--   \`withAutoDerive\` False                 -- Disable automatic instance generation
-- @

module Warlock.Tweak
  ( -- * Type Derivation (TypeScript-style)
    pick
  , omit
  , pick'
  , omit'

    -- * Legacy API
  , tweakType
  , TweakStrategy(..)

    -- * Configuration
  , TweakConfig(..)
  , defaultTweakConfig
  , FieldSelector(..)
  , PrefixOp(..)

    -- * Configuration Helpers
  , withFields
  , withoutFields
  , withRenames
  , withPrefixOps
  , withAutoDerive
  , withRules
  , withDefaults
  , withNormalize
  , withRuleGens

    -- * Prefix Operations
  , stripPrefix
  , addPrefix
  , replacePrefix

    -- * Re-exports from Warlock (Field Rules)
  , Warlock.FieldRule
  , Warlock.rename
  , Warlock.defaultTo
  , Warlock.ignore
  , Warlock.mapField
  , Warlock.virtualField
  , Warlock.computedField
  , Warlock.combineFields
  , Warlock.disassembleFields
  , Warlock.DisassembledField
  , (Warlock..=)

    -- * Field Getters
  , Warlock.FieldGetter
  , Warlock.get
  , Warlock.getName
  , Warlock.DisassembleGetter
  , Warlock.getSource
  , Warlock.getSourceName

    -- * Name Normalizers
  , Warlock.snakeToCamel
  , Warlock.camelToSnake
  , Warlock.lowerFirst

    -- * Field Rule Generators
  , Warlock.datatypePrefixRules
  , Warlock.constructorPrefixRules
  , Warlock.customPrefixRules
  , Warlock.RuleGenContext(..)
  ) where

import Language.Haskell.TH
import Control.Monad (when, forM)
import Data.List (isPrefixOf)
import Data.Char (toLower, toUpper)
import qualified Warlock

--------------------------------------------------------------------------------
-- Core Types

-- | Strategy for deriving a new type from an existing one.
data TweakStrategy
  = KeepOnly TweakConfig   -- ^ Keep only specified fields
  | DropFrom TweakConfig   -- ^ Drop specified fields, keep rest
  | TransformWith TweakConfig  -- ^ Reserved for future type transformations

-- | Configuration for type derivation.
--
-- Build configurations using 'defaultTweakConfig' and helper functions:
--
-- @
-- defaultTweakConfig
--   \`withFields\` ['field1, 'field2]
--   \`stripPrefix\` "user"
-- @
data TweakConfig = TweakConfig
  { fieldSelector :: FieldSelector
    -- ^ Which fields to select/keep
  , warlockConfig :: Warlock.Config
    -- ^ Embedded Warlock configuration for field rules
  , autoDerive :: Bool
    -- ^ Whether to automatically generate From instances (default: True)
  }

-- | Selector for which fields to include.
data FieldSelector
  = AllFields
    -- ^ Select all fields (used with DropFrom)
  | SelectFields [Name]
    -- ^ Select only these specific fields (used with KeepOnly)
  deriving (Show, Eq)

-- | Prefix operations that can be applied to field names.
data PrefixOp
  = StripPrefix String
    -- ^ Remove prefix from field names
  | AddPrefix String
    -- ^ Add prefix to field names
  | ReplacePrefix String String
    -- ^ Replace one prefix with another
  deriving (Show, Eq)

--------------------------------------------------------------------------------
-- Default Configuration

-- | Default configuration: all fields, no renames, no prefix ops, auto-derive enabled.
defaultTweakConfig :: TweakConfig
defaultTweakConfig = TweakConfig
  { fieldSelector = AllFields
  , warlockConfig = Warlock.defaultConfig
  , autoDerive = True
  }

--------------------------------------------------------------------------------
-- Configuration Helpers

-- | Keep only the specified fields.
--
-- @
-- defaultTweakConfig \`withFields\` ['userId, 'userName]
-- @
withFields :: TweakConfig -> [Name] -> TweakConfig
withFields cfg fields = cfg { fieldSelector = SelectFields fields }

-- | Drop the specified fields (inverse of 'withFields').
--
-- @
-- defaultTweakConfig \`withoutFields\` ['userPassword, 'userCreatedAt]
-- @
withoutFields :: TweakConfig -> [Name] -> TweakConfig
withoutFields cfg fields = cfg { fieldSelector = SelectFields fields }

-- | Add explicit field renames.
--
-- @
-- defaultTweakConfig \`withRenames\` [('userId, 'id), ('userName, 'name)]
-- @
withRenames :: TweakConfig -> [(Name, Name)] -> TweakConfig
withRenames cfg renames = cfg { warlockConfig = Warlock.withRenames (warlockConfig cfg) renames }

-- | Set prefix operations (legacy support - now implemented via Warlock normalize).
--
-- @
-- defaultTweakConfig \`withPrefixOps\` [StripPrefix "user", AddPrefix "dto"]
-- @
withPrefixOps :: TweakConfig -> [PrefixOp] -> TweakConfig
withPrefixOps cfg ops =
  let oldNormalize = Warlock.normalize (warlockConfig cfg)
      newNormalize = composePrefixOps oldNormalize ops
  in cfg { warlockConfig = Warlock.withNormalize (warlockConfig cfg) newNormalize }
  where
    composePrefixOps :: (String -> String) -> [PrefixOp] -> (String -> String)
    composePrefixOps baseNorm ops' s = foldl applyOp (baseNorm s) ops'

    applyOp :: String -> PrefixOp -> String
    applyOp name (StripPrefix prefix) =
      if prefix `isPrefixOf` name
        then lowercaseFirst2 $ drop (length prefix) name
        else name
    applyOp name (AddPrefix prefix) = prefix ++ uppercaseFirst2 name
    applyOp name (ReplacePrefix old new) =
      if old `isPrefixOf` name
        then new ++ uppercaseFirst2 (drop (length old) name)
        else name

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

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

-- | Enable or disable automatic instance generation.
--
-- @
-- defaultTweakConfig \`withAutoDerive\` False
-- @
withAutoDerive :: TweakConfig -> Bool -> TweakConfig
withAutoDerive cfg enabled = cfg { autoDerive = enabled }

-- | Add custom field rules (combineFields, disassembleFields, mapField, etc.).
--
-- @
-- defaultTweakConfig \`withRules\`
--   [ combineFields 'fullName $ do
--       first <- get 'firstName
--       last <- get 'lastName
--       pure [| $first ++ " " ++ $last |]
--   , mapField 'email (\src _dst -> [| toLower <$> $(varE src) |])
--   ]
-- @
withRules :: TweakConfig -> [Warlock.FieldRule] -> TweakConfig
withRules cfg rules = cfg { warlockConfig = Warlock.withRules (warlockConfig cfg) rules }

-- | Add default values for new fields.
--
-- @
-- defaultTweakConfig \`withDefaults\` [('newField, [| \"default\" |])]
-- @
withDefaults :: TweakConfig -> [(Name, Q Exp)] -> TweakConfig
withDefaults cfg defaults = cfg { warlockConfig = Warlock.withDefaults (warlockConfig cfg) defaults }

-- | Set the field name normalization function.
--
-- @
-- defaultTweakConfig \`withNormalize\` snakeToCamel
-- @
withNormalize :: TweakConfig -> (String -> String) -> TweakConfig
withNormalize cfg f = cfg { warlockConfig = Warlock.withNormalize (warlockConfig cfg) f }

-- | Set field rule generators.
--
-- @
-- defaultTweakConfig \`withRuleGens\` [datatypePrefixRules]
-- @
withRuleGens :: TweakConfig -> [Warlock.RuleGenContext -> String -> Maybe Warlock.FieldRule] -> TweakConfig
withRuleGens cfg gens = cfg { warlockConfig = Warlock.withRuleGens (warlockConfig cfg) gens }

--------------------------------------------------------------------------------
-- Convenience Functions

-- | Strip a prefix from all field names.
--
-- @
-- defaultTweakConfig \`stripPrefix\` "user"  -- userName -> name
-- @
stripPrefix :: TweakConfig -> String -> TweakConfig
stripPrefix cfg prefix = withPrefixOps cfg [StripPrefix prefix]

-- | Add a prefix to all field names.
--
-- @
-- defaultTweakConfig \`addPrefix\` "dto"  -- name -> dtoName
-- @
addPrefix :: TweakConfig -> String -> TweakConfig
addPrefix cfg prefix = withPrefixOps cfg [AddPrefix prefix]

-- | Replace one prefix with another.
--
-- @
-- defaultTweakConfig \`replacePrefix\` ("user", "dto")  -- userName -> dtoName
-- @
replacePrefix :: TweakConfig -> (String, String) -> TweakConfig
replacePrefix cfg (old, new) = withPrefixOps cfg [ReplacePrefix old new]

--------------------------------------------------------------------------------
-- TypeScript-style API

-- | Pick specific fields from a type to create a new type.
--
-- Similar to TypeScript's @Pick<T, K>@ utility type.
--
-- @
-- data User = User
--   { userId :: Int
--   , userName :: String
--   , userEmail :: String
--   , userPassword :: String
--   }
--
-- -- Pick only id and name fields
-- pick ''User "UserSummary" ['userId, 'userName]
--
-- -- With additional config
-- pick ''User "UserDTO" ['userId, 'userName, 'userEmail]
--   \`stripPrefix\` "user"
-- @
pick
  :: Name          -- ^ Source type
  -> String        -- ^ New type name
  -> [Name]        -- ^ Fields to pick
  -> Q [Dec]
pick sourceType newTypeName fields =
  tweakType (KeepOnly $ defaultTweakConfig `withFields` fields) newTypeName sourceType

-- | Omit specific fields from a type to create a new type.
--
-- Similar to TypeScript's @Omit<T, K>@ utility type.
--
-- @
-- data User = User
--   { userId :: Int
--   , userName :: String
--   , userEmail :: String
--   , userPassword :: String
--   }
--
-- -- Omit sensitive fields
-- omit ''User "UserDTO" ['userPassword]
--
-- -- With additional config
-- omit ''User "UserResponse" ['userPassword, 'userCreatedAt]
--   \`addPrefix\` "dto"
-- @
omit
  :: Name          -- ^ Source type
  -> String        -- ^ New type name
  -> [Name]        -- ^ Fields to omit
  -> Q [Dec]
omit sourceType newTypeName fields =
  tweakType (DropFrom $ defaultTweakConfig `withoutFields` fields) newTypeName sourceType

-- | Configuration-first version of 'pick' for use with config helpers.
--
-- @
-- pick' (defaultTweakConfig \`stripPrefix\` "user") ''User "UserClean" ['userName, 'userEmail]
-- @
pick'
  :: TweakConfig   -- ^ Configuration
  -> Name          -- ^ Source type
  -> String        -- ^ New type name
  -> [Name]        -- ^ Fields to pick
  -> Q [Dec]
pick' config sourceType newTypeName fields =
  tweakType (KeepOnly $ config `withFields` fields) newTypeName sourceType

-- | Configuration-first version of 'omit' for use with config helpers.
--
-- @
-- omit' (defaultTweakConfig \`addPrefix\` "dto") ''User "UserDTO" ['userPassword]
-- @
omit'
  :: TweakConfig   -- ^ Configuration
  -> Name          -- ^ Source type
  -> String        -- ^ New type name
  -> [Name]        -- ^ Fields to omit
  -> Q [Dec]
omit' config sourceType newTypeName fields =
  tweakType (DropFrom $ config `withoutFields` fields) newTypeName sourceType

--------------------------------------------------------------------------------
-- Main API (Legacy)

-- | Generate a new data type from an existing one with modifications.
--
-- @
-- -- Drop sensitive fields
-- tweakType (DropFrom $ defaultTweakConfig \`withoutFields\` ['userPassword])
--           "UserDTO"
--           ''User
--
-- -- Keep only specific fields
-- tweakType (KeepOnly $ defaultTweakConfig \`withFields\` ['userId, 'userName])
--           "UserMinimal"
--           ''User
-- @
tweakType
  :: TweakStrategy  -- ^ How to select/modify fields
  -> String         -- ^ New type name
  -> Name          -- ^ Source type name
  -> Q [Dec]       -- ^ Generated: data type + optional From instances
tweakType strategy newTypeName sourceTypeName = do
  -- Extract source type info
  sourceInfo <- reify sourceTypeName
  (sourceConName, sourceFields) <- extractFieldsWithCon sourceInfo

  -- Get config from strategy
  let config = case strategy of
        KeepOnly cfg -> cfg
        DropFrom cfg -> cfg
        TransformWith cfg -> cfg

  -- Apply field selection
  selectedFields <- applyFieldSelection strategy sourceFields

  -- Apply field renames
  renamedFields <- applyFieldRenames config selectedFields

  -- Generate new data type
  let newTypeNameQ = mkName newTypeName
  newTypeDec <- generateDataType newTypeNameQ renamedFields

  -- Generate Warlock instances if enabled
  instanceDecs <- if autoDerive config
    then generateInstances config sourceTypeName sourceConName newTypeNameQ selectedFields renamedFields
    else pure []

  pure $ newTypeDec : instanceDecs

--------------------------------------------------------------------------------
-- Implementation: Extract Source Type Info

extractFieldsWithCon :: Info -> Q (Name, [(Name, Bang, Type)])
extractFieldsWithCon info = case info of
  TyConI (DataD _ _ _ _ [RecC conName fields] _) ->
    pure (conName, fields)
  TyConI (NewtypeD _ _ _ _ (RecC conName fields) _) ->
    pure (conName, fields)
  TyConI (DataD _ _ _ _ [NormalC conName _fields] _) ->
    fail $ "Warlock.Tweak does not support positional constructors. "
        ++ "Constructor " ++ nameBase conName ++ " has no named fields."
  TyConI (DataD _ _ _ _ cons _) | length cons > 1 ->
    fail $ "Warlock.Tweak does not support multi-constructor ADTs (yet). "
        ++ "Found " ++ show (length cons) ++ " constructors."
  TyConI (DataD _ _ _ _ [] _) ->
    fail "Warlock.Tweak: Source type has no constructors."
  _ -> fail $ "Warlock.Tweak: Cannot derive from " ++ show info

--------------------------------------------------------------------------------
-- Implementation: Field Selection

applyFieldSelection :: TweakStrategy -> [(Name, Bang, Type)] -> Q [(Name, Bang, Type)]
applyFieldSelection strategy sourceFields =
  case strategy of
    KeepOnly cfg -> case fieldSelector cfg of
      AllFields -> pure sourceFields
      SelectFields names -> do
        let fieldMap = [(fname, (fname, b, ty)) | (fname, b, ty) <- sourceFields]
        forM names $ \fieldName ->
          case lookup fieldName fieldMap of
            Just field -> pure field
            Nothing -> fail $ "Field '" ++ nameBase fieldName ++ "' not found in source type"

    DropFrom cfg -> case fieldSelector cfg of
      AllFields -> pure sourceFields
      SelectFields namesToDrop -> do
        let shouldKeep (fname, _, _) = fname `notElem` namesToDrop
        pure $ filter shouldKeep sourceFields

    TransformWith _ ->
      fail "TransformWith strategy not yet implemented"

--------------------------------------------------------------------------------
-- Implementation: Field Renaming

applyFieldRenames :: TweakConfig -> [(Name, Bang, Type)] -> Q [(Name, Bang, Type)]
applyFieldRenames TweakConfig{..} inputFields = do
  -- Apply field transformations:
  -- 1. First check for explicit renames in staticRules
  -- 2. Then apply the normalize function for prefix operations

  let staticRules = Warlock.staticRules warlockConfig
      normalize = Warlock.normalize warlockConfig

      -- Build a map of explicit renames from staticRules
      -- In Warlock, a rename rule stores dst (populated field) and src (source field)
      -- For Tweak, we use renames to transform field names, so we reverse the mapping:
      -- the rule's dst is the OLD name (what we're renaming FROM)
      -- the rule's src is the NEW name (what we're renaming TO)
      explicitRenameMap =
        [ (dstName, srcName)  -- dst is old name, src is new name
        | rule <- staticRules
        , let dstName = Warlock.dst rule
        , Just (Warlock.RealField srcName) <- [Warlock.src rule]
        ]

      transformedFields = map applyTransform inputFields

      applyTransform :: (Name, Bang, Type) -> (Name, Bang, Type)
      applyTransform (fname, b, ty) =
        -- First check for explicit rename
        case lookup fname explicitRenameMap of
          Just newName -> (newName, b, ty)
          Nothing ->
            -- Otherwise apply normalize function
            let normalized = normalize (nameBase fname)
            in (mkName normalized, b, ty)

  -- Check for name collisions
  checkNameCollisions transformedFields

  pure transformedFields
  where

    checkNameCollisions :: [(Name, Bang, Type)] -> Q ()
    checkNameCollisions flds = do
      let names = map (\(n, _, _) -> nameBase n) flds
      let duplicates = findDuplicates names
      when (not $ null duplicates) $
        fail $ "Warlock.Tweak: Field name collision after renaming: "
            ++ show duplicates

    findDuplicates :: Eq a => [a] -> [a]
    findDuplicates [] = []
    findDuplicates (x:xs)
      | x `elem` xs = x : findDuplicates xs
      | otherwise = findDuplicates xs

--------------------------------------------------------------------------------
-- Implementation: Generate Data Type

generateDataType :: Name -> [(Name, Bang, Type)] -> Q Dec
generateDataType typeName fields = do
  let conName = typeName  -- Constructor name same as type name
  let conDec = RecC conName fields
  pure $ DataD [] typeName [] Nothing [conDec] []

--------------------------------------------------------------------------------
-- Implementation: Generate Warlock Instances

generateInstances
  :: TweakConfig
  -> Name  -- Source type
  -> Name  -- Source constructor
  -> Name  -- New type
  -> [(Name, Bang, Type)]  -- Source fields (before renaming)
  -> [(Name, Bang, Type)]  -- New fields (after renaming)
  -> Q [Dec]
generateInstances TweakConfig{..} sourceType sourceConName newType sourceFields newFields = do
  -- For Tweak, we need to create explicit rename rules for each field since the
  -- destination fields may have been renamed (e.g., userId -> dtoUserId).
  -- We create a config with identity normalization and explicit field mappings.
  --
  -- sourceFields and newFields are parallel arrays (same length, same order):
  -- sourceFields are the selected fields from the source type (before renaming)
  -- newFields are those same fields after applying field transformations

  let existingRules = Warlock.staticRules warlockConfig
      -- Extract destination field names from existing rules
      existingDstFields = [Warlock.dst r | r <- existingRules]

      -- Only add positional renames for fields that don't have explicit rules
      automaticRules =
        [ Warlock.rename destName srcName
        | ((srcName, _, _), (destName, _, _)) <- zip sourceFields newFields
        , destName `notElem` existingDstFields  -- Don't override explicit rules
        ]

      -- Use the warlock config but add automatic renames and set normalize to identity
      -- (since we've already applied transformations to the field names)
      tweakInstanceConfig = warlockConfig
        { Warlock.staticRules = existingRules ++ automaticRules
        , Warlock.normalize = id
        }

  inst <- Warlock.resolveFieldRulesAndBuildInstance tweakInstanceConfig sourceType sourceConName sourceFields newType newType newFields

  pure [inst]