{-# LANGUAGE LambdaCase 
           , ScopedTypeVariables
           #-}
module Language.Haskell.Tools.Refactor.OrganizeImports (organizeImports) where

import SrcLoc
import Name hiding (Name)
import GHC (Ghc, GhcMonad, lookupGlobalName, TyThing(..), moduleNameString, moduleName)
import qualified GHC
import TyCon
import ConLike
import DataCon
import Outputable (ppr, showSDocUnsafe)

import Control.Reference hiding (element)
import Control.Monad
import Control.Monad.IO.Class
import Data.Function hiding ((&))
import Data.String
import Data.Maybe
import Data.Data
import Data.List
import Data.Generics.Uniplate.Data
import Language.Haskell.Tools.AST
import Language.Haskell.Tools.AST.FromGHC
import Language.Haskell.Tools.AnnTrf.SourceTemplate
import Language.Haskell.Tools.AnnTrf.SourceTemplateHelpers
import Language.Haskell.Tools.PrettyPrint
import Language.Haskell.Tools.Refactor.DebugGhcAST
import Language.Haskell.Tools.AST.Gen
import Language.Haskell.Tools.Refactor.RefactorBase
import Debug.Trace

organizeImports :: forall n . (NamedThing n, Data n) => Ann Module (STWithNames n) -> RefactoredModule n
organizeImports mod
  = element&modImports&annListElems !~ narrowImports usedNames . sortImports $ mod
  where usedNames = map getName $ catMaybes
                                $ map (^? (annotation&semanticInfo&nameInfo))
                                $ (universeBi (mod ^. element&modHead) ++ universeBi (mod ^. element&modDecl) :: [Ann SimpleName (STWithNames n)])
        
-- | Sorts the imports in alphabetical order
sortImports :: [Ann ImportDecl (STWithNames n)] -> [Ann ImportDecl (STWithNames n)]
sortImports = sortBy (ordByOccurrence `on` (^. element&importModule&element))

-- | Modify an import to only import  names that are used.
narrowImports :: forall n . NamedThing n => [GHC.Name] -> [Ann ImportDecl (STWithNames n)] -> Refactor n [Ann ImportDecl (STWithNames n)]
narrowImports usedNames imps = foldM (narrowOneImport usedNames) imps imps 
  where narrowOneImport :: [GHC.Name] -> [Ann ImportDecl (STWithNames n)] -> Ann ImportDecl (STWithNames n) -> Refactor n [Ann ImportDecl (STWithNames n)]
        narrowOneImport names all one =
          (\case Just x -> map (\e -> if e == one then x else e) all
                 Nothing -> delete one all) <$> narrowImport names (map (^. semantics) all) one 
        
narrowImport :: NamedThing n => [GHC.Name] -> [SemanticInfo n] -> Ann ImportDecl (STWithNames n) 
                             -> Refactor n (Maybe (Ann ImportDecl (STWithNames n)))
narrowImport usedNames otherModules imp
  | importIsExact (imp ^. element) 
  = Just <$> (element&importSpec&annJust&element&importSpecList !~ narrowImportSpecs usedNames $ imp)
  | otherwise 
  = if null actuallyImported
      then if length (otherModules ^? traversal&importedModule&filtered (== importedMod) :: [GHC.Module]) > 1 
              then pure Nothing
              else Just <$> (element&importSpec !- replaceWithJust (mkImportSpecList []) $ imp)
      else pure (Just imp)
  where actuallyImported = map getName (fromJust (imp ^? annotation&semanticInfo&importedNames)) `intersect` usedNames
        Just importedMod = imp ^? annotation&semanticInfo&importedModule
    
-- | Narrows the import specification (explicitely imported elements)
narrowImportSpecs :: forall n . NamedThing n => [GHC.Name] -> AnnList IESpec (STWithNames n) -> Refactor n (AnnList IESpec (STWithNames n))
narrowImportSpecs usedNames 
  = (annList&element !~ narrowSpecSubspec usedNames) 
       >=> return . filterList isNeededSpec
  where narrowSpecSubspec :: [GHC.Name] -> IESpec (STWithNames n) -> Refactor n (IESpec (STWithNames n))
        narrowSpecSubspec usedNames spec 
          = do let Just specName = spec ^? ieName&element&simpleName&annotation&semanticInfo&nameInfo
               Just tt <- GHC.lookupName (getName specName)
               let subspecsInScope = case tt of ATyCon tc | not (isClassTyCon tc) 
                                                  -> map getName (tyConDataCons tc) `intersect` usedNames
                                                _ -> usedNames
               ieSubspec&annJust !- narrowImportSubspecs subspecsInScope $ spec
  
        isNeededSpec :: Ann IESpec (STWithNames n) -> Bool
        isNeededSpec ie = 
          -- if the name is used, it is needed
          fmap getName (ie ^? element&ieName&element&simpleName&annotation&semanticInfo&nameInfo) `elem` map Just usedNames
          -- if the name is not used, but some of its constructors are used, it is needed
            || ((ie ^? element&ieSubspec&annJust&element&essList&annList) /= [])
            || (case ie ^? element&ieSubspec&annJust&element of Just SubSpecAll -> True; _ -> False)     
  
narrowImportSubspecs :: NamedThing n => [GHC.Name] -> Ann SubSpec (STWithNames n) -> Ann SubSpec (STWithNames n)
narrowImportSubspecs [] (Ann _ SubSpecAll) = mkSubList []
narrowImportSubspecs _ ss@(Ann _ SubSpecAll) = ss
narrowImportSubspecs usedNames ss@(Ann _ (SubSpecList _)) 
  = element&essList .- filterList (\n -> fmap getName (n ^? element&simpleName&annotation&semanticInfo&nameInfo) `elem` map Just usedNames) $ ss