{-# LANGUAGE
    TemplateHaskell
  , UnicodeSyntax
  #-}
-- |FIXME: doc
module Data.Collections.Newtype.TH
    ( derive
    )
    where
import Control.Applicative hiding (empty)
import Control.Monad.Unicode
import Data.Collections
import Data.Collections.BaseInstances ()
import Data.Data
import Data.Generics.Aliases
import Data.Generics.Schemes
import Data.Maybe
import Language.Haskell.TH.Lib
import Language.Haskell.TH.Ppr
import Language.Haskell.TH.Syntax
import Prelude hiding ( concat, concatMap, exp, filter
                      , foldl, foldr, foldl1, foldr1, null)
import Prelude.Unicode

type Deriver = Q Cxt → Q Type → Q Exp → Q Exp → Q Dec

-- |FIXME: doc
derive ∷ Q [Dec] → Q [Dec]
derive = (concat <$>) ∘ (mapM go =≪)
    where
      go ∷ Dec → Q [Dec]
      go (InstanceD c ty _) = deriveInstance c ty
      go _ = fail "derive: usage: derive [d| instance A; instance B; ... |]"

deriveInstance ∷ Cxt → Type → Q [Dec]
deriveInstance c ty
    = do (wrapperTy, deriver) ← inspectInstance ty
         (wrap     , wrapD  ) ← genWrap   wrapperTy
         (unwrap   , unwrapD) ← genUnwrap wrapperTy
         instanceDecl         ← deriver (return c     )
                                        (return ty    )
                                        (return wrap  )
                                        (return unwrap)
         return $ [ d | d ← wrapD  , wrap   `isUsedIn` instanceDecl ]
                ⧺ [ d | d ← unwrapD, unwrap `isUsedIn` instanceDecl ]
                ⧺ [ instanceDecl ]

isUsedIn ∷ (Eq α, Typeable α, Data β) ⇒ α → β → Bool
isUsedIn α = (> 0) ∘ gcount (mkQ False (≡ α))

inspectInstance ∷ Type → Q (Type, Deriver)
inspectInstance (AppT (AppT (ConT classTy) wrapperTy) _)
    | classTy ≡ ''Unfoldable
        = return (wrapperTy, deriveUnfoldable)
    | classTy ≡ ''Foldable
        = return (wrapperTy, deriveFoldable)
    | classTy ≡ ''Collection
        = return (wrapperTy, deriveCollection)
inspectInstance (AppT (AppT (AppT (ConT classTy) wrapperTy) _) _)
    | classTy ≡ ''Indexed
        = return (wrapperTy, deriveIndexed)
inspectInstance ty
    = fail $ "deriveInstance: unsupported type: " ⧺ pprint ty

genWrap ∷ Type → Q (Exp, [Dec])
genWrap wrapperTy
    = do name      ← newName "wrap"
         (con, ty) ← wrapperConTy wrapperTy
         decls     ← sequence
                     [ sigD name [t| $(return ty) → $(return wrapperTy) |]
                     , pragInlD name (inlineSpecNoPhase True True)
                     , funD name [clause [] (normalB (conE con)) []]
                     ]
         return (VarE name, decls)

genUnwrap ∷ Type → Q (Exp, [Dec])
genUnwrap wrapperTy
    = do name      ← newName "unwrap"
         i         ← newName "i"
         (con, ty) ← wrapperConTy wrapperTy
         decls     ← sequence
                     [ sigD name [t| $(return wrapperTy) → $(return ty) |]
                     , pragInlD name (inlineSpecNoPhase True True)
                     , funD name [clause [conP con [varP i]] (normalB (varE i)) []]
                     ]
         return (VarE name, decls)

wrapperConTy ∷ Type → Q (Name, Type)
wrapperConTy = (conTy =≪) ∘ tyInfo
    where
      tyInfo ∷ Type → Q Info
      tyInfo (ConT name) = reify name
      tyInfo (AppT ty _) = tyInfo ty
      tyInfo (SigT ty _) = tyInfo ty
      tyInfo ty
          = fail $ "wrapperConTy: unsupported type: " ⧺ pprint ty

      conTy ∷ Info → Q (Name, Type)
      conTy (TyConI (NewtypeD [] _ [] (NormalC con [(NotStrict, ty)]) []))
          = return (con, ty)
      conTy info
          = fail $ "wrapperConTy: unsupported type: " ⧺ pprint info

methodNames ∷ Name → Q [Name]
methodNames = (names =≪) ∘ reify
    where
      names ∷ Info → Q [Name]
      names (ClassI (ClassD _ _ _ _ decls) _)
              = return ∘ catMaybes $ map name decls
      names c = fail $ "methodNames: not a class: " ⧺ pprint c

      name ∷ Dec → Maybe Name
      name (SigD n _) = Just n
      name _          = Nothing

pointfreeMethod ∷ (Name → Q Exp) → Name → [Q Dec]
pointfreeMethod f name
    = [ funD name [clause [] (normalB (f name)) []]
      -- THINKME: Inserting PragmaD in an InstanceD causes an error
      -- least GHC 7.0.3. Why?
      -- , pragInlD name (inlineSpecNoPhase True False)
      ]

deriveUnfoldable ∷ Q Cxt → Q Type → Q Exp → Q Exp → Q Dec
deriveUnfoldable c ty wrap unwrap
    = do names ← methodNames ''Unfoldable
         instanceD c ty $ concatMap (pointfreeMethod exp) names
    where
      exp ∷ Name → Q Exp
      exp name
          | name ≡ 'insert
              = [| ($wrap ∘) ∘ (∘ $unwrap) ∘ insert |]
          | name ≡ 'empty
              = [| $wrap empty |]
          | name ≡ 'singleton
              = [| $wrap ∘ singleton |]
          | name ≡ 'insertMany
              = [| ($wrap ∘) ∘ (∘ $unwrap) ∘ insertMany |]
          | name ≡ 'insertManySorted
              = [| ($wrap ∘) ∘ (∘ $unwrap) ∘ insertManySorted |]
          | otherwise
              = fail $ "deriveUnfoldable: unknown method: " ⧺ pprint name

deriveFoldable ∷ Q Cxt → Q Type → Q Exp → Q Exp → Q Dec
deriveFoldable c ty _ unwrap
    = do names ← methodNames ''Foldable
         instanceD c ty $ concatMap (pointfreeMethod exp) names
    where
      exp ∷ Name → Q Exp
      exp name
          | name ≡ 'fold
              = [| fold ∘ $unwrap |]
          | name ≡ 'foldMap
              = [| (∘ $unwrap) ∘ foldMap |]
          | name ≡ 'foldr
              = [| flip flip $unwrap ∘ ((∘) ∘) ∘ foldr |]
          | name ≡ 'foldl
              = [| flip flip $unwrap ∘ ((∘) ∘) ∘ foldl |]
          | name ≡ 'foldr1
              = [| (∘ $unwrap) ∘ foldr1 |]
          | name ≡ 'foldl1
              = [| (∘ $unwrap) ∘ foldl1 |]
          | name ≡ 'null
              = [| null ∘ $unwrap |]
          | name ≡ 'size
              = [| size ∘ $unwrap |]
          | name ≡ 'isSingleton
              = [| isSingleton ∘ $unwrap |]
          | otherwise
              = fail $ "deriveFoldable: unknown method: " ⧺ pprint name

deriveCollection ∷ Q Cxt → Q Type → Q Exp → Q Exp → Q Dec
deriveCollection c ty wrap unwrap
    = do names ← methodNames ''Collection
         instanceD c ty $ concatMap (pointfreeMethod exp) names
    where
      exp ∷ Name → Q Exp
      exp name
          | name ≡ 'filter
              = [| ($wrap ∘) ∘ (∘ $unwrap) ∘ filter |]
          | otherwise
              = fail $ "deriveCollection: unknown method: " ⧺ pprint name

deriveIndexed ∷ Q Cxt → Q Type → Q Exp → Q Exp → Q Dec
deriveIndexed c ty wrap unwrap
    = do names ← methodNames ''Indexed
         instanceD c ty $ concatMap (pointfreeMethod exp) names
    where
      exp ∷ Name → Q Exp
      exp name
          | name ≡ 'index
              = [| (∘ $unwrap) ∘ index |]
          | name ≡ 'adjust
              = [| (($wrap ∘) ∘) ∘ flip flip $unwrap ∘ ((∘) ∘) ∘ adjust |]
          | name ≡ 'inDomain
              = [| (∘ $unwrap) ∘ inDomain |]
          | name ≡ '(//)
              = [| ($wrap ∘) ∘ (//) ∘ $unwrap |]
          | name ≡ 'accum
              = [| (($wrap ∘) ∘) ∘ (∘ $unwrap) ∘ accum |]
          | otherwise
              = fail $ "deriveIndexed: unknown method: " ⧺ pprint name