HKD.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Warlock.HKD
-- Description : Higher-Kinded Data (HKD) type generation
-- License     : MIT
--
-- = Overview
--
-- Generate HKD versions of your types for use with @barbies@ and other
-- HKD libraries. Automatically creates 'Witch.From' instances for conversion
-- between regular types and their Identity-wrapped HKD equivalents.
--
-- = Quick Example
--
-- @
-- {-# LANGUAGE TemplateHaskell #-}
-- {-# LANGUAGE TypeFamilies #-}
-- import Warlock.HKD
-- import Data.Functor.Identity
-- import Witch (from)
--
-- data Person = Person
--   { name :: String
--   , age :: Int
--   }
--
-- deriveHKD' ''Person
--
-- -- Generates:
-- -- type instance HKD Person f = Person' f
-- -- data Person' f = Person' { name :: f String, age :: f Int }
-- -- instance From (HKD Person Identity) Person
-- -- instance From Person (HKD Person Identity)
--
-- -- Usage:
-- let person = Person "Alice" 30
-- let hkdPerson = from person :: HKD Person Identity
-- let unwrapped = from hkdPerson :: Person
--
-- -- Usage with barbies:
-- import qualified Barbies as B
-- validatePerson :: Person' Maybe -> Either [String] Person
-- validatePerson p = B.btraverse (maybe (Left ["missing"]) Right) p
-- @
--
-- = Use Cases
--
-- * **Validation**: Wrap fields in Maybe for partial construction
-- * **Form handling**: Build up data incrementally
-- * **Serialization**: Handle optional fields gracefully
-- * **Testing**: Generate test data with default values
-- * **Barbies integration**: Use with btraverse, bpure, etc.

module Warlock.HKD
  ( -- * Type Family
    HKD

    -- * Template Haskell Generation
  , deriveHKD
  , deriveHKD'

    -- * Configuration
  , HKDConfig(..)
  , defaultHKDConfig
  , withFieldTransform
  , withConstructorTransform
  , withFieldPrefix
  , withConstructorSuffix
  , withoutFromInstances
  ) where

import Language.Haskell.TH
import qualified Data.Kind as Kind
import Data.Functor.Identity (Identity(..))
import qualified Witch

--------------------------------------------------------------------------------
-- Type Family

-- | Higher-Kinded Data type family.
--
-- Transforms a regular type into its HKD equivalent where each field
-- is wrapped in a type constructor @f@.
--
-- @
-- type instance HKD Person f = Person' f
-- @
type HKD :: Kind.Type -> (Kind.Type -> Kind.Type) -> Kind.Type
type family HKD a f

--------------------------------------------------------------------------------
-- Configuration

-- | Configuration for HKD generation.
--
-- Controls how the HKD type and its components are named, and whether
-- to generate From instances.
data HKDConfig = HKDConfig
  { hkdFieldNameTransform :: String -> String
    -- ^ Transform field names. Default: id (keep original names)
  , hkdConstructorNameTransform :: String -> String
    -- ^ Transform constructor names. Default: (++ "'") (add prime suffix)
  , hkdGenerateFromInstances :: Bool
    -- ^ Auto-generate From instances. Default: True
  }

-- | Default HKD configuration.
--
-- - Field names: unchanged
-- - Constructor names: add prime suffix (@Person@ → @Person'@)
-- - Generate From instances: yes
defaultHKDConfig :: HKDConfig
defaultHKDConfig = HKDConfig
  { hkdFieldNameTransform = id
  , hkdConstructorNameTransform = (++ "'")
  , hkdGenerateFromInstances = True
  }

-- | Set custom field name transform function.
--
-- @
-- deriveHKD (defaultHKDConfig \`withFieldTransform\` (\"hkd\" ++)) ''Person
-- @
withFieldTransform :: (String -> String) -> HKDConfig -> HKDConfig
withFieldTransform f cfg = cfg { hkdFieldNameTransform = f }

-- | Set custom constructor name transform function.
--
-- @
-- deriveHKD (defaultHKDConfig \`withConstructorTransform\` (++ "HKD")) ''Person
-- @
withConstructorTransform :: (String -> String) -> HKDConfig -> HKDConfig
withConstructorTransform f cfg = cfg { hkdConstructorNameTransform = f }

-- | Add prefix to field names.
--
-- @
-- deriveHKD (defaultHKDConfig \`withFieldPrefix\` "hkd") ''Person
-- -- name → hkdname
-- @
withFieldPrefix :: HKDConfig -> String -> HKDConfig
withFieldPrefix cfg prefix = withFieldTransform (prefix ++) cfg

-- | Add suffix to constructor names.
--
-- @
-- deriveHKD (defaultHKDConfig \`withConstructorSuffix\` "HKD") ''Person
-- -- Person → PersonHKD
-- @
withConstructorSuffix :: HKDConfig -> String -> HKDConfig
withConstructorSuffix cfg suffix = withConstructorTransform (++ suffix) cfg

-- | Disable automatic From instance generation.
withoutFromInstances :: HKDConfig -> HKDConfig
withoutFromInstances cfg = cfg { hkdGenerateFromInstances = False }

--------------------------------------------------------------------------------
-- Template Haskell Generation

-- | Generate HKD type with custom configuration.
--
-- @
-- deriveHKD defaultHKDConfig ''Person
-- deriveHKD (defaultHKDConfig \`withFieldPrefix\` "hkd") ''Person
-- @
deriveHKD :: HKDConfig -> Name -> Q [Dec]
deriveHKD config typeName = do
  -- Extract type information
  info <- reify typeName
  (typeVars, constructors) <- extractTypeInfo info

  -- Generate HKD data type
  let hkdTypeName = mkName $ hkdConstructorNameTransform config (nameBase typeName)
  hkdDataDecl <- generateHKDDataType config hkdTypeName typeVars constructors

  -- Generate type family instance
  let fVar = mkName "f"
  let typeFamilyInst = generateTypeFamilyInstance typeName hkdTypeName typeVars fVar

  -- Generate From instances if enabled
  fromInstances <- if hkdGenerateFromInstances config
    then generateFromInstances config typeName hkdTypeName typeVars constructors
    else pure []

  pure $ hkdDataDecl : typeFamilyInst : fromInstances

-- | Generate HKD type with default configuration.
--
-- @
-- deriveHKD' ''Person
-- @
deriveHKD' :: Name -> Q [Dec]
deriveHKD' = deriveHKD defaultHKDConfig

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

extractTypeInfo :: Info -> Q ([TyVarBndr BndrVis], [Con])
extractTypeInfo info = case info of
  TyConI (DataD _ _ typeVars _ cons _) ->
    pure (typeVars, cons)
  TyConI (NewtypeD _ _ typeVars _ con _) ->
    pure (typeVars, [con])
  _ -> fail $ "Warlock.HKD: Cannot derive HKD for " ++ show info

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

generateHKDDataType :: HKDConfig -> Name -> [TyVarBndr BndrVis] -> [Con] -> Q Dec
generateHKDDataType config hkdTypeName typeVars constructors = do
  let fVar = mkName "f"
  let fTyVar = PlainTV fVar BndrReq

  -- Add 'f' as the last type variable
  let allTyVars = typeVars ++ [fTyVar]

  -- Transform constructors
  hkdCons <- mapM (transformConstructor config fVar) constructors

  pure $ DataD [] hkdTypeName allTyVars Nothing hkdCons []

transformConstructor :: HKDConfig -> Name -> Con -> Q Con
transformConstructor config fVar con = case con of
  RecC conName fields -> do
    hkdFields <- mapM (transformField config fVar) fields
    let hkdConName = mkName $ hkdConstructorNameTransform config (nameBase conName)
    pure $ RecC hkdConName hkdFields

  NormalC conName fields -> do
    hkdFields <- mapM (transformBangType config fVar) fields
    let hkdConName = mkName $ hkdConstructorNameTransform config (nameBase conName)
    pure $ NormalC hkdConName hkdFields

  InfixC left conName right -> do
    hkdLeft <- transformBangType config fVar left
    hkdRight <- transformBangType config fVar right
    let hkdConName = mkName $ hkdConstructorNameTransform config (nameBase conName)
    pure $ InfixC hkdLeft hkdConName hkdRight

  _ -> fail "Warlock.HKD: GADT constructors not supported"

transformField :: HKDConfig -> Name -> VarBangType -> Q VarBangType
transformField config fVar (fieldName, bndr, fieldType) = do
  let hkdFieldName = mkName $ hkdFieldNameTransform config (nameBase fieldName)
  let hkdFieldType = AppT (VarT fVar) fieldType
  pure (hkdFieldName, bndr, hkdFieldType)

transformBangType :: HKDConfig -> Name -> BangType -> Q BangType
transformBangType _config fVar (bndr, fieldType) = do
  let hkdFieldType = AppT (VarT fVar) fieldType
  pure (bndr, hkdFieldType)

--------------------------------------------------------------------------------
-- Implementation: Generate Type Family Instance

generateTypeFamilyInstance :: Name -> Name -> [TyVarBndr BndrVis] -> Name -> Dec
generateTypeFamilyInstance origTypeName hkdTypeName typeVars fVar =
  let -- Build the original type applied to its type variables
      origTypeVarNames = map getTyVarName typeVars
      origType = foldl AppT (ConT origTypeName) (map VarT origTypeVarNames)

      -- Build the HKD type applied to type variables + f
      hkdType = foldl AppT (ConT hkdTypeName) (map VarT (origTypeVarNames ++ [fVar]))

      -- type instance HKD OrigType f = HKDType' f
      lhs = AppT (AppT (ConT ''HKD) origType) (VarT fVar)
  in TySynInstD (TySynEqn Nothing lhs hkdType)
  where
    getTyVarName :: TyVarBndr BndrVis -> Name
    getTyVarName (PlainTV n _) = n
    getTyVarName (KindedTV n _ _) = n

--------------------------------------------------------------------------------
-- Implementation: Generate From Instances

generateFromInstances :: HKDConfig -> Name -> Name -> [TyVarBndr BndrVis] -> [Con] -> Q [Dec]
generateFromInstances config origTypeName hkdTypeName typeVars constructors = do
  -- Generate: From (HKD OrigType Identity) OrigType
  unwrapInst <- generateUnwrapInstance config origTypeName hkdTypeName typeVars constructors

  -- Generate: From OrigType (HKD OrigType Identity)
  wrapInst <- generateWrapInstance config origTypeName hkdTypeName typeVars constructors

  pure [unwrapInst, wrapInst]

generateUnwrapInstance :: HKDConfig -> Name -> Name -> [TyVarBndr BndrVis] -> [Con] -> Q Dec
generateUnwrapInstance config origTypeName hkdTypeName typeVars constructors = do
  let typeVarNames = map getTyVarName typeVars
  let origType = foldl AppT (ConT origTypeName) (map VarT typeVarNames)
  -- Use concrete HKD type instead of type family for instance head
  -- Apply Identity as a type constructor, not a type variable
  let hkdType = foldl AppT (ConT hkdTypeName) (map VarT typeVarNames ++ [ConT ''Identity])

  -- Generate from clauses for each constructor
  clauses <- mapM (generateUnwrapClause config) constructors

  let fromMethod = FunD (mkName "from") clauses

  pure $ InstanceD Nothing [] (AppT (AppT (ConT ''Witch.From) hkdType) origType) [fromMethod]
  where
    getTyVarName (PlainTV n _) = n
    getTyVarName (KindedTV n _ _) = n

generateUnwrapClause :: HKDConfig -> Con -> Q Clause
generateUnwrapClause config con = case con of
  RecC conName fields -> do
    -- Generate pattern: HKDCon' (Identity f1) (Identity f2) ...
    fieldVars <- mapM (\_ -> newName "x") fields
    let hkdConName = mkName $ hkdConstructorNameTransform config (nameBase conName)
    -- Use transformed field names for pattern matching
    let transformedFields = [(mkName $ hkdFieldNameTransform config (nameBase fname), var) | ((fname, _, _), var) <- zip fields fieldVars]
    let hkdPattern = RecP hkdConName [(fname, ConP 'Identity [] [VarP var]) | (fname, var) <- transformedFields]

    -- Generate expression: OrigCon f1 f2 ...
    let origConExp = foldl AppE (ConE conName) (map VarE fieldVars)

    pure $ Clause [hkdPattern] (NormalB origConExp) []

  NormalC conName fields -> do
    fieldVars <- mapM (\_ -> newName "x") fields
    let hkdConName = mkName $ hkdConstructorNameTransform config (nameBase conName)
    let hkdPattern = ConP hkdConName [] [ConP 'Identity [] [VarP var] | var <- fieldVars]
    let origConExp = foldl AppE (ConE conName) (map VarE fieldVars)

    pure $ Clause [hkdPattern] (NormalB origConExp) []

  InfixC _ conName _ -> do
    leftVar <- newName "l"
    rightVar <- newName "r"
    let hkdConName = mkName $ hkdConstructorNameTransform config (nameBase conName)
    let hkdPattern = InfixP (ConP 'Identity [] [VarP leftVar]) hkdConName (ConP 'Identity [] [VarP rightVar])
    let origConExp = InfixE (Just (VarE leftVar)) (ConE conName) (Just (VarE rightVar))

    pure $ Clause [hkdPattern] (NormalB origConExp) []

  _ -> fail "Warlock.HKD: GADT constructors not supported in From instances"

generateWrapInstance :: HKDConfig -> Name -> Name -> [TyVarBndr BndrVis] -> [Con] -> Q Dec
generateWrapInstance config origTypeName hkdTypeName typeVars constructors = do
  let typeVarNames = map getTyVarName typeVars
  let origType = foldl AppT (ConT origTypeName) (map VarT typeVarNames)
  -- Use concrete HKD type instead of type family for instance head
  -- Apply Identity as a type constructor, not a type variable
  let hkdType = foldl AppT (ConT hkdTypeName) (map VarT typeVarNames ++ [ConT ''Identity])

  -- Generate from clauses for each constructor
  clauses <- mapM (generateWrapClause config) constructors

  let fromMethod = FunD (mkName "from") clauses

  pure $ InstanceD Nothing [] (AppT (AppT (ConT ''Witch.From) origType) hkdType) [fromMethod]
  where
    getTyVarName (PlainTV n _) = n
    getTyVarName (KindedTV n _ _) = n

generateWrapClause :: HKDConfig -> Con -> Q Clause
generateWrapClause config con = case con of
  RecC conName fields -> do
    -- Generate pattern: OrigCon f1 f2 ...
    fieldVars <- mapM (\_ -> newName "x") fields
    let origPattern = RecP conName [(fname, VarP var) | ((fname, _, _), var) <- zip fields fieldVars]

    -- Generate expression: HKDCon' (Identity f1) (Identity f2) ...
    let hkdConName = mkName $ hkdConstructorNameTransform config (nameBase conName)
    let transformedFields = [(mkName $ hkdFieldNameTransform config (nameBase fname), AppE (ConE 'Identity) (VarE var)) | ((fname, _, _), var) <- zip fields fieldVars]
    let hkdConExp = RecConE hkdConName transformedFields

    pure $ Clause [origPattern] (NormalB hkdConExp) []

  NormalC conName fields -> do
    fieldVars <- mapM (\_ -> newName "x") fields
    let origPattern = ConP conName [] [VarP var | var <- fieldVars]
    let hkdConName = mkName $ hkdConstructorNameTransform config (nameBase conName)
    let hkdConExp = foldl AppE (ConE hkdConName) [AppE (ConE 'Identity) (VarE var) | var <- fieldVars]

    pure $ Clause [origPattern] (NormalB hkdConExp) []

  InfixC _ conName _ -> do
    leftVar <- newName "l"
    rightVar <- newName "r"
    let origPattern = InfixP (VarP leftVar) conName (VarP rightVar)
    let hkdConName = mkName $ hkdConstructorNameTransform config (nameBase conName)
    let hkdConExp = InfixE
          (Just (AppE (ConE 'Identity) (VarE leftVar)))
          (ConE hkdConName)
          (Just (AppE (ConE 'Identity) (VarE rightVar)))

    pure $ Clause [origPattern] (NormalB hkdConExp) []

  _ -> fail "Warlock.HKD: GADT constructors not supported in From instances"