{-# LANGUAGE RecordWildCards #-}
-- Copyright (c) Facebook, Inc. and its affiliates.
--
-- This source code is licensed under the MIT license found in the
-- LICENSE file in the root directory of this source tree.
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
module Retrie.Rewrites
  ( RewriteSpec(..)
  , QualifiedName
  , parseRewriteSpecs
  , parseQualified
  , parseAdhocs
  ) where

import Control.Exception
import qualified Data.Map as Map
import Data.Maybe
import Data.Data hiding (Fixity)
import qualified Data.Text as Text
import Data.Traversable
import System.FilePath

import Retrie.CPP
import Retrie.ExactPrint
import Retrie.Fixity
#if __GLASGOW_HASKELL__ < 904
import Retrie.GHC
#else
import Retrie.GHC hiding (Pattern)
#endif
import Retrie.Rewrites.Function
import Retrie.Rewrites.Patterns
import Retrie.Rewrites.Rules
import Retrie.Rewrites.Types
import Retrie.Types
import Retrie.Universe
import Retrie.Util

-- | A qualified name. (e.g. @"Module.Name.functionName"@)
type QualifiedName = String

-- | Possible ways to specify rewrites to 'parseRewrites'.
data RewriteSpec
  = Adhoc String
    -- ^ Equation in RULES-format. (e.g. @"forall x. succ (pred x) = x"@)
    -- Will be applied left-to-right.
  | AdhocPattern String
    -- ^ Equation in pattern-synonym format, _without_ the keyword 'pattern'.
  | AdhocType String
    -- ^ Equation in type-synonym format, _without_ the keyword 'type'.
  | Fold QualifiedName
    -- ^ Fold a function definition. The inverse of unfolding/inlining.
    -- Replaces instances of the function body with calls to the function.
  | RuleBackward QualifiedName
    -- ^ Apply a GHC RULE right-to-left.
  | RuleForward QualifiedName
    -- ^ Apply a GHC RULE left-to-right.
  | TypeBackward QualifiedName
    -- ^ Apply a type synonym right-to-left.
  | TypeForward QualifiedName
    -- ^ Apply a type synonym left-to-right.
  | Unfold QualifiedName
    -- ^ Unfold, or inline, a function definition.
  | PatternForward QualifiedName
    -- ^ Unfold a pattern synonym
  | PatternBackward QualifiedName
    -- ^ Fold a pattern synonym, replacing instances of the rhs with the synonym


data ClassifiedRewrites = ClassifiedRewrites
  { adhocRules :: [String]
  , adhocPatterns :: [String]
  , adhocTypes :: [String]
  , fileBased :: [(FilePath, [(FileBasedTy,[(FastString, Direction)])])]
  }

instance Monoid ClassifiedRewrites where
  mempty = ClassifiedRewrites [] [] [] []

instance Semigroup ClassifiedRewrites where
  ClassifiedRewrites a b c d <> ClassifiedRewrites a' b' c' d' =
    ClassifiedRewrites (a <> a') (b <> b') (c <> c') (d <> d')

parseRewriteSpecs
  :: LibDir
  -> (FilePath -> IO (CPP AnnotatedModule))
  -> FixityEnv
  -> [RewriteSpec]
  -> IO [Rewrite Universe]
parseRewriteSpecs libdir parser fixityEnv specs = do
  ClassifiedRewrites{..} <- mconcat <$> sequence
    [ case spec of
        Adhoc rule -> return mempty{adhocRules = [rule]}
        AdhocPattern pSyn -> return mempty{adhocPatterns = [pSyn]}
        AdhocType tySyn -> return mempty{adhocTypes = [tySyn]}
        Fold name -> mkFileBased FoldUnfold RightToLeft name
        RuleBackward name -> mkFileBased Rule RightToLeft name
        RuleForward name -> mkFileBased Rule LeftToRight name
        TypeBackward name -> mkFileBased Type RightToLeft name
        TypeForward name -> mkFileBased Type LeftToRight name
        PatternBackward name -> mkFileBased Pattern RightToLeft name
        PatternForward name -> mkFileBased Pattern LeftToRight name
        Unfold name -> mkFileBased FoldUnfold LeftToRight name
    | spec <- specs
    ]
  fbRewrites <- parseFileBased libdir parser fileBased
  adhocExpressionRewrites <- parseAdhocs libdir fixityEnv adhocRules
  -- debugPrint Loud "parseRewriteSpecs" (["adhocExpressionRewrites:" ++ show adhocRules]  ++ map (\r -> showAst ((astA . qPattern) r)) adhocExpressionRewrites)
  adhocTypeRewrites <- parseAdhocTypes libdir fixityEnv adhocTypes
  -- debugPrint Loud "parseRewriteSpecs" (["adhocTypeRewrites:"] ++ map (\r -> showAst ((astA . qPattern) r)) adhocTypeRewrites)
  adhocPatternRewrites <- parseAdhocPatterns libdir fixityEnv adhocPatterns
  -- debugPrint Loud "parseRewriteSpecs" (["adhocPatternRewrites:"] ++ map (\r -> showAst ((astA . qPattern) r)) adhocPatternRewrites)
  return $
    fbRewrites ++
    adhocExpressionRewrites ++
    adhocTypeRewrites ++
    adhocPatternRewrites
  where
    mkFileBased ty dir name =
      case parseQualified name of
        Left err -> throwIO $ ErrorCall $ "parseRewriteSpecs: " ++ err
        Right (fp, fs) -> return mempty{fileBased = [(fp, [(ty, [(fs, dir)])])]}

data FileBasedTy = FoldUnfold | Rule | Type | Pattern
  deriving (Eq, Ord)

parseFileBased
  :: LibDir
  -> (FilePath -> IO (CPP AnnotatedModule))
  -> [(FilePath, [(FileBasedTy, [(FastString, Direction)])])]
  -> IO [Rewrite Universe]
parseFileBased _ _ [] = return []
parseFileBased libdir parser specs = concat <$> mapM (uncurry goFile) (gather specs)
  where
    gather :: Ord a => [(a,[b])] -> [(a,[b])]
    gather = Map.toList . Map.fromListWith (++)

    goFile
      :: FilePath
      -> [(FileBasedTy, [(FastString, Direction)])]
      -> IO [Rewrite Universe]
    goFile fp rules = do
      cpp <- parser fp
      concat <$> mapM (uncurry $ constructRewrites libdir cpp) (gather rules)

parseAdhocs :: LibDir -> FixityEnv -> [String] -> IO [Rewrite Universe]
parseAdhocs _ _ [] = return []
parseAdhocs libdir fixities adhocs = do
  -- debugPrint Loud "parseAdhocs:adhocs" adhocs
  -- debugPrint Loud "parseAdhocs:adhocRules" (map show adhocRules)
  cpp <-
    parseCPP (parseContent libdir fixities "parseAdhocs") (Text.unlines adhocRules)
  -- debugPrint Loud "parseAdhocs:cpp" [showCpp cpp]
  constructRewrites libdir cpp Rule adhocSpecs
  where
    -- In search mode, there is no need to specify a right-hand side, but we
    -- need one to parse as a RULE, so add it if necessary.
    addRHS s
      | '=' `elem` s = s
      | otherwise = s ++ " = undefined"
    (adhocSpecs, adhocRules) = unzip
      [ ( (mkFastString nm, LeftToRight)
        , "{-# RULES \"" <> Text.pack nm <> "\" " <> Text.pack s <> " #-}"
        )
      | (i,s) <- zip [1..] $ map addRHS adhocs
      , let nm = "adhoc" ++ show (i::Int)
      ]


showCpp :: (Data ast, ExactPrint ast) => CPP (Annotated ast) -> String
showCpp (NoCPP c) = showAstA c
showCpp (CPP{}) = "CPP{}"

parseAdhocTypes :: LibDir -> FixityEnv -> [String] -> IO [Rewrite Universe]
parseAdhocTypes _ _ [] = return []
parseAdhocTypes libdir fixities tySyns = do
  print adhocTySyns
  cpp <-
    parseCPP (parseContent libdir fixities "parseAdhocTypes") (Text.unlines adhocTySyns)
  constructRewrites libdir cpp Type adhocSpecs
  where
    (adhocSpecs, adhocTySyns) = unzip
      [ ( (mkFastString nm, LeftToRight), "type " <> Text.pack s)
      | s <- tySyns
      , Just nm <- [listToMaybe $ words s]
      ]

parseAdhocPatterns :: LibDir -> FixityEnv -> [String] -> IO [Rewrite Universe]
parseAdhocPatterns _ _ [] = return []
parseAdhocPatterns libdir fixities patSyns = do
  cpp <-
    parseCPP (parseContent libdir fixities "parseAdhocPatterns")
             (Text.unlines $ pragma : adhocPatSyns)
  constructRewrites libdir cpp Pattern adhocSpecs
  where
    pragma = "{-# LANGUAGE PatternSynonyms #-}"
    (adhocSpecs, adhocPatSyns) = unzip
      [ ( (mkFastString nm, LeftToRight), "pattern " <> Text.pack s)
      | s <- patSyns
      , Just nm <- [listToMaybe $ words s]
      ]

constructRewrites
  :: LibDir
  -> CPP AnnotatedModule
  -> FileBasedTy
  -> [(FastString, Direction)]
  -> IO [Rewrite Universe]
constructRewrites libdir cpp ty specs = do
  cppM <- traverse (tyBuilder libdir ty specs) cpp
  let
    names = nonDetEltsUniqSet $ mkUniqSet $ map fst specs

    nameOf FoldUnfold = "definition"
    nameOf Rule = "rule"
    nameOf Type = "type synonym"
    nameOf Pattern = "pattern synonym"

    m = foldr (plusUFM_C (++)) emptyUFM cppM

  fmap concat $ forM names $ \fs ->
    case lookupUFM m fs of
      Nothing ->
        fail $ "could not find " ++ nameOf ty ++ " named " ++ unpackFS fs
      Just rrs -> do
        -- debugPrint Loud "constructRewrites:cppM" ["enter"]
        return rrs

tyBuilder
  :: LibDir
  -> FileBasedTy
  -> [(FastString, Direction)]
  -> AnnotatedModule
#if __GLASGOW_HASKELL__ < 900
  -> IO (UniqFM [Rewrite Universe])
#else
  -> IO (UniqFM FastString [Rewrite Universe])
#endif
tyBuilder libdir FoldUnfold specs am = promote <$> dfnsToRewrites libdir specs am
tyBuilder _libdir Rule specs am = promote <$> rulesToRewrites specs am
tyBuilder _libdir Type specs am = promote <$> typeSynonymsToRewrites specs am
tyBuilder libdir Pattern specs am = patternSynonymsToRewrites libdir specs am

#if __GLASGOW_HASKELL__ < 900
promote :: Matchable a => UniqFM [Rewrite a] -> UniqFM [Rewrite Universe]
#else
promote :: Matchable a => UniqFM k [Rewrite a] -> UniqFM k [Rewrite Universe]
#endif
promote = fmap (map toURewrite)

parseQualified :: String -> Either String (FilePath, FastString)
parseQualified [] = Left "qualified name is empty"
parseQualified fqName =
  case span isHsSymbol reversed of
    (_,[]) -> mkError "unqualified operator name"
    ([],_) ->
      case span (/='.') reversed of
        (_,[]) -> mkError "unqualified function name"
        (rname,_:rmod) -> mkResult (reverse rmod) (reverse rname)
    (rop,rmod) ->
      case reverse rop of
        '.':op -> mkResult (reverse rmod) op
        _ -> mkError "malformed qualified operator"
  where
    reversed = reverse fqName
    mkError str = Left $ str ++ ": " ++ fqName
    mkResult moduleNameStr occNameStr = Right
      -- 'moduleNameSlashes' gives us system-dependent path separator
      ( moduleNameSlashes (mkModuleName moduleNameStr) <.> "hs"
      , mkFastString occNameStr
      )

isHsSymbol :: Char -> Bool
isHsSymbol = (`elem` symbols)
  -- see https://www.haskell.org/onlinereport/lexemes.html
  where
    symbols :: String
    symbols = "!#$%&*+./<=>?@\\^|-~"