{-# LANGUAGE LambdaCase, OverloadedStrings, RecordWildCards, ViewPatterns #-}
module TypeLevel.Rewrite (plugin) where

import Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Trans.Writer
import Data.Foldable
import Data.Traversable

-- GHC API
import Coercion (Role(Representational), mkUnivCo)
import Constraint (CtEvidence(ctev_loc), Ct, ctEvExpr, ctLoc, mkNonCanonical)
import GhcPlugins (PredType, SDoc, eqType, fsep, ppr)
import Plugins (Plugin(pluginRecompile, tcPlugin), CommandLineOption, defaultPlugin, purePlugin)
import TcEvidence (EvExpr, EvTerm, evCast)
import TcPluginM (newWanted)
import TcRnTypes
import TyCoRep (UnivCoProvenance(PluginProv))
import TyCon (synTyConDefn_maybe)

import TypeLevel.Rewrite.Internal.ApplyRules
import TypeLevel.Rewrite.Internal.DecomposedConstraint
import TypeLevel.Rewrite.Internal.Lookup
import TypeLevel.Rewrite.Internal.PrettyPrint
import TypeLevel.Rewrite.Internal.TypeEq
import TypeLevel.Rewrite.Internal.TypeRule
import TypeLevel.Rewrite.Internal.TypeTerm

-- printf-debugging:
--import TcPluginM (tcPluginIO)
--import Outputable
----tcPluginIO $ print ("foo", showSDocUnsafe $ ppr foo)


data ReplaceCt = ReplaceCt
  { evidenceOfCorrectness  :: EvTerm
  , replacedConstraint     :: Ct
  , replacementConstraints :: [Ct]
  }

combineReplaceCts
  :: [ReplaceCt]
  -> TcPluginResult
combineReplaceCts replaceCts
  = TcPluginOk (fmap solvedConstraint replaceCts)
               (foldMap replacementConstraints replaceCts)
  where
    solvedConstraint :: ReplaceCt -> (EvTerm, Ct)
    solvedConstraint = (,) <$> evidenceOfCorrectness <*> replacedConstraint


usage
  :: String  -- ^ expected
  -> String  -- ^ actual
  -> TcPluginM a
usage expected actual
  = error $ "usage:\n"
         ++ "  {-# OPTIONS_GHC -fplugin TypeLevel.Rewrite\n"
         ++ "                  -fplugin-opt=TypeLevel.Rewrite:TypeLevel.Append.RightIdentity\n"
         ++ "                  -fplugin-opt=TypeLevel.Rewrite:TypeLevel.Append.RightAssociative #-}\n"
         ++ "Where 'TypeLevel.Append' is a module containing a type synonym named 'RightIdentity':\n"
         ++ "  type RightIdentity as = (as ++ '[]) ~ as\n"
         ++ "Type expressions which match the left of the '~' will get rewritten to the type\n"
         ++ "expression on the right of the '~'. Be careful not to introduce cycles!\n"
         ++ "\n"
         ++ "expected:\n"
         ++ "  " ++ expected ++ "\n"
         ++ "got:\n"
         ++ "  " ++ actual

lookupTypeRules
  :: [CommandLineOption]
  -> TcPluginM [TypeRule]
lookupTypeRules [] = do
  usage (show [ "TypeLevel.Append.RightIdentity" :: String
              , "TypeLevel.Append.RightAssociative"
              ])
        "[]"
lookupTypeRules fullyQualifiedTypeSynonyms = do
  -- ["TypeLevel.Append.RightIdentity", "TypeLevel.Append.RightAssociative"]
  for fullyQualifiedTypeSynonyms $ \fullyQualifiedTypeSynonym -> do
    -- "TypeLevel.Append.RightIdentity"
    case splitLastDot fullyQualifiedTypeSynonym of
      Nothing -> do
        usage (show ("TypeLevel.Append.RightIdentity" :: String))
              (show fullyQualifiedTypeSynonym)
      Just (moduleNameStr, tyConNameStr) -> do
        -- ("TypeLevel.Append", "RightIdentity")
        tyCon <- lookupTyCon moduleNameStr tyConNameStr  -- FIXME: if tyConNameStr is not found in
                                                         -- the module, the error message is poor
        case synTyConDefn_maybe tyCon of
          Nothing -> do
            usage ("type " ++ pprTyCon tyCon ++ " ... = ...")
                  (pprTyCon tyCon ++ " is not a type synonym")
          Just (_tyVars, definition) -> do
            -- ([TyVar "as"], Type "(as ++ '[]) ~ as")
            case toTypeRule_maybe definition of
              Nothing -> do
                usage "... ~ ..."
                      (pprType definition)
              Just typeRule -> do
                -- Rule (TypeTree "(as ++ '[])")
                --      (TypeTree "as")
                pure typeRule


plugin
  :: Plugin
plugin = defaultPlugin
  { tcPlugin = \args -> Just $ TcPlugin
    { tcPluginInit  = lookupTypeRules args
    , tcPluginSolve = solve
    , tcPluginStop  = \_ -> pure ()
    }
  , pluginRecompile = purePlugin
  }


mkErrCtx
  :: SDoc
  -> ErrCtxt
mkErrCtx errDoc = (True, \env -> pure (env, errDoc))

newRuleInducedWanted
  :: Ct
  -> TypeRule
  -> PredType
  -> TcPluginM CtEvidence
newRuleInducedWanted oldCt rule newPredType = do
  let loc = ctLoc oldCt

  -- include the rewrite rule in the error message, if any
  let errMsg = fsep [ "From the typelevel rewrite rule:"
                    , ppr (fromTypeRule rule)
                    ]
  let loc' = pushErrCtxtSameOrigin (mkErrCtx errMsg) loc

  wanted <- newWanted loc' newPredType

  -- ctLoc only copies the "arising from function X" part but not the location
  -- etc., so we need to copy the rest of it manually
  pure $ wanted { ctev_loc = loc' }

solve
  :: [TypeRule]
  -> [Ct]  -- ^ Given constraints
  -> [Ct]  -- ^ Derived constraints
  -> [Ct]  -- ^ Wanted constraints
  -> TcPluginM TcPluginResult
solve _ _ _ [] = do
  pure $ TcPluginOk [] []
solve rules givens _ wanteds = do
  typeSubst <- execWriterT $ do
    for_ givens $ \given -> do
      for_ (asEqualityConstraint given) $ \(lhs, rhs) -> do
        -- lhs ~ rhs
        -- where lhs is typically an expression and rhs is typically a variable
        let var = TypeEq rhs
        let val = toTypeTerm lhs
        tell [(var, val)]

  replaceCts <- execWriterT $ do
    for_ wanteds $ \wanted -> do
      -- wanted => ...
      for_ (asDecomposedConstraint wanted) $ \types -> do
        -- C a b c => ...

        -- C a b c
        let typeTerms = fmap toTypeTerm types
        let predType = fromDecomposeConstraint types

        for_ (applyRules typeSubst rules typeTerms) $ \(rule, typeTerms') -> do
          -- C a' b' c'
          let types' = fmap fromTypeTerm typeTerms'
          let predType' = fromDecomposeConstraint types'

          unless (eqType predType' predType) $ do
            -- co :: C a' b' c'  ~R  C a b c
            let co = mkUnivCo
                       (PluginProv "TypeLevel.Rewrite")
                       Representational
                       predType'
                       predType
            evWanted' <- lift $ newRuleInducedWanted wanted rule predType'
            let wanted' = mkNonCanonical evWanted'
            let futureDict :: EvExpr
                futureDict = ctEvExpr evWanted'
            let replaceCt :: ReplaceCt
                replaceCt = ReplaceCt
                  { evidenceOfCorrectness  = evCast futureDict co
                  , replacedConstraint     = wanted
                  , replacementConstraints = [wanted']
                  }
            tell [replaceCt]
  pure $ combineReplaceCts replaceCts