content: {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE RecordWildCards #-}
  
  module Ormolu.Imports.Grouping
    ( Import (..),
      prepareExistingGroups,
      groupImports,
    )
  where
  
  import Data.Bifunctor (Bifunctor (..))
  import Data.Foldable (toList)
  import Data.Function (on)
  import Data.List (groupBy, minimumBy, sortOn)
  import Data.List.NonEmpty (NonEmpty)
  import Data.List.NonEmpty qualified as NonEmpty
  import Data.Set (Set)
  import Data.Set qualified as Set
  import Distribution.ModuleName qualified as Cabal
  import GHC.Hs (GhcPs, getLocA)
  import Language.Haskell.Syntax (LImportDecl, ModuleName, moduleNameString)
  import Ormolu.Config (ImportGroup (..), ImportGroupRule (..), ImportGrouping (..), ImportModuleMatcher (..))
  import Ormolu.Config qualified as Config
  import Ormolu.Utils (ghcModuleNameToCabal, groupBy', separatedByBlank)
  import Ormolu.Utils.Glob (matchesGlob)
  
  newtype ImportGroups = ImportGroups (NonEmpty ImportGroup)
  
  data Import = Import
    { importName :: ModuleName,
      importQualified :: Bool
    }
  
  importGroupSingleStrategy :: ImportGroups
  importGroupSingleStrategy =
    ImportGroups $
      pure
        ImportGroup
          { igName = Nothing,
            igRules = pure matchAllImportRule
          }
  
  importGroupByQualifiedStrategy :: ImportGroups
  importGroupByQualifiedStrategy =
    ImportGroups $
      NonEmpty.fromList
        [ ImportGroup
            { igName = Nothing,
              igRules = pure $ withUnqualifiedOnly matchAllImportRule
            },
          ImportGroup
            { igName = Nothing,
              igRules = pure $ withQualifiedOnly matchAllImportRule
            }
        ]
  
  importGroupByScopeStrategy :: ImportGroups
  importGroupByScopeStrategy =
    ImportGroups $
      NonEmpty.fromList
        [ ImportGroup
            { igName = Nothing,
              igRules = pure matchAllImportRule
            },
          ImportGroup
            { igName = Nothing,
              igRules = pure matchLocalModulesRule
            }
        ]
  
  importGroupByScopeThenQualifiedStrategy :: ImportGroups
  importGroupByScopeThenQualifiedStrategy =
    ImportGroups $
      NonEmpty.fromList
        [ ImportGroup
            { igName = Nothing,
              igRules = pure $ withQualified matchModule
            }
        | matchModule <- [matchAllImportRule, matchLocalModulesRule],
          withQualified <- [withUnqualifiedOnly, withQualifiedOnly]
        ]
  
  groupsFromConfig :: Config.ImportGrouping -> ImportGroups
  groupsFromConfig =
    \case
      Config.ImportGroupLegacy -> importGroupSingleStrategy
      Config.ImportGroupPreserve -> importGroupSingleStrategy
      Config.ImportGroupSingle -> importGroupSingleStrategy
      Config.ImportGroupByQualified -> importGroupByQualifiedStrategy
      Config.ImportGroupByScope -> importGroupByScopeStrategy
      Config.ImportGroupByScopeThenQualified -> importGroupByScopeThenQualifiedStrategy
      Config.ImportGroupCustom igs -> ImportGroups igs
  
  matchAllImportRule :: ImportGroupRule
  matchAllImportRule =
    ImportGroupRule
      { igrModuleMatcher = MatchAllModules,
        igrQualifiedMatcher = Config.MatchBothQualifiedAndUnqualified,
        igrPriority = Config.matchAllRulePriority
      }
  
  matchLocalModulesRule :: ImportGroupRule
  matchLocalModulesRule =
    ImportGroupRule
      { igrModuleMatcher = MatchLocalModules,
        igrQualifiedMatcher = Config.MatchBothQualifiedAndUnqualified,
        igrPriority = Config.matchLocalRulePriority
      }
  
  withQualifiedOnly :: ImportGroupRule -> ImportGroupRule
  withQualifiedOnly ImportGroupRule {..} =
    ImportGroupRule
      { igrQualifiedMatcher = Config.MatchQualifiedOnly,
        ..
      }
  
  withUnqualifiedOnly :: ImportGroupRule -> ImportGroupRule
  withUnqualifiedOnly ImportGroupRule {..} =
    ImportGroupRule
      { igrQualifiedMatcher = Config.MatchUnqualifiedOnly,
        ..
      }
  
  matchesRule :: Set Cabal.ModuleName -> Import -> ImportGroupRule -> Bool
  matchesRule localMods Import {..} ImportGroupRule {..} = matchesModules && matchesQualified
    where
      matchesModules = case igrModuleMatcher of
        MatchAllModules -> True
        MatchLocalModules -> ghcModuleNameToCabal importName `Set.member` localMods
        MatchGlob gl -> moduleNameString importName `matchesGlob` gl
      matchesQualified = case igrQualifiedMatcher of
        Config.MatchQualifiedOnly -> importQualified
        Config.MatchUnqualifiedOnly -> not importQualified
        Config.MatchBothQualifiedAndUnqualified -> True
  
  prepareExistingGroups :: ImportGrouping -> Bool -> [LImportDecl GhcPs] -> [[LImportDecl GhcPs]]
  prepareExistingGroups ig respectful =
    case ig of
      ImportGroupPreserve -> preserveGroups
      ImportGroupLegacy | respectful -> preserveGroups
      _ -> flattenGroups
    where
      preserveGroups = map toList . groupBy' (\x y -> not $ separatedByBlank getLocA x y)
      flattenGroups = pure
  
  groupImports :: forall x. ImportGrouping -> Set Cabal.ModuleName -> (x -> Import) -> [x] -> [[x]]
  groupImports ig localModules fToImport = regroup . fmap (breakTies . matchRules)
    where
      ImportGroups igs = groupsFromConfig ig
  
      indexedGroupRules :: [(Int, [ImportGroupRule])]
      indexedGroupRules = zip [0 ..] (toList . igRules <$> toList igs)
  
      matchRules :: x -> ([(Int, [ImportGroupRule])], x)
      matchRules x =
        let imp = fToImport x
            testRule (_, rules) = any (matchesRule localModules imp) rules
         in (filter testRule indexedGroupRules, x)
  
      breakTies :: ([(Int, [ImportGroupRule])], x) -> (Int, x)
      breakTies ([], x) =
        (maxBound, x) -- Any non-matched imports will be in the last group
      breakTies (matches, x) =
        (fst . minimumBy (compare `on` snd) $ second (minimum . fmap igrPriority) <$> matches, x)
  
      regroup :: [(Int, x)] -> [[x]]
      regroup = fmap (fmap snd) . groupBy ((==) `on` fst) . sortOn fst
  