{-# LANGUAGE
    FlexibleInstances
  , GeneralizedNewtypeDeriving
  , MultiParamTypeClasses
  , RecordWildCards
  , TemplateHaskell
  , UnicodeSyntax
  #-}
{-# OPTIONS_GHC -fno-warn-missing-methods #-}
-- |An internal module for rewriting 'Name's in Template Haskell AST.
module Network.HTTP.Lucu.Implant.Rewrite
    ( NamePat(..)
    , RewriteOp(..)

    , Imports
    , ImportOp(..)

    , Rules
    , RewriteRule(..)
    , qualifyAll
    , unqualify
    , unqualifyIn
    , unqualifyAll

    , rewriteNames
    )
    where
import Control.Applicative hiding (empty)
import Control.Monad.State
import Data.Collections
import Data.Collections.BaseInstances ()
import qualified Data.Collections.Newtype.TH as C
import Data.Data
import Data.Generics.Aliases hiding (GT)
import Data.Generics.Schemes
import Data.Monoid
import Data.Monoid.Unicode
import qualified Data.Set as S (Set)
import Language.Haskell.TH.Syntax
import Prelude
import Prelude.Unicode

-- |Pattern for 'Name's. 'Just' represents a perfect matching pattern,
-- and 'Nothing' represensts a wildcard.
data NamePat
    = NamePat !(Maybe ModName) !(Maybe OccName)

-- |Instruction for rewriting 'Name's.
data RewriteOp
    = Identity
    | Unqualify
    | Qualify !ModName

-- |A 'Set' of modules and names to be imported.
newtype Imports = Imports (S.Set ImportOp)

-- |Instruction for declaring module imports.
data ImportOp
    = -- |> import qualified M as A
      QualifiedImp {
        impModule ∷ !ModName
      , impAlias  ∷ !ModName
      }
      -- |> import M
      --
      -- or
      --
      -- > import M (a, b, c, ...)
    | UnqualifiedImp {
        impModule ∷ !ModName
      , impNames  ∷ !(Maybe (S.Set (NameSpace, OccName)))
      }
    deriving Eq

-- |List of 'RewriteRule's.
type Rules = [RewriteRule]

-- |Instruction for rewriting 'Name's and declaring module imports.
data RewriteRule
    = RewriteRule {
        rrPat  ∷ !NamePat
      , rrOp   ∷ !RewriteOp
      , rrImps ∷ !Imports
      }

C.derive [d| instance Foldable   Imports ImportOp
             instance Collection Imports ImportOp
             instance Map        Imports ImportOp ()
             instance Set        Imports ImportOp
             instance SortingCollection Imports ImportOp
           |]

-- |@'insert' imp@ merges @imp@ with an existing one if any.
instance Unfoldable Imports ImportOp where
    insert qi@(QualifiedImp   {}) (Imports s) = Imports $ insert qi s
    insert ui@(UnqualifiedImp {}) (Imports s)
        = case find sameMod s of
            Nothing  → Imports $ insert ui s
            Just ui' → Imports $ insert (merge ui') (delete ui' s)
        where
          sameMod ∷ ImportOp → Bool
          sameMod ui'@(UnqualifiedImp {})
              = impModule ui ≡ impModule ui'
          sameMod _
              = False

          merge ∷ ImportOp → ImportOp
          merge ui'
              = case (impNames ui, impNames ui') of
                  (Nothing, _       ) → ui
                  (_      , Nothing ) → ui'
                  (Just ns, Just ns') → ui { impNames = Just (ns ⊕ ns') }

    empty     = Imports empty
    singleton = Imports ∘ singleton

instance Monoid Imports where
    mempty  = empty
    mappend = insertMany

instance Ord ImportOp where
    α `compare` β
        | impModule α < impModule β = LT
        | impModule α > impModule β = GT
        | otherwise
            = case (α, β) of
                (QualifiedImp   {}, QualifiedImp   {})
                    → impAlias α `compare` impAlias β
                (QualifiedImp   {}, _                )
                    → GT
                (UnqualifiedImp {}, UnqualifiedImp {})
                    → impNames α `compare` impNames β
                (UnqualifiedImp {}, _                )
                    → LT

-- |@'qualifyAll' module alias@: qualify every symbols defined in
-- @module@ with @alias@.
qualifyAll ∷ String → String → RewriteRule
qualifyAll m a
    = let pat = NamePat (Just (mkModName m)) Nothing
          rop = Qualify (mkModName a)
          iop = QualifiedImp (mkModName m) (mkModName a)
      in
        RewriteRule pat rop (singleton iop)

-- |@'unqualify' name module@: unqualify the symbol @name@ with
-- importing @module@.
unqualify ∷ Name → String → RewriteRule
unqualify (Name o _) m
    = let pat = NamePat Nothing (Just o)
          iop = UnqualifiedImp (mkModName m) ∘ Just
                $ singleton (VarName, o)
      in
        RewriteRule pat Unqualify (singleton iop)

-- |@'unqualifyIn' name tycl module@: unqualify a constructor, field
-- name, or whatever resides in the type or class @tycl@ with
-- importing @module@.
unqualifyIn ∷ Name → Name → String → RewriteRule
unqualifyIn (Name name _) (Name tycl _) m
    = let pat = NamePat Nothing (Just name)
          iop = UnqualifiedImp (mkModName m) ∘ Just
                $ singleton (TcClsName, tycl)
      in
        RewriteRule pat Unqualify (singleton iop)

-- |@'unqualifyAll' origMod impMod@: unqualify every symbols
-- defined in @origMod@ with importing @impMod@.
unqualifyAll ∷ String → String → RewriteRule
unqualifyAll origMod impMod
    = let pat = NamePat (Just (mkModName origMod)) Nothing
          iop = UnqualifiedImp (mkModName impMod) Nothing
      in
        RewriteRule pat Unqualify (singleton iop)

-- |@'rewriteNames' rules d@ rewrites each and every 'Name's included
-- in @d@ according to the name-rewriting @rules@ while at the same
-- time building a set of modules to be imported.
rewriteNames ∷ Data d ⇒ Rules → d → (d, Imports)
rewriteNames rules = flip runState (∅) ∘ gmapM (everywhereM (mkM f))
    where
      f ∷ (Functor m, Monad m) ⇒ Name → StateT Imports m Name
      f n = case findRule rules n of
              Nothing → fail $ "No rules matches to name: " ⧺ showName n
              Just r  → applyRule r n

findRule ∷ Rules → Name → Maybe RewriteRule
findRule _  (Name _  NameS       ) = Just identityRule
findRule rs (Name o (NameQ     m)) = find (matchPat m o ∘ rrPat) rs
findRule _  (Name _ (NameU _    )) = Just identityRule
findRule rs (Name o (NameG _ _ m)) = find (matchPat m o ∘ rrPat) rs
findRule _  _                      = Nothing

identityRule ∷ RewriteRule
identityRule = RewriteRule {
                 rrPat  = NamePat Nothing Nothing
               , rrOp   = Identity
               , rrImps = (∅)
               }

matchPat ∷ ModName → OccName → NamePat → Bool
matchPat m o (NamePat mp op)
    = maybe True (≡ m) mp ∧ maybe True (≡ o) op

applyRule ∷ (Functor m, Monad m)
          ⇒ RewriteRule
          → Name
          → StateT Imports m Name
applyRule (RewriteRule {..}) n
    = modify (⊕ rrImps) *> pure (rewrite rrOp n)

rewrite ∷ RewriteOp → Name → Name
rewrite Identity    n          = n
rewrite Unqualify   (Name o _) = Name o NameS
rewrite (Qualify m) (Name o _) = Name o (NameQ m)