{-# LANGUAGE
  CPP,
  DeriveDataTypeable,
  GADTs,
  KindSignatures,
  PolyKinds,
  ScopedTypeVariables #-}

-- | Derive instances without spelling out "deriving".
--
-- = Usage
--
-- __Step 1__: add this pragma at the top of the file to load the plugin:
--
-- @
-- {-\# OPTIONS_GHC -fplugin=Driving.Classes \#-}
-- @
--
-- __Step 2__: enable @DerivingStrategies@ and other relevant extensions as needed
-- (@DerivingVia@, @GeneralizedNewtypeDeriving@, @DeriveAnyClass@):
--
-- @
-- {-\# LANGUAGE DerivingStrategies \#-}
-- @
--
-- __Step 3__: add an @ANN@ pragma after imports to configure the classes to auto-derive:
--
-- @
-- {-\# ANN module (Driving :: Driving '[ \<LIST OF OPTIONS\> ]) \#-}
-- @
--
-- = Example
--
-- This automatically declares instances of @Eq@, @Ord@, @Show@ for @T@, @U@, @V@,
-- and disables auto-deriving for @MyEndo@.
--
-- @
-- {-\# ANN module (Driving :: Driving
--   '[ Stock '(Eq, Ord, Show)
--    , NoDriving '(Eq MyEndo, Ord MyEndo, Show MyEndo)
--    ]) \#-}
--
-- data T = C1 | C2
-- data U = D1 | D2
-- data V = E1 | E2
--
-- newtype MyEndo a = MyEndo (a -> a)
-- @
--
-- Available options:
--
-- - 'Stock'
-- - 'Anyclass'
-- - 'Newtype'
-- - 'Via'
-- - 'ViaF'
-- - 'NoDriving'
--
-- See more examples below.

#if __GLASGOW_HASKELL__ >= 902
#define PRE902(x)
#define POST902(x) x
#else
#define PRE902(x) x
#define POST902(x)
#endif

module Driving.Classes
  ( -- * Options
    Driving(..)
  , Stock
  , Newtype
  , Anyclass
  , Via
  , ViaF
  , NoDriving

    -- * Plugin
  , plugin) where

import Data.Bifunctor (first)
import Data.Data (Data)
import Data.Kind (Type)

import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set

#if __GLASGOW_HASKELL__ >= 900
import GHC.Plugins hiding (Type)
#else
import GhcPlugins hiding (Type)
#endif
#if __GLASGOW_HASKELL__ >= 810
import GHC.Hs
#else
import HsSyn
#define NoExtField NoExt
#endif

-- * User configuration

-- | Type constructor for configuring the plugin in a source annotation.
--
-- Argument: list of types using the constructors below.
--
-- === Example
--
-- @
-- {-\# ANN module (Driving :: 'Driving' '[ 'Stock' '(Eq, Ord), 'Newtype' Num ]) \#-}
-- @
data Driving :: k -> Type where
  -- | Dummy constructor
  Driving :: Driving x
  deriving Data

-- | Auto-derive classes using the @stock@ deriving strategy.
--
-- Argument: a class (of kind @k -> Constraint@ for some @k@), or a tuple of classes.
--
-- === Examples
--
-- @
-- {-\# ANN module (Driving :: 'Driving' '[ 'Stock' Show ]) \#-}
-- {-\# ANN module (Driving :: 'Driving' '[ 'Stock' '(Eq, Ord) ]) \#-}
-- @
data Stock :: k -> Type

-- | Auto-derive classes using the @newtype@ deriving strategy.
-- Enable the extension @GeneralizedNewtypeDeriving@ to use this.
--
-- Argument: a class (of kind @k -> Constraint@ for some @k@), or a tuple of classes.
--
-- === Examples
--
-- @
-- {-\# ANN module (Driving :: 'Driving' '[ 'Newtype' Num ]) \#-}
-- {-\# ANN module (Driving :: 'Driving' '[ 'Newtype' '(Semigroup, Monoid)]) \#-}
-- @
data Newtype :: k -> Type

-- | Auto-derive classes using the @anyclass@ deriving strategy.
-- Enable the extension @DeriveAnyClass@ to use this.
--
-- Argument: a class (of kind @k -> Constraint@ for some @k@), or a tuple of classes.
--
-- === Examples
--
-- @
-- {-\# ANN module (Driving :: 'Driving' '[ 'Anyclass' Binary ]) \#-}
-- {-\# ANN module (Driving :: 'Driving' '[ 'Anyclass' '(ToJSON, FromJSON) ]) \#-}
-- -- Classes from the packages binary and aeson
-- @
data Anyclass :: k -> Type

-- | Auto-derive classes using the @via@ deriving strategy, for a given via-type.
-- Enable the extension @DerivingVia@ to use this.
--
-- Arguments:
--
-- 1. a class (of kind @k -> Constraint@ for some @k@), or a tuple of classes;
-- 2. a type.
--
-- === Examples
--
-- @
-- {-\# ANN module (Driving :: 'Driving' '[ Num `'Via'` Int ]) \#-}
-- {-\# ANN module (Driving :: 'Driving' '[ '(Eq, Ord) `'Via'` Int ]) \#-}
-- @
data Via :: k -> l -> Type

-- | Auto-derive classes using the @via@ deriving strategy, where the via-type
-- is an application of a given type constructor to each newly declared type.
-- Enable the extension @DerivingVia@ to use this.
--
-- Arguments:
--
-- 1. a class (of kind @k -> Constraint@ for some @k@), or a tuple of classes;
-- 2. a type constructor.
--
-- === Examples
--
-- @
-- {-\# ANN module (Driving :: 'Driving' '[ '(Functor, Applicative) `'ViaF'` WrappedMonad ]) \#-}
-- {-\# ANN module (Driving :: 'Driving' '[ '(Semigroup, Monoid) `'ViaF'` Generically ]) \#-}
-- -- Generically from the package generic-data
-- @
data ViaF :: k -> l -> Type

-- | Cancel auto-deriving for a particular instance.
--
-- Argument: an application of a class to a type, or a tuple of those.
--
-- === Example
--
-- Derive @Show@ for all types except @MyType@:
--
-- @
-- {-\# ANN module (Driving :: 'Driving' '[ 'Stock' Show, 'NoDriving' (Show MyType) ]) \#-}
-- @
data NoDriving :: k -> Type

-- * Plugin

-- | For the compiler.
plugin :: Plugin
plugin = defaultPlugin
  { parsedResultAction = parsed
  }

-- ** Implementation

parsed :: [CommandLineOption] -> ModSummary -> HsParsedModule -> Hsc HsParsedModule
parsed _opts _modsum m = pure (driving m)

type Endo a = a -> a
type DrivingPass a = Config -> a -> a

driving :: Endo HsParsedModule
driving m = m { hpm_module = fmap drivingMod (hpm_module m) }

#if __GLASGOW_HASKELL__ >= 900
drivingMod :: Endo HsModule
#else
drivingMod :: Endo (HsModule GhcPs)
#endif
drivingMod m@HsModule{ hsmodDecls = ds } = m { hsmodDecls = drivingDecls emptyConfig ds }

-- *** AST Traversal

-- | Traverse the source top-down, any annotation using @Driving@ overrides the
-- configuration.
drivingDecls :: DrivingPass [LHsDecl GhcPs]
drivingDecls _conf [] = []
drivingDecls conf (d : ds)
    -- Erase plugin annotations. They can't go through the renamer because they break
    -- the staging restriction by refering to types in the current module.
    -- Also some annotations are ill-kinded. Very sloppy API...
  | Just newConf <- getConf d = drivingDecls newConf ds
  | otherwise = fmap (drivingDecl conf) d : drivingDecls conf ds

drivingDecl :: DrivingPass (HsDecl GhcPs)
drivingDecl conf (TyClD x d@DataDecl{ tcdDataDefn = dd }) =
  TyClD x (d { tcdDataDefn = drivingDataDefn (tcdName d) conf dd })
drivingDecl _conf decl = decl

drivingDataDefn :: RdrName -> DrivingPass (HsDataDefn GhcPs)
drivingDataDefn tyname conf dd@HsDataDefn{ dd_derivs = derivs } =
  dd { dd_derivs = PRE902(fmap) (drivingDerivs tyname conf) derivs }

drivingDerivs :: RdrName -> DrivingPass [LHsDerivingClause GhcPs]
drivingDerivs tyname conf derivs = extraDerivingClauses tyname conf ++ derivs

extraDerivingClauses :: RdrName -> Config -> [LHsDerivingClause GhcPs]
extraDerivingClauses tyname conf = hsClauses
  where
    clauses =
      let clauses0 = drivingClauses conf in
      case Map.lookup tyname (exceptions conf) of
        Nothing -> clauses0
        Just excs -> DrivingClauses
          { drivingStock    = filter (headNoMatch excs) (drivingStock clauses0)
          , drivingNewtype  = filter (headNoMatch excs) (drivingNewtype clauses0)
          , drivingAnyclass = filter (headNoMatch excs) (drivingAnyclass clauses0)
          , drivingVia      = (fmap . first) (filter (headNoMatch excs)) (drivingVia clauses0)
          , drivingViaF     = (fmap . first) (filter (headNoMatch excs)) (drivingViaF clauses0)
          }
    hsClauses =
         mkDerivingClauses (StockStrategy POST902(noAnn)) (drivingStock clauses)
      ++ mkDerivingClauses (NewtypeStrategy POST902(noAnn)) (drivingNewtype clauses)
      ++ mkDerivingClauses (AnyclassStrategy POST902(noAnn)) (drivingAnyclass clauses)
      ++ (mkDerivingViaClauses =<< drivingVia clauses)
      ++ (mkDerivingViaClauses =<< (fmap . fmap) applyToTyname (drivingViaF clauses))
    applyToTyname f = noLocA (HsAppTy NoExtField f (noLocA (hsTyVar tyname)))

headNoMatch :: Set RdrName -> LHsType GhcPs -> Bool
headNoMatch excs (L _ (HsParTy _ t)) = headNoMatch excs t
headNoMatch excs (L _ (HsAppTy _ t _)) = headNoMatch excs t
headNoMatch excs (L _ (HsTyVar _ _ (L _ v))) = Set.notMember v excs
headNoMatch _ _ = True

mkDerivingClauses :: DerivStrategy GhcPs -> [LHsType GhcPs] -> [LHsDerivingClause GhcPs]
mkDerivingClauses _ [] = []
mkDerivingClauses strat cls =
  [ noLoc (HsDerivingClause
      { deriv_clause_ext = noAnn
      , deriv_clause_strategy = Just (noLoc strat)
      , deriv_clause_tys = noLocA (POST902(mkDerivingClausesTys) (map hsTypeToHsSigType cls))
      })
  ]

#if __GLASGOW_HASKELL__ >= 902
-- Input: one or more
mkDerivingClausesTys :: [LHsSigType GhcPs] -> DerivClauseTys GhcPs
mkDerivingClausesTys [c] = DctSingle NoExtField c
mkDerivingClausesTys cls = DctMulti NoExtField cls
#endif

mkDerivingViaClauses :: ([LHsType GhcPs], LHsType GhcPs) -> [LHsDerivingClause GhcPs]
mkDerivingViaClauses (cls, v) =
#if __GLASGOW_HASKELL__ >= 902
  let s = XViaStrategyPs noAnn (hsTypeToHsSigType v) in
#else
  let s = mkHsImplicitBndrs v in
#endif
  mkDerivingClauses (ViaStrategy s) cls

hsTyVar :: RdrName -> HsType GhcPs
hsTyVar = HsTyVar noAnn NotPromoted . noLocA

-- *** Configuration

data Config = Config
  { drivingClauses :: DrivingClauses
  , exceptions :: Map RdrName (Set RdrName)
  }

data DrivingClauses = DrivingClauses
  { drivingStock    :: [LHsType GhcPs]
  , drivingNewtype  :: [LHsType GhcPs]
  , drivingAnyclass :: [LHsType GhcPs]
  , drivingVia      :: [([LHsType GhcPs], LHsType GhcPs)]
  , drivingViaF     :: [([LHsType GhcPs], LHsType GhcPs)]
  }

addException :: RdrName -> RdrName -> Config -> Config
addException ty cls config = config { exceptions = Map.alter add ty (exceptions config) } where
  add Nothing = Just (Set.singleton cls)
  add (Just clss) = Just (Set.insert cls clss)

updateDrivingClauses :: (DrivingClauses -> DrivingClauses) -> Config -> Config
updateDrivingClauses f conf = conf { drivingClauses = f (drivingClauses conf) }

addStock, addNewtype, addAnyclass :: [LHsType GhcPs] -> DrivingClauses -> DrivingClauses
addStock    names dc = dc { drivingStock    = names ++ drivingStock dc }
addNewtype  names dc = dc { drivingNewtype  = names ++ drivingNewtype dc }
addAnyclass names dc = dc { drivingAnyclass = names ++ drivingAnyclass dc }

addVia :: [LHsType GhcPs] -> LHsType GhcPs -> DrivingClauses -> DrivingClauses
addVia names v dc = dc { drivingVia = (names, v) : drivingVia dc }

addViaF :: [LHsType GhcPs] -> LHsType GhcPs -> DrivingClauses -> DrivingClauses
addViaF names v dc = dc { drivingViaF = (names, v) : drivingViaF dc }

emptyDrivingClauses :: DrivingClauses
emptyDrivingClauses = DrivingClauses
  { drivingStock = []
  , drivingNewtype = []
  , drivingAnyclass = []
  , drivingVia = []
  , drivingViaF = []
  }

emptyConfig :: Config
emptyConfig = Config
  { drivingClauses = emptyDrivingClauses
  , exceptions = Map.empty
  }

getConf :: LHsDecl GhcPs -> Maybe Config
getConf (L _ (AnnD _ (HsAnnotation _ _ prov ann_)))
    | ModuleAnnProvenance <- prov = getConfExpr ann_
getConf _ = Nothing

unParTy :: LHsType GhcPs -> HsType GhcPs
unParTy (L _ (HsParTy _ t)) = unParTy t
unParTy (L _ t) = t

getConfExpr :: LHsExpr GhcPs -> Maybe Config
getConfExpr = addModuleAnns_ . unPar where
  addModuleAnns_ :: HsExpr GhcPs -> Maybe Config
#if __GLASGOW_HASKELL__ >= 808
  addModuleAnns_ (ExprWithTySig _ _ t) =
#else
  addModuleAnns_ (ExprWithTySig t _) =
#endif
#if __GLASGOW_HASKELL__ >= 902
    let hsImplicitBody = sig_body . unLoc in
#endif
    case unParTy (hsImplicitBody (hswc_body t)) of
      HsAppTy _ (L _ (HsTyVar _ _ (L _ con))) t'
        | con `eqTyOcc` "Driving" -> Just (mkConfig t' emptyConfig)
      _ -> Nothing
  addModuleAnns_ _ = Nothing

unPar :: LHsExpr GhcPs -> HsExpr GhcPs
unPar (L _ (HsPar _ e)) = unPar e
unPar (L _ e) = e

mkConfig :: LHsType GhcPs -> Config -> Config
mkConfig = go . unParTy where
  go (HsExplicitListTy _ _ ts) conf = foldr (go . unParTy) conf ts
  go (HsAppTy _ (L _ (HsTyVar _ _ (L _ con))) t) conf
    | con `eqTyOcc` "Stock"     = updateDrivingClauses (addStock (extractClasses t)) conf
    | con `eqTyOcc` "Newtype"   = updateDrivingClauses (addNewtype (extractClasses t)) conf
    | con `eqTyOcc` "Anyclass"  = updateDrivingClauses (addAnyclass (extractClasses t)) conf
    | con `eqTyOcc` "NoDriving" = updExceptions t conf
  go (HsAppTy _ (L _ (HsAppTy _ (L _ (HsTyVar _ _ (L _ con))) t)) t') conf
    | con `eqTyOcc` "Via" = updateDrivingClauses (addVia (extractClasses t) t') conf
    | con `eqTyOcc` "ViaF" = updateDrivingClauses (addViaF (extractClasses t) t') conf
  go (HsOpTy _ t (L _ con) t') conf
    | con `eqTyOcc` "Via" = updateDrivingClauses (addVia (extractClasses t) t') conf
    | con `eqTyOcc` "ViaF" = updateDrivingClauses (addViaF (extractClasses t) t') conf
  go _ _ = error "Unrecognized syntax"

eqTyOcc :: RdrName -> String -> Bool
eqTyOcc con cname = rdrNameOcc con == mkTcOcc cname

updExceptions :: LHsType GhcPs -> Config -> Config
updExceptions (L _ (HsParTy _ t)) conf = updExceptions t conf
updExceptions (L _ (HsExplicitTupleTy _ ts)) conf = foldr updExceptions conf ts
updExceptions (L _ (HsAppTy _ (L _ (HsTyVar _ _ (L _ cname))) (L _ (HsTyVar _ _ (L _ tname))))) conf =
  addException tname cname conf
updExceptions _ _ = error "Unrecognized syntax"

extractClasses :: LHsType GhcPs -> [LHsType GhcPs]
extractClasses e = case unLoc e of
  HsParTy _ t -> extractClasses t
  HsExplicitTupleTy _ ts -> ts >>= extractClasses
  _ -> [e]

#if __GLASGOW_HASKELL__ < 902
noLocA :: e -> Located e
noLocA = noLoc

noAnn :: NoExtField
noAnn = NoExtField

hsTypeToHsSigType :: e -> HsImplicitBndrs GhcPs e
hsTypeToHsSigType = mkHsImplicitBndrs
#endif