module Data.Packed.TH.Write (genWrite, writeFName) where

import Data.Packed.Needs
import Data.Packed.Packable
import Data.Packed.TH.Flag (PackingFlag)
import Data.Packed.TH.Utils
import Data.Packed.TH.WriteCon
import Language.Haskell.TH

-- For a data type 'Tree', will generate the function name 'writeTree'
writeFName :: Name -> Name
writeFName tyName = mkName $ "write" ++ nameBase tyName

-- | Generates a function that serialises and writes a value into a 'Needs'
--
-- The function simply calls the functions generated by 'Data.Packed.TH.genConWrite'
--
-- __Example:__
--
-- For the 'Tree' data type, it generates the following function
--
-- @
-- writeTree :: ('Packable' a) => Tree a -> 'NeedsWriter' (Tree a) r t
-- writeTree (Leaf n) = writeConLeaf n
-- writeTree (Node l r) = writeConNode l r
-- @
genWrite ::
    [PackingFlag] ->
    -- | The name of the type to generate the function for
    Name ->
    Q [Dec]
genWrite flags tyName = do
    (TyConI (DataD _ _ _ _ cs _)) <- reify tyName
    -- For each data constructor, we generate the corresponding clause
    clauses <-
        mapM
            ( \con -> do
                let (conName, types) = getNameAndBangTypesFromCon con
                -- Generate names for each variable in the constructor
                paramNames <- mapM (const $ newName "t") types
                -- We apply each parameter of the constructor and the 'Needs' to the 'writeConXXX' function
                body <- foldl (\f arg -> [|$f $(varE arg)|]) (varE $ conWriteFName conName) paramNames
                return $ Clause [ConP conName [] (VarP <$> paramNames)] (NormalB body) []
            )
            cs
    -- For each of the data constructor of the type, we generate the corresponding `writeConXXX`
    -- We define the Tag using the index of the data constructor
    conWriter <-
        mapM
            ( \(index, constructor) ->
                let (conName, types) = getNameAndBangTypesFromCon constructor
                 in genConWrite flags conName index types
            )
            $ zip [0 ..] cs
    signature <- genWriteSignature tyName
    return $ concat conWriter ++ [signature, FunD (writeFName tyName) clauses]

-- Generates the following function signature for a data type 'Tree'
-- writeTree :: ('Packable' a) => Tree a -> 'NeedsWriter' (Tree a) r t
genWriteSignature :: Name -> Q Dec
genWriteSignature tyName = do
    (sourceType, typeParameterNames) <- resolveAppliedType tyName
    let fName = writeFName tyName
        -- Type variables for Needs
        r = varT $ mkName "r"
        t = varT $ mkName "t"
        -- Define Packable constraints on each of the type parameters
        constraints = mapM (\tyVarName -> [t|Packable $(varT tyVarName)|]) typeParameterNames
        signature = [t|$(return sourceType) -> NeedsWriter $(return sourceType) $r $t|]
    sigD fName (forallT [] constraints signature)