{-# LANGUAGE TemplateHaskellQuotes #-}
module Optics.TH.Internal.Sum
  ( makePrisms
  , makePrismLabels
  , makeClassyPrisms
  , makeDecPrisms
  ) where

import Data.Char
import Data.List
import Data.Maybe
import Data.Traversable
import Language.Haskell.TH
import Language.Haskell.TH.Datatype.TyVarBndr
import qualified Data.Map as M
import qualified Data.Set as S
import qualified Language.Haskell.TH.Datatype as D

import Data.Set.Optics
import Language.Haskell.TH.Optics.Internal
import Optics.Core hiding (cons)
import Optics.Internal.Magic
import Optics.TH.Internal.Utils

-- | Generate a 'Prism' for each constructor of a data type. Isos generated when
-- possible. Reviews are created for constructors with existentially quantified
-- constructors and GADTs.
--
-- /e.g./
--
-- @
-- data FooBarBaz a
--   = Foo Int
--   | Bar a
--   | Baz Int Char
-- makePrisms ''FooBarBaz
-- @
--
-- will create
--
-- @
-- _Foo :: Prism' (FooBarBaz a) Int
-- _Bar :: Prism (FooBarBaz a) (FooBarBaz b) a b
-- _Baz :: Prism' (FooBarBaz a) (Int, Char)
-- @
makePrisms :: Name {- ^ Type constructor name -} -> DecsQ
makePrisms = makePrisms' True

-- | Generate a 'Prism' for each constructor of a data type and combine them
-- into a single class. No Isos are created. Reviews are created for
-- constructors with existentially quantified constructors and GADTs.
--
-- /e.g./
--
-- @
-- data FooBarBaz a
--   = Foo Int
--   | Bar a
--   | Baz Int Char
-- makeClassyPrisms ''FooBarBaz
-- @
--
-- will create
--
-- @
-- class AsFooBarBaz s a | s -> a where
--   _FooBarBaz :: Prism' s (FooBarBaz a)
--   _Foo :: Prism' s Int
--   _Bar :: Prism' s a
--   _Baz :: Prism' s (Int,Char)
--
--   _Foo = _FooBarBaz % _Foo
--   _Bar = _FooBarBaz % _Bar
--   _Baz = _FooBarBaz % _Baz
--
-- instance AsFooBarBaz (FooBarBaz a) a
-- @
--
-- Generate an "As" class of prisms. Names are selected by prefixing the
-- constructor name with an underscore. Constructors with multiple fields will
-- construct Prisms to tuples of those fields.
makeClassyPrisms :: Name {- ^ Type constructor name -} -> DecsQ
makeClassyPrisms = makePrisms' False

makePrismLabels :: Name -> DecsQ
makePrismLabels typeName = do
  requireExtensionsForLabels
  info <- D.reifyDatatype typeName
  let cons = map (normalizeCon info) $ D.datatypeCons info
  catMaybes <$> traverse (makeLabel info cons) cons
  where
    makeLabel :: D.DatatypeInfo -> [NCon] -> NCon -> Q (Maybe Dec)
    makeLabel info cons con = do
      stab@(Stab tvsCovered cx otype s t a b) <- computeOpticType labelConfig ty cons con
      (k,  cxtK) <- eqSubst (ConT $ opticTypeToTag otype) "k"
      (a', cxtA) <- eqSubst a "a"
      (b', cxtB) <- eqSubst b "b"
      let label = nameBase . prismName $ view nconName con
          tyArgs = [LitT (StrTyLit label), k, s, t, a', b']
          context = concat
            [ -- If some of the type variables are not covered, instance is
              -- dysfunctional.
              if tvsCovered then [] else [conAppsT ''Dysfunctional tyArgs]
            , [cxtK, cxtA, cxtB]
            , cx
            ]
      Just <$> instanceD (pure context)
                         (pure $ conAppsT ''LabelOptic tyArgs)
                         (fun stab 'labelOptic)
      where
        ty        = addKindInfo info $ D.datatypeType info
        isNewtype = D.datatypeVariant info == D.Newtype

        opticTypeToTag IsoType    = ''An_Iso
        opticTypeToTag PrismType  = ''A_Prism
        opticTypeToTag ReviewType = ''A_Review -- for complete match

        fun :: Stab -> Name -> [DecQ]
        fun stab n = valD (varP n) (normalB $ funDef stab) [] : inlinePragma n

        funDef :: Stab -> ExpQ
        funDef stab
          | isNewtype = varE 'coerced
          | otherwise = makeConOpticExp stab cons con

-- | Main entry point into Prism generation for a given type constructor name.
makePrisms' :: Bool -> Name -> DecsQ
makePrisms' normal typeName =
  do info <- D.reifyDatatype typeName
     let cls | normal    = Nothing
             | otherwise = Just (D.datatypeName info)
         cons = D.datatypeCons info
     makeConsPrisms info (map (normalizeCon info) cons) cls


-- | Generate prisms for the given 'Dec'
makeDecPrisms :: Bool {- ^ generate top-level definitions -} -> Dec -> DecsQ
makeDecPrisms normal dec =
  do info <- D.normalizeDec dec
     let cls | normal    = Nothing
             | otherwise = Just (D.datatypeName info)
         cons = D.datatypeCons info
     makeConsPrisms info (map (normalizeCon info) cons) cls

-- | Generate prisms for the given type, normalized constructors, and an
-- optional name to be used for generating a prism class. This function
-- dispatches between Iso generation, normal top-level prisms, and classy
-- prisms.
makeConsPrisms :: D.DatatypeInfo -> [NCon] -> Maybe Name -> DecsQ

-- top-level definitions
makeConsPrisms info cons Nothing = fmap concat . for cons $ \con -> do
  stab <- computeOpticType defaultConfig ty cons con
  let n = prismName $ view nconName con
      body = if isNewtype
             then varE 'coerced
             else makeConOpticExp stab cons con
  sequenceA $
    [ sigD n . pure . close $ stabToType stab
    , valD (varP n) (normalB body) []
    ] ++ inlinePragma n
  where
    ty        = addKindInfo info $ D.datatypeType info
    isNewtype = D.datatypeVariant info == D.Newtype

-- classy prism class and instance
makeConsPrisms info cons (Just typeName) =
  sequenceA
    [ makeClassyPrismClass ty className methodName cons
    , makeClassyPrismInstance ty className methodName cons
    ]
  where
    ty = D.datatypeType info
    className = mkName ("As" ++ nameBase typeName)
    methodName = prismName typeName

----------------------------------------

data StabConfig = StabConfig
  { scForLabelInstance :: Bool
  , scAllowIsos        :: Bool
  }

defaultConfig :: StabConfig
defaultConfig = StabConfig
  { scForLabelInstance = False
  , scAllowIsos        = True
  }

classyConfig :: StabConfig
classyConfig = StabConfig
  { scForLabelInstance = False
  , scAllowIsos        = False
  }

labelConfig :: StabConfig
labelConfig = StabConfig
  { scForLabelInstance = True
  , scAllowIsos        = True
  }

data OpticType = IsoType | PrismType | ReviewType
  deriving Eq
data Stab  = Stab Bool Cxt OpticType Type Type Type Type

simplifyStab :: Stab -> Stab
simplifyStab (Stab tvsCovered cx ty _ t _ b) = Stab tvsCovered cx ty t t b b
  -- simplification uses t and b because those types
  -- are interesting in the Review case

stabSimple :: Stab -> Bool
stabSimple (Stab _ _ _ s t a b) = s == t && a == b

stabToType :: Stab -> Type
stabToType stab@(Stab _ cx ty s t a b) = ForallT vs cx $
  case ty of
    IsoType    | stabSimple stab -> ''Iso'    `conAppsT` [s,a]
               | otherwise       -> ''Iso     `conAppsT` [s,t,a,b]
    PrismType  | stabSimple stab -> ''Prism'  `conAppsT` [s,a]
               | otherwise       -> ''Prism   `conAppsT` [s,t,a,b]
    ReviewType                   -> ''Review  `conAppsT` [t,b]

  where
    vs = changeTVFlags SpecifiedSpec
       . D.freeVariablesWellScoped
       . S.toList
       $ setOf (folded % typeVarsKinded) cx

stabType :: Stab -> OpticType
stabType (Stab _ _ o _ _ _ _) = o

computeOpticType :: StabConfig -> Type -> [NCon] -> NCon -> Q Stab
computeOpticType conf t cons con =
  do let cons' = delete con cons
     if null (_nconVars con)
         then computePrismType conf t (view nconCxt con) cons' con
         else computeReviewType t (view nconCxt con) (view nconTypes con)

computeReviewType :: Type -> Cxt -> [Type] -> Q Stab
computeReviewType t cx tys = do
  b <- toTupleT (map return tys)
  return (Stab False cx ReviewType t t b b)

-- | Compute the full type-changing Prism type given an outer type, list of
-- constructors, and target constructor name.
computePrismType :: StabConfig -> Type -> Cxt -> [NCon] -> NCon -> Q Stab
computePrismType conf s cx cons con = do
  let ts       = view nconTypes con
      free     = setOf typeVars s
      fixed    = setOf typeVars cons
      phantoms = free S.\\ setOf (folded % nconTypes % typeVars) (con : cons)

      unbound    = free S.\\ fixed
      tvsCovered = if scForLabelInstance conf
                   then S.null phantoms
                   else True
  sub <- sequenceA (M.fromSet (newName . nameBase) unbound)
  a   <- toTupleT (map return ts)
  b   <- toTupleT (map return (substTypeVars sub ts))
  --runIO $ do
  --  putStrLn $ "S:        " ++ show s
  --  putStrLn $ "A:        " ++ show a
  --  putStrLn $ "FREE:     " ++ show free
  --  putStrLn $ "FIXED:    " ++ show fixed
  --  putStrLn $ "PHANTOMS: " ++ show phantoms
  --  putStrLn $ "UNBOUND:  " ++ show unbound
  let t = substTypeVars sub s
      cx' = substTypeVars sub cx
      otype = if null cons && scAllowIsos conf
              then IsoType
              else PrismType
  return (Stab tvsCovered cx' otype s t a b)

-- | Construct either a Review or Prism as appropriate
makeConOpticExp :: Stab -> [NCon] -> NCon -> ExpQ
makeConOpticExp stab cons con =
  case stabType stab of
    IsoType    -> makeConIsoExp con
    PrismType  -> makeConPrismExp stab cons con
    ReviewType -> makeConReviewExp con

-- | Construct prism expression
--
-- prism <<reviewer>> <<remitter>>
makeConPrismExp ::
  Stab ->
  [NCon] {- ^ constructors       -} ->
  NCon   {- ^ target constructor -} ->
  ExpQ
makeConPrismExp stab cons con = appsE [varE 'prism, reviewer, remitter]
  where
  ts = view nconTypes con
  fields  = length ts
  conName = view nconName con

  reviewer                   = makeReviewer       conName fields
  remitter | stabSimple stab = makeSimpleRemitter conName fields
           | otherwise       = makeFullRemitter cons conName


-- | Construct an Iso expression
--
-- iso <<reviewer>> <<remitter>>
makeConIsoExp :: NCon -> ExpQ
makeConIsoExp con = appsE [varE 'iso, remitter, reviewer]
  where
  conName = view nconName con
  fields  = length (view nconTypes con)

  reviewer = makeReviewer    conName fields
  remitter = makeIsoRemitter conName fields


-- | Construct a Review expression
--
-- unto (\(x,y,z) -> Con x y z)
makeConReviewExp :: NCon -> ExpQ
makeConReviewExp con = appE (varE 'unto) reviewer
  where
  conName = view nconName con
  fields  = length (view nconTypes con)

  reviewer = makeReviewer conName fields


------------------------------------------------------------------------
-- Prism and Iso component builders
------------------------------------------------------------------------


-- | Construct the review portion of a prism.
--
-- (\(x,y,z) -> Con x y z) :: b -> t
makeReviewer :: Name -> Int -> ExpQ
makeReviewer conName fields =
  do xs <- newNames "x" fields
     lam1E (toTupleP (map varP xs))
           (conE conName `appsE1` map varE xs)


-- | Construct the remit portion of a prism.
-- Pattern match only target constructor, no type changing
--
-- (\s -> case s of
--          Con x y z -> Right (x,y,z)
--          _         -> Left s
-- ) :: s -> Either s a
makeSimpleRemitter :: Name -> Int -> ExpQ
makeSimpleRemitter conName fields =
  do x  <- newName "x"
     xs <- newNames "y" fields
     let matches =
           [ match (conP conName (map varP xs))
                   (normalB (appE (conE 'Right) (toTupleE (map varE xs))))
                   []
           , match wildP (normalB (appE (conE 'Left) (varE x))) []
           ]
     lam1E (varP x) (caseE (varE x) matches)


-- | Pattern match all constructors to enable type-changing
--
-- (\s -> case s of
--          Con x y z -> Right (x,y,z)
--          Other_n w   -> Left (Other_n w)
-- ) :: s -> Either t a
makeFullRemitter :: [NCon] -> Name -> ExpQ
makeFullRemitter cons target =
  do x <- newName "x"
     lam1E (varP x) (caseE (varE x) (map mkMatch cons))
  where
  mkMatch (NCon conName _ _ n) =
    do xs <- newNames "y" (length n)
       match (conP conName (map varP xs))
             (normalB
               (if conName == target
                  then appE (conE 'Right) (toTupleE (map varE xs))
                  else appE (conE 'Left) (conE conName `appsE1` map varE xs)))
             []


-- | Construct the remitter suitable for use in an 'Iso'
--
-- (\(Con x y z) -> (x,y,z)) :: s -> a
makeIsoRemitter :: Name -> Int -> ExpQ
makeIsoRemitter conName fields =
  do xs <- newNames "x" fields
     lam1E (conP conName (map varP xs))
           (toTupleE (map varE xs))


------------------------------------------------------------------------
-- Classy prisms
------------------------------------------------------------------------


-- | Construct the classy prisms class for a given type and constructors.
--
-- class ClassName r <<vars in type>> | r -> <<vars in Type>> where
--   topMethodName   :: Prism' r Type
--   conMethodName_n :: Prism' r conTypes_n
--   conMethodName_n = topMethodName . conMethodName_n
makeClassyPrismClass ::
  Type   {- Outer type      -} ->
  Name   {- Class name      -} ->
  Name   {- Top method name -} ->
  [NCon] {- Constructors    -} ->
  DecQ
makeClassyPrismClass t className methodName cons =
  do r <- newName "r"
     let methodType = appsT (conT ''Prism') [varT r,return t]
     methodss <- traverse (mkMethod (VarT r)) cons'
     classD (cxt[]) className (map plainTV (r : vs)) (fds r)
       ( sigD methodName methodType
       : map return (concat methodss)
       )

  where
  mkMethod r con =
    do Stab tvsCovered cx o _ _ _ b <- computeOpticType classyConfig t cons con
       let stab' = Stab tvsCovered cx o r r b b
           defName = view nconName con
           body    = appsE [varE '(%), varE methodName, varE defName]
       sequenceA
         [ sigD defName        (return (stabToType stab'))
         , valD (varP defName) (normalB body) []
         ]

  cons'         = map (over nconName prismName) cons
  vs            = S.toList (setOf typeVars t)
  fds r
    | null vs   = []
    | otherwise = [FunDep [r] vs]



-- | Construct the classy prisms instance for a given type and constructors.
--
-- instance Classname OuterType where
--   topMethodName = id
--   conMethodName_n = <<prism>>
makeClassyPrismInstance ::
  Type ->
  Name     {- Class name      -} ->
  Name     {- Top method name -} ->
  [NCon] {- Constructors    -} ->
  DecQ
makeClassyPrismInstance s className methodName cons =
  do let vs = S.toList (setOf typeVars s)
         cls = className `conAppsT` (s : map VarT vs)

     instanceD (cxt[]) (return cls)
       (   valD (varP methodName)
                (normalB (varE 'castOptic `appE` varE 'equality)) []
       : [ do stab <- computeOpticType classyConfig s cons con
              let stab' = simplifyStab stab
              valD (varP (prismName conName))
                (normalB (makeConOpticExp stab' cons con)) []
           | con <- cons
           , let conName = view nconName con
           ]
       )


------------------------------------------------------------------------
-- Utilities
------------------------------------------------------------------------


-- | Normalized constructor
data NCon = NCon
  { _nconName :: Name
  , _nconVars :: [Name]
  , _nconCxt  :: Cxt
  , _nconTypes :: [Type]
  }
  deriving (Eq, Show)

instance HasTypeVars NCon where
  typeVarsEx s = traversalVL $ \f (NCon x vars y z) ->
    let s' = foldl' (flip S.insert) s vars
    in NCon x vars <$> traverseOf (typeVarsEx s') f y
                   <*> traverseOf (typeVarsEx s') f z

nconName :: Lens' NCon Name
nconName = lensVL $ \f x -> fmap (\y -> x {_nconName = y}) (f (_nconName x))

nconCxt :: Lens' NCon Cxt
nconCxt = lensVL $ \f x -> fmap (\y -> x {_nconCxt = y}) (f (_nconCxt x))

nconTypes :: Lens' NCon [Type]
nconTypes = lensVL $ \f x -> fmap (\y -> x {_nconTypes = y}) (f (_nconTypes x))


-- | Normalize a single 'Con' to its constructor name and field types.
normalizeCon :: D.DatatypeInfo -> D.ConstructorInfo -> NCon
normalizeCon di info = NCon
  { _nconName  = D.constructorName info
  , _nconVars  = D.tvName <$> D.constructorVars info
  , _nconCxt   = D.constructorContext info
  , _nconTypes = let tyVars = map tyVarBndrToType (D.constructorVars info)
                 in addKindInfo' tyVars di <$> D.constructorFields info
  }


-- | Compute a prism's name by prefixing an underscore for normal
-- constructors and period for operators.
prismName :: Name -> Name
prismName n = case nameBase n of
                [] -> error "prismName: empty name base?"
                x:xs | isUpper x -> mkName ('_':x:xs)
                     | otherwise -> mkName ('.':x:xs) -- operator


-- | Quantify all the free variables in a type.
close :: Type -> Type
close (ForallT vars cx ty) = quantifyType vars cx ty
close ty                   = quantifyType []   [] ty