{-# LANGUAGE
  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.

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

import GHC.Plugins hiding (Type)
import GHC.Hs

-- * 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) }

drivingMod :: Endo HsModule
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 = 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 (drivingStock clauses)
      ++ mkDerivingClauses NewtypeStrategy (drivingNewtype clauses)
      ++ mkDerivingClauses AnyclassStrategy (drivingAnyclass clauses)
      ++ (mkDerivingViaClauses =<< drivingVia clauses)
      ++ (mkDerivingViaClauses =<< (fmap . fmap) applyToTyname (drivingViaF clauses))
    applyToTyname f = noLoc (HsAppTy NoExtField f (noLoc (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 = NoExtField
      , deriv_clause_strategy = Just (noLoc strat)
      , deriv_clause_tys = noLoc (map mkHsImplicitBndrs cls)
      })
  ]

mkDerivingViaClauses :: ([LHsType GhcPs], LHsType GhcPs) -> [LHsDerivingClause GhcPs]
mkDerivingViaClauses (cls, v) = mkDerivingClauses (ViaStrategy (mkHsImplicitBndrs v)) cls

hsTyVar :: RdrName -> HsType GhcPs
hsTyVar = HsTyVar NoExtField NotPromoted . noLoc

-- *** 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 e = addModuleAnns_ (unPar e) where
  addModuleAnns_ (ExprWithTySig _ _ t) =
    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]