{-# LANGUAGE
    FlexibleInstances
  , GeneralizedNewtypeDeriving
  , RecordWildCards
  , UnicodeSyntax
  #-}
-- |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
import Control.Monad.State
import Data.Data
import Data.Foldable
import Data.Generics.Aliases hiding (GT)
import Data.Generics.Schemes
import Data.Monoid
import Data.Monoid.Unicode
import Data.Set (Set)
import qualified Data.Set as S
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 (Set α)
    deriving Foldable

-- |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 (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 ImportOp)
      }

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

instance Monoid (Imports ImportOp) where
    mempty
        = Imports (∅)
    mappend (Imports α) (Imports β)
        = Imports (foldl' insertImp α β)

insertImp ∷ Set ImportOp → ImportOp → Set ImportOp
insertImp α qi@(QualifiedImp   {}) = S.insert qi α
insertImp α ui@(UnqualifiedImp {})
    = case find sameMod α of
        Nothing  → S.insert ui α
        Just ui' → S.insert (merge ui') (S.delete ui' α)
    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 s , Just s') → ui { impNames = Just (s ⊕ s') }

-- |@'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 (Imports (S.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
                $ S.singleton (VarName, o)
      in
        RewriteRule pat Unqualify (Imports (S.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
                $ S.singleton (TcClsName, tycl)
      in
        RewriteRule pat Unqualify (Imports (S.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 (Imports (S.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 ImportOp)
rewriteNames rules = flip runState (∅) ∘ gmapM (everywhereM (mkM f))
    where
      f ∷ (Functor m, Monad m) ⇒ Name → StateT (Imports ImportOp) 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 ImportOp) 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)