{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RankNTypes #-}
module Language.Haskell.Refact.Utils.ExactPrint
  (
    replace
  , replaceAnnKey
  , copyAnn

  , setAnnKeywordDP
  , clearPriorComments
  , balanceAllComments
  , locate
  , addEmptyAnn
  , addAnnVal
  , addAnn
  , zeroDP
  , setDP
  , handleParseResult
  , removeAnns
  , synthesizeAnns
  , addNewKeyword
  , addNewKeywords
  ) where

import qualified GHC           as GHC

import qualified Data.Generics as SYB
import Control.Monad

import Language.Haskell.GHC.ExactPrint.Transform
import Language.Haskell.GHC.ExactPrint.Types
import Language.Haskell.GHC.ExactPrint.Utils
import Language.Haskell.Refact.Utils.GhcUtils

import Language.Haskell.Refact.Utils.Monad
import Language.Haskell.Refact.Utils.MonadFunctions

import qualified Data.Map as Map

-- ---------------------------------------------------------------------

-- ++AZ++:TODO: Move this to ghc-exactprint
-- |The annotations are keyed to the constructor, so if we replace a qualified
-- with an unqualified RdrName or vice versa we have to rebuild the key for the
-- appropriate annotation.
replaceAnnKey :: (SYB.Data old,SYB.Data new)
  => GHC.Located old -> GHC.Located new -> Anns -> Anns
replaceAnnKey old new ans =
  case Map.lookup (mkAnnKey old) ans of
    Nothing -> ans
    Just v ->  anns'
      where
        anns1 = Map.delete (mkAnnKey old) ans
        anns' = Map.insert (mkAnnKey new) v anns1


-- ---------------------------------------------------------------------

-- ++AZ++ TODO: migrate this to ghc-exactprint
copyAnn :: (SYB.Data old,SYB.Data new)
  => GHC.Located old -> GHC.Located new -> Anns -> Anns
copyAnn old new ans =
  case Map.lookup (mkAnnKey old) ans of
    Nothing -> ans
    Just v  -> Map.insert (mkAnnKey new) v ans

-- ---------------------------------------------------------------------

-- | Replaces an old expression with a new expression
replace :: AnnKey -> AnnKey -> Anns -> Maybe Anns
replace old new ans = do
  let as = ans
  oldan <- Map.lookup old as
  newan <- Map.lookup new as
  let newan' = Ann
                { annEntryDelta        = annEntryDelta oldan
                -- , annDelta             = annDelta oldan
                -- , annTrueEntryDelta    = annTrueEntryDelta oldan
                , annPriorComments     = annPriorComments oldan
                , annFollowingComments = annFollowingComments oldan
                , annsDP               = moveAnns (annsDP oldan) (annsDP newan)
                , annSortKey           = annSortKey oldan
                , annCapturedSpan      = annCapturedSpan oldan
                }
  return ((\anns -> Map.delete old . Map.insert new newan' $ anns) ans)

-- ---------------------------------------------------------------------

-- | Shift the first output annotation into the correct place
moveAnns :: [(KeywordId, DeltaPos)] -> [(KeywordId, DeltaPos)] -> [(KeywordId, DeltaPos)]
moveAnns [] xs        = xs
moveAnns ((_, dp): _) ((kw, _):xs) = (kw,dp) : xs
moveAnns _ []         = []

-- ---------------------------------------------------------------------

-- |Change the @DeltaPos@ for a given @KeywordId@ if it appears in the
-- annotation for the given item.
setAnnKeywordDP :: (SYB.Data a) => GHC.Located a -> KeywordId -> DeltaPos -> Transform ()
setAnnKeywordDP la kw dp = modifyAnnsT changer
  where
    changer ans = case Map.lookup (mkAnnKey la) ans of
      Nothing -> ans
      Just an -> Map.insert (mkAnnKey la) (an {annsDP = map update (annsDP an)}) ans
    update (kw',dp')
      | kw == kw' = (kw',dp)
      | otherwise = (kw',dp')

-- ---------------------------------------------------------------------

-- |Remove any preceding comments from the given item
clearPriorComments :: (SYB.Data a) => GHC.Located a -> Transform ()
clearPriorComments la = do
  edp <- getEntryDPT la
  modifyAnnsT $ \ans ->
    case Map.lookup (mkAnnKey la) ans of
      Nothing -> ans
      Just an -> Map.insert (mkAnnKey la) (an {annPriorComments = [] }) ans
  setEntryDPT la edp

-- ---------------------------------------------------------------------

balanceAllComments :: SYB.Data a => GHC.Located a -> Transform (GHC.Located a)
balanceAllComments la
  -- Must be top-down
  = everywhereM' (SYB.mkM inMod
                     `SYB.extM` inExpr
                     `SYB.extM` inMatch
                     `SYB.extM` inStmt
                   ) la
  where
    inMod :: GHC.ParsedSource -> Transform (GHC.ParsedSource)
    inMod m = doBalance m

    inExpr :: GHC.LHsExpr GHC.RdrName -> Transform (GHC.LHsExpr GHC.RdrName)
    inExpr e = doBalance e

    inMatch :: (GHC.LMatch GHC.RdrName (GHC.LHsExpr GHC.RdrName)) -> Transform (GHC.LMatch GHC.RdrName (GHC.LHsExpr GHC.RdrName))
    inMatch m = doBalance m

    inStmt :: GHC.LStmt GHC.RdrName (GHC.LHsExpr GHC.RdrName) -> Transform (GHC.LStmt GHC.RdrName (GHC.LHsExpr GHC.RdrName))
    inStmt s = doBalance s

    -- |Balance all comments between adjacent decls, as well as pushing all
    -- trailing comments to the right place.
    {-
    e.g., for

        foo = do
          return x
            where
               x = ['a'] -- do

        bar = undefined

    the "-- do" comment must end up in the trailing comments for "x = ['a']"
    -}
    doBalance t = do
      decls <- hsDecls t
      let
        go [] = return []
        go [x] = return [x]
        go (x1:x2:xs) = do
          balanceComments x1 x2
          go (x2:xs)
      _ <- go decls
      -- replaceDecls t decls'
      unless (null decls) $ moveTrailingComments t (last decls)
      return t

--This generates a unique location and wraps the given ast chunk with that location
--Also adds an empty annotation at that location
locate :: (SYB.Data a) => a -> RefactGhc (GHC.Located a)
locate ast = do
  loc <- liftT uniqueSrcSpanT
  let res = (GHC.L loc ast)
  addEmptyAnn res
  return res

--Adds an empty annotation at the provided location
addEmptyAnn :: (SYB.Data a) => GHC.Located a -> RefactGhc ()
addEmptyAnn a = addAnn a annNone

--Adds an "AnnVal" annotation at the provided location
addAnnVal :: (SYB.Data a) => GHC.Located a -> RefactGhc ()
addAnnVal a = addAnn a valAnn
  where valAnn = annNone {annEntryDelta = DP (0,1), annsDP = [(G GHC.AnnVal, DP (0,0))]}

--Adds the given annotation at the provided location
addAnn :: (SYB.Data a) => GHC.Located a -> Annotation -> RefactGhc ()
addAnn a ann = do
  currAnns <- fetchAnnsFinal
  let k = mkAnnKey a
  setRefactAnns $ Map.insert k ann currAnns

--Sets the entry delta position of an ast chunk
setDP :: (SYB.Data a) => DeltaPos -> GHC.Located a -> RefactGhc ()
setDP dp ast = do
  currAnns <- fetchAnnsFinal
  let k = mkAnnKey ast
      mv = Map.lookup k currAnns
  case mv of
    Nothing -> return ()
    Just v -> addAnn ast (v {annEntryDelta = dp})

--Resets the given AST chunk's delta position to zero.
zeroDP :: (SYB.Data a) => GHC.Located a -> RefactGhc ()
zeroDP = setDP (DP (0,0))

--This just pulls out the successful result from an exact print parser or throws an error if the parse was unsuccessful.
handleParseResult :: String -> Either (GHC.SrcSpan, String) (Anns, a) -> RefactGhc (Anns, a)
handleParseResult msg e = case e of
  (Left (_, errStr)) -> error $ "The parse from: " ++ msg ++ " with error: " ++ errStr
  (Right res) -> return res

-- This creates an empty annotation for every located item where an annotation does not already exist in the given AST chunk
synthesizeAnns :: (SYB.Data a) => a -> RefactGhc a
synthesizeAnns = generic `SYB.ext2M` located
  where generic :: SYB.Data a => a -> RefactGhc a
        generic a = do
          _ <- SYB.gmapM synthesizeAnns a
          return a
        located :: (SYB.Data b, SYB.Data loc) => GHC.GenLocated loc b -> RefactGhc (GHC.GenLocated loc b)
        located b@(GHC.L ss a) = case SYB.cast ss of
          Just (s :: GHC.SrcSpan) -> do
            --logm $ "Located found: " ++ (show $ toConstr a)
            anns <- fetchAnnsFinal
            let castRes = (GHC.L s a)
                ann = getAnnotationEP castRes anns
            --logm $ "Found ann: " ++ show ann
            case ann of
              Nothing -> do
                --logm "No ann found for located item"
                let newKey = mkAnnKey castRes
                    newAnns = Map.insert newKey annNone anns
                setRefactAnns newAnns
                return ()
              _ -> return ()
            _ <- SYB.gmapM synthesizeAnns b
            return b
          Nothing ->
            return b

-- This removes all the annotations associated with the given AST chunk.
removeAnns :: (SYB.Data a) => a -> RefactGhc a
removeAnns = generic `SYB.ext2M` located
  where generic :: SYB.Data a => a -> RefactGhc a
        generic a = do
          _ <- SYB.gmapM synthesizeAnns a
          return a
        located :: (SYB.Data b, SYB.Data loc) => GHC.GenLocated loc b -> RefactGhc (GHC.GenLocated loc b)
        located b@(GHC.L ss a) = case SYB.cast ss of
          Just (s :: GHC.SrcSpan) -> do
            anns <- fetchAnnsFinal
            let k = mkAnnKey (GHC.L s a)
            logm $ "Deleting ann at: " ++ (show s)
            setRefactAnns $ Map.delete k anns
            _ <- SYB.gmapM removeAnns b
            return b
          Nothing -> return b

--This takes in a located ast chunk and adds the provided keyword and delta position into the annsDP list
--If there is not annotation associated with the chunk nothing happens
addNewKeyword :: (SYB.Data a) => (KeywordId, DeltaPos) -> GHC.Located a -> RefactGhc ()
addNewKeyword entry a = do
  anns <- liftT getAnnsT
  let key = mkAnnKey a
      mAnn = Map.lookup key anns
  case mAnn of
    Nothing -> return ()
    (Just ann) -> do
      let newAnn = ann{annsDP = (entry:(annsDP ann))}
      setRefactAnns $ Map.insert key newAnn anns

addNewKeywords :: (SYB.Data a) => [(KeywordId, DeltaPos)] -> GHC.Located a -> RefactGhc ()
addNewKeywords entries a = mapM_ ((flip addNewKeyword) a) entries