{-# LANGUAGE OverloadedStrings #-}

module Aihc.Cli.Repl
  ( ReplError (..),
    ReplSettings (..),
    ReplSession (..),
    ReplStep (..),
    defaultReplSettings,
    evaluateExpression,
    handleReplInput,
    loadReplSession,
    replCompletion,
    runRepl,
  )
where

import Aihc.Cli.PackageInterface
  ( PackageInterface (..),
    PackageInterfaceBinding (..),
    PackageInterfaceTcModule (..),
    packageInterfaceExports,
    readPackageInterface,
  )
import Aihc.Fc (DesugarConfig (..), DesugarResult (..), FcModuleId (..), FcProgram (..), desugarModuleWithInterface, evalProgramBinding, mergePrograms, renderProgram, renderValue)
import Aihc.Parser (ParseResult (..), ParserConfig (..), defaultConfig, parseExpr, parseModule)
import Aihc.Parser.Shorthand (Shorthand (..))
import Aihc.Parser.Syntax
  ( Decl (..),
    Expr,
    ImportDecl (..),
    Match (..),
    MatchHeadForm (..),
    Module (..),
    ModuleHead (..),
    NameType (..),
    Rhs (..),
    ValueDecl (..),
    mkUnqualifiedName,
    nameQualifier,
    nameText,
    qualifyName,
    unqualifiedNameFromText,
  )
import Aihc.Parser.Syntax qualified as Syntax
import Aihc.Resolve (ModuleExports, ModuleKey (..), Package (..), PackageId (..), ResolveError (..), ResolveResult (..), ResolvedName (..), Scope (..), extractInterface, resolveWithDeps)
import Aihc.Tc
  ( InstanceInfo,
    TcBindingResult (..),
    TcDiagnostic (..),
    TcErrorKind (..),
    TcInterface,
    TcType (..),
    TyCon (..),
    TyVarId (..),
    TypeScheme (..),
    Unique (..),
    emptyTcInterface,
    renderPred,
    renderTcSignature,
    renderTcType,
    tcConfig,
    tcModuleBindings,
    tcModuleDiagnostics,
    tcModuleInstances,
    tcModuleSuccess,
    typecheckModuleWithEnvAndInstances,
    typecheckModulesWithInterfaceConfig,
  )
import Control.Monad.IO.Class (liftIO)
import Data.Aeson ((.:))
import Data.Aeson qualified as Aeson
import Data.ByteString.Lazy qualified as BL
import Data.Char (isAlpha, isSpace)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.List (find, intercalate, isPrefixOf, sortOn, stripPrefix)
import Data.List.NonEmpty qualified as NonEmpty
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.IO.Utf8 qualified as Utf8
import Prettyprinter (defaultLayoutOptions, layoutPretty, pretty)
import Prettyprinter.Render.String (renderString)
import System.Console.Haskeline qualified as Haskeline
import System.Directory (doesDirectoryExist, doesFileExist, getCurrentDirectory, listDirectory)
import System.Environment (lookupEnv)
import System.FilePath ((</>))
import System.FilePath qualified as FilePath

data ReplSession = ReplSession
  { replModuleExports :: !ModuleExports,
    replImportedTerms :: ![(Text, TypeScheme)],
    replBindingTypes :: !(Map.Map Text (Map.Map Text TcType)),
    replImportedInstances :: ![InstanceInfo],
    replTcInterface :: !TcInterface,
    replDependencyProgram :: !FcProgram,
    replSettings :: !(IORef ReplSettings)
  }

data ReplSettings = ReplSettings
  { replShowParsedPretty :: !Bool,
    replShowParsedShorthand :: !Bool,
    replShowInferredType :: !Bool,
    replShowSystemFc :: !Bool
  }
  deriving (Eq, Show)

defaultReplSettings :: ReplSettings
defaultReplSettings =
  ReplSettings
    { replShowParsedPretty = False,
      replShowParsedShorthand = False,
      replShowInferredType = False,
      replShowSystemFc = False
    }

data ReplError
  = ReplMissingBase FilePath
  | ReplInvalidInterface FilePath String
  | ReplParseError
  | ReplResolveError [ResolveError]
  | ReplTypeError [String]
  | ReplDesugarError [String]
  | ReplEvalError String
  deriving (Eq, Show)

data ReplStep
  = ReplContinue !(Maybe String)
  | ReplExit !(Maybe String)
  deriving (Eq, Show)

data CheckedExpression = CheckedExpression
  { checkedExpr :: !Expr,
    checkedResolvedModule :: !Module,
    checkedTcModule :: !Module,
    checkedType :: !TcType
  }

runRepl :: Maybe FilePath -> IO ()
runRepl maybeStoreRoot = do
  session <- loadReplSession maybeStoreRoot
  let settings :: Haskeline.Settings IO
      settings = (Haskeline.defaultSettings :: Haskeline.Settings IO) {Haskeline.complete = replCompletion session}
  Haskeline.runInputT settings (loop session)
  where
    loop session = do
      minput <- Haskeline.getInputLine "aihc> "
      case minput of
        Nothing -> pure ()
        Just input -> do
          step <- liftIO (handleReplInput session input)
          case step of
            ReplContinue output -> do
              mapM_ Haskeline.outputStrLn output
              loop session
            ReplExit output ->
              mapM_ Haskeline.outputStrLn output

loadReplSession :: Maybe FilePath -> IO ReplSession
loadReplSession maybeStoreRoot = do
  baseContext <- loadAihcBaseContext
  installedInterface <- loadExplicitStoreInterface maybeStoreRoot
  settingsRef <- newIORef defaultReplSettings
  let installedExports = maybe Map.empty packageInterfaceExports installedInterface
      installedTerms = maybe [] interfaceImportedTerms installedInterface
      installedBindingTypes = maybe Map.empty interfaceBindingTypes installedInterface
  pure
    ReplSession
      { replModuleExports = ensurePreludeMvpScope (unionModuleExports (replBaseExports baseContext) installedExports),
        replImportedTerms = mergeImportedTerms (replBaseImportedTerms baseContext) (ensurePreludeMvpTerms installedTerms),
        replBindingTypes =
          ensurePreludeMvpBindingTypes
            (unionBindingTypes (replBaseBindingTypes baseContext) installedBindingTypes),
        replImportedInstances = replBaseImportedInstances baseContext,
        replTcInterface = replBaseTcInterface baseContext,
        replDependencyProgram = replBaseProgram baseContext,
        replSettings = settingsRef
      }

handleReplInput :: ReplSession -> String -> IO ReplStep
handleReplInput session input =
  let trimmedInput = trim input
   in case trimmedInput of
        "" -> pure (ReplContinue Nothing)
        ":quit" -> pure (ReplExit Nothing)
        ":q" -> pure (ReplExit Nothing)
        ":help" -> pure (ReplContinue (Just helpText))
        ":?" -> pure (ReplContinue (Just helpText))
        ":set" -> ReplContinue . Just . renderReplSettings <$> readIORef (replSettings session)
        _ | Just command <- stripPrefix ":set " trimmedInput -> handleSetCommand session command
        _ | Just moduleName' <- replCommandArgument ":browse" trimmedInput -> handleBrowseCommand session moduleName'
        _ | Just expression <- replCommandArgument ":type" trimmedInput -> handleTypeCommand session expression
        _ | Just expression <- replCommandArgument ":t" trimmedInput -> handleTypeCommand session expression
        command@(':' : _) -> pure (ReplContinue (Just ("unknown command: " <> command)))
        _ -> do
          result <- evaluateExpression session (T.pack trimmedInput)
          pure (ReplContinue (Just (either renderReplError T.unpack result)))

evaluateExpression :: ReplSession -> Text -> IO (Either ReplError Text)
evaluateExpression session input = do
  settings <- readIORef (replSettings session)
  case typecheckExpression session input of
    Left err -> pure (Left err)
    Right checked -> do
      let expr = checkedExpr checked
          tcResult = checkedTcModule checked
          inferredType = checkedType checked
          allBindings = importedTermBindings (replImportedTerms session) <> tcModuleBindings tcResult
          dsResult = desugarModuleWithInterface (DesugarConfig {primPackageId = PackageId "aihc-prim"}) allBindings (replTcInterface session) tcResult
      if not (dsSuccess dsResult)
        then pure (Left (ReplDesugarError (dsErrors dsResult)))
        else do
          valueResult <- evalProgramBinding replBindingName (concatPrograms [replDependencyProgram session, dsProgram dsResult])
          case mapLeft (ReplEvalError . show) valueResult of
            Left err -> pure (Left err)
            Right value -> do
              renderedResult <- renderValue value
              pure $ do
                renderedValue <- mapLeft (ReplEvalError . show) renderedResult
                Right (renderEvaluation settings expr inferredType dsResult renderedValue)

typecheckExpression :: ReplSession -> Text -> Either ReplError CheckedExpression
typecheckExpression session input = do
  expr <-
    case parseExpr defaultConfig {parserSourceName = "<repl>"} input of
      ParseOk parsed -> Right parsed
      ParseErr _ -> Left ReplParseError
  let parsedModule = replModule expr
      resolved = resolveWithDeps (replModuleExports session) [(replPackage, parsedModule)]
  case resolveErrors resolved of
    [] -> pure ()
    errors -> Left (ReplResolveError errors)
  resolvedModule <-
    case resolvedModules resolved of
      [(_, modu)] -> Right modu
      _ -> Left (ReplResolveError [ResolveNotImplemented "REPL resolver returned no module"])
  let tcResult = typecheckModuleWithEnvAndInstances (replImportedTerms session) (replImportedInstances session) resolvedModule
  if tcModuleSuccess tcResult
    then pure ()
    else Left (ReplTypeError (map renderReplTcDiagnostic (tcModuleDiagnostics tcResult)))
  inferredType <-
    case find ((== replBindingName) . tbName) (tcModuleBindings tcResult) of
      Just binding -> Right (tbType binding)
      Nothing -> Left (ReplTypeError ["missing inferred type for " <> T.unpack replBindingName])
  Right
    CheckedExpression
      { checkedExpr = expr,
        checkedResolvedModule = resolvedModule,
        checkedTcModule = tcResult,
        checkedType = inferredType
      }

handleTypeCommand :: ReplSession -> String -> IO ReplStep
handleTypeCommand session expression =
  pure . ReplContinue . Just $
    case typecheckExpression session (T.pack expression) of
      Left err -> renderReplError err
      Right checked -> renderTcSignature (T.pack expression) (checkedType checked)

handleBrowseCommand :: ReplSession -> String -> IO ReplStep
handleBrowseCommand session rawModuleName
  | null rawModuleName =
      pure (ReplContinue (Just "usage: :browse <module>"))
  | otherwise =
      pure . ReplContinue . Just $
        case matchingScopes of
          [scope] -> renderBrowseScope session moduleName' scope
          _ -> "unknown module: " <> rawModuleName
  where
    moduleName' = T.pack rawModuleName
    matchingScopes =
      [ scope
      | (moduleKey, scope) <- Map.toList (replModuleExports session),
        moduleKeyName moduleKey == moduleName'
      ]

renderBrowseScope :: ReplSession -> Text -> Scope -> String
renderBrowseScope session browsingModule scope =
  intercalate "\n" (map snd (sortOn fst (typeEntries <> termEntries)))
  where
    typeEntries =
      [ (name, renderBrowseType name)
      | name <- Map.keys (scopeTypes scope)
      ]
    termEntries =
      [ (name, renderBrowseTerm name (lookupExportType resolved name))
      | (name, resolved) <- Map.toAscList (scopeTerms scope)
      ]

    lookupExportType resolved exportedName = do
      (definingModule, bindingName) <- resolvedBindingKey browsingModule exportedName resolved
      Map.lookup definingModule (replBindingTypes session) >>= Map.lookup bindingName

renderBrowseType :: Text -> String
renderBrowseType = ("type " <>) . T.unpack . renderBrowseName

renderBrowseTerm :: Text -> Maybe TcType -> String
renderBrowseTerm name maybeType =
  case maybeType of
    Just ty -> renderTcSignature (renderBrowseName name) ty
    Nothing -> T.unpack (renderBrowseName name)

renderBrowseName :: Text -> Text
renderBrowseName name =
  case T.uncons name of
    Just (first, _)
      | isAlpha first || first == '_' -> name
    _
      | name == "[]" || T.isPrefixOf "(" name -> name
      | otherwise -> "(" <> name <> ")"

resolvedBindingKey :: Text -> Text -> ResolvedName -> Maybe (Text, Text)
resolvedBindingKey browsingModule exportedName resolved =
  case resolved of
    ResolvedTopLevel _ name ->
      Just
        ( fromMaybe browsingModule (nameQualifier name),
          normalizeImportedBindingName (nameText name)
        )
    ResolvedBuiltin _ -> Just (browsingModule, exportedName)
    _ -> Nothing

replCompletion :: ReplSession -> Haskeline.CompletionFunc IO
replCompletion session =
  Haskeline.completeWordWithPrev Nothing " \t" $ \reversedBeforeWord word ->
    pure
      [ Haskeline.simpleCompletion candidate
      | trim (reverse reversedBeforeWord) == ":browse",
        candidate <- map (T.unpack . moduleKeyName) (Map.keys (replModuleExports session)),
        word `isPrefixOf` candidate
      ]

replCommandArgument :: String -> String -> Maybe String
replCommandArgument command input =
  case stripPrefix command input of
    Just "" -> Just ""
    Just rest@(first : _)
      | isSpace first -> Just (trim rest)
    _ -> Nothing

handleSetCommand :: ReplSession -> String -> IO ReplStep
handleSetCommand session rawCommand =
  case words rawCommand of
    [command] ->
      case parseSettingCommand command of
        Just updateSettings -> do
          let settingsRef = replSettings session
          settings <- readIORef settingsRef
          let settings' = updateSettings settings
          writeIORef settingsRef settings'
          pure (ReplContinue (Just (renderReplSettings settings')))
        Nothing ->
          pure (ReplContinue (Just ("unknown setting: " <> command)))
    _ ->
      pure (ReplContinue (Just "usage: :set [+|-]parsed-pretty|parsed-shorthand|type|system-fc"))

parseSettingCommand :: String -> Maybe (ReplSettings -> ReplSettings)
parseSettingCommand command =
  case command of
    "+parsed-pretty" -> Just (\settings -> settings {replShowParsedPretty = True})
    "-parsed-pretty" -> Just (\settings -> settings {replShowParsedPretty = False})
    "+pretty" -> Just (\settings -> settings {replShowParsedPretty = True})
    "-pretty" -> Just (\settings -> settings {replShowParsedPretty = False})
    "+parsed-shorthand" -> Just (\settings -> settings {replShowParsedShorthand = True})
    "-parsed-shorthand" -> Just (\settings -> settings {replShowParsedShorthand = False})
    "+shorthand" -> Just (\settings -> settings {replShowParsedShorthand = True})
    "-shorthand" -> Just (\settings -> settings {replShowParsedShorthand = False})
    "+type" -> Just (\settings -> settings {replShowInferredType = True})
    "-type" -> Just (\settings -> settings {replShowInferredType = False})
    "+system-fc" -> Just (\settings -> settings {replShowSystemFc = True})
    "-system-fc" -> Just (\settings -> settings {replShowSystemFc = False})
    "+fc" -> Just (\settings -> settings {replShowSystemFc = True})
    "-fc" -> Just (\settings -> settings {replShowSystemFc = False})
    _ -> Nothing

renderReplSettings :: ReplSettings -> String
renderReplSettings settings =
  unwords
    [ "settings:",
      flag replShowParsedPretty "parsed-pretty",
      flag replShowParsedShorthand "parsed-shorthand",
      flag replShowInferredType "type",
      flag replShowSystemFc "system-fc"
    ]
  where
    flag selector name =
      (if selector settings then "+" else "-") <> name

renderEvaluation :: ReplSettings -> Expr -> TcType -> DesugarResult -> Text -> Text
renderEvaluation settings expr inferredType dsResult renderedValue =
  T.intercalate "\n" (debugSections <> [renderedValue])
  where
    debugSections =
      concat
        [ section "parsed" (T.pack (renderPrettyExpr expr)) (replShowParsedPretty settings),
          section "shorthand" (T.pack (show (shorthand expr))) (replShowParsedShorthand settings),
          section "type" (T.pack (renderTcType inferredType)) (replShowInferredType settings),
          section "system-fc" (T.pack (renderProgram (dsProgram dsResult))) (replShowSystemFc settings)
        ]
    section label body enabled =
      [label <> ":\n" <> body | enabled]

renderPrettyExpr :: Expr -> String
renderPrettyExpr =
  renderString . layoutPretty defaultLayoutOptions . pretty

replModule :: Expr -> Module
replModule expr =
  Module
    { moduleAnns = [],
      moduleHead =
        Just
          ModuleHead
            { moduleHeadAnns = [],
              moduleHeadName = "Aihc.Repl",
              moduleHeadWarningPragma = Nothing,
              moduleHeadExports = Nothing
            },
      moduleLanguagePragmas = [],
      moduleImports = [],
      moduleDecls =
        [ DeclValue $
            FunctionBind
              (mkUnqualifiedName NameVarId replBindingName)
              [ Match
                  { matchAnns = [],
                    matchHeadForm = MatchHeadPrefix,
                    matchPats = [],
                    matchRhs = UnguardedRhs [] expr Nothing
                  }
              ]
        ]
    }

replBindingName :: Text
replBindingName = "__aihc_repl_it"

replPackage :: Package
replPackage = Package "repl" (PackageId "repl")

aihcBasePackage :: Package
aihcBasePackage = Package "aihc-base" (PackageId "aihc-base")

aihcPrimPackage :: Package
aihcPrimPackage = Package "aihc-prim" (PackageId "aihc-prim")

renderReplError :: ReplError -> String
renderReplError err =
  case err of
    ReplMissingBase storeRoot ->
      "repl error: installed base package not found in " <> storeRoot <> "; run `aihc install base` first"
    ReplInvalidInterface path message ->
      "repl error: could not load " <> path <> ": " <> message
    ReplParseError ->
      "parse error"
    ReplResolveError errors ->
      "resolve error: " <> unwords (map show errors)
    ReplTypeError errors ->
      "type error: " <> unwords errors
    ReplDesugarError errors ->
      "desugar error: " <> unwords errors
    ReplEvalError message ->
      "eval error: " <> message

renderReplTcDiagnostic :: TcDiagnostic -> String
renderReplTcDiagnostic diagnostic =
  case diagKind diagnostic of
    UnsolvedWanted pred' _ ->
      "no instance for " <> renderPred pred'
    other -> show other

helpText :: String
helpText =
  unlines
    [ "Commands:",
      "  :quit, :q  Exit the REPL",
      "  :help, :?  Show this help",
      "  :browse     Show a module's exported types and terms",
      "  :type, :t  Show an expression's type"
    ]

trim :: String -> String
trim = dropWhileEnd isSpace . dropWhile isSpace

dropWhileEnd :: (a -> Bool) -> [a] -> [a]
dropWhileEnd predicate = reverse . dropWhile predicate . reverse

data ReplBaseContext = ReplBaseContext
  { replBaseExports :: !ModuleExports,
    replBaseImportedTerms :: ![(Text, TypeScheme)],
    replBaseBindingTypes :: !(Map.Map Text (Map.Map Text TcType)),
    replBaseImportedInstances :: ![InstanceInfo],
    replBaseTcInterface :: !TcInterface,
    replBaseProgram :: !FcProgram
  }

loadAihcBaseContext :: IO ReplBaseContext
loadAihcBaseContext = do
  baseRoot <- defaultAihcBaseRoot
  primRoot <- defaultAihcPrimRoot
  modulesResult <-
    loadTransitiveModules
      [ (aihcBasePackage, baseRoot),
        (aihcPrimPackage, primRoot)
      ]
      (Set.singleton "Prelude")
  case modulesResult of
    Left err -> ioError (userError ("repl error: could not load bundled aihc-base Prelude: " <> err))
    Right modules -> buildBaseContext modules

buildBaseContext :: [(Package, Module)] -> IO ReplBaseContext
buildBaseContext modules =
  case resolveWithDeps Map.empty modules of
    resolved@ResolveResult {resolveErrors = [], resolvedModules} -> do
      let moduleAsts = map snd resolvedModules
          (tcResults, tcInterface) =
            typecheckModulesWithInterfaceConfig
              (tcConfig (PackageId "aihc-prim"))
              emptyTcInterface
              moduleAsts
      if all tcModuleSuccess tcResults
        then do
          let allBindings = concatMap tcModuleBindings tcResults
              dsResults = map (desugarModuleWithInterface (DesugarConfig {primPackageId = PackageId "aihc-prim"}) allBindings tcInterface) tcResults
              bindingTypes = moduleBindingTypes moduleAsts tcResults
          if all dsSuccess dsResults
            then
              pure
                ReplBaseContext
                  { replBaseExports = extractInterface resolved,
                    replBaseImportedTerms = map bindingImportedTerm allBindings,
                    replBaseBindingTypes = bindingTypes,
                    replBaseImportedInstances = concatMap tcModuleInstances tcResults,
                    replBaseTcInterface = tcInterface,
                    replBaseProgram = concatPrograms (map dsProgram dsResults)
                  }
            else ioError (userError ("repl error: could not desugar bundled aihc-base Prelude: " <> unwords (concatMap dsErrors dsResults)))
        else ioError (userError ("repl error: could not type-check bundled aihc-base Prelude: " <> unwords (concatMap (map show . tcModuleDiagnostics) tcResults)))
    ResolveResult {resolveErrors} ->
      ioError (userError ("repl error: could not resolve bundled aihc-base Prelude: " <> show resolveErrors))

bindingImportedTerm :: TcBindingResult -> (Text, TypeScheme)
bindingImportedTerm binding =
  (tbName binding, tcTypeScheme (tbType binding))

mergeImportedTerms :: [(Text, TypeScheme)] -> [(Text, TypeScheme)] -> [(Text, TypeScheme)]
mergeImportedTerms preferred fallback =
  Map.toList (Map.fromList fallback <> Map.fromList preferred)

unionModuleExports :: ModuleExports -> ModuleExports -> ModuleExports
unionModuleExports preferred fallback =
  preferred `Map.union` Map.filterWithKey keepFallback fallback
  where
    preferredNames = Set.fromList (map moduleKeyName (Map.keys preferred))
    keepFallback moduleKey _ = moduleKeyName moduleKey `Set.notMember` preferredNames

unionBindingTypes :: Map.Map Text (Map.Map Text TcType) -> Map.Map Text (Map.Map Text TcType) -> Map.Map Text (Map.Map Text TcType)
unionBindingTypes = Map.unionWith Map.union

moduleBindingTypes :: [Module] -> [Module] -> Map.Map Text (Map.Map Text TcType)
moduleBindingTypes sourceModules checkedModules =
  Map.fromList
    [ (moduleName', bindingTypesForModule checkedModule)
    | (sourceModule, checkedModule) <- zip sourceModules checkedModules,
      Just moduleName' <- [Syntax.moduleName sourceModule]
    ]

bindingTypesForModule :: Module -> Map.Map Text TcType
bindingTypesForModule =
  Map.fromList
    . map (\binding -> (normalizeImportedBindingName (tbName binding), tbType binding))
    . tcModuleBindings

importedTermBindings :: [(Text, TypeScheme)] -> [TcBindingResult]
importedTermBindings terms =
  [ TcBindingResult
      { tbName = name,
        tbDisplayName = name,
        tbType = instantiateSchemeBody scheme
      }
  | (name, scheme) <- terms
  ]

instantiateSchemeBody :: TypeScheme -> TcType
instantiateSchemeBody (ForAll tvs preds body) =
  foldr TcForAllTy (qualify preds body) tvs
  where
    qualify [] ty = ty
    qualify constraints ty = TcQualTy constraints ty

concatPrograms :: [FcProgram] -> FcProgram
concatPrograms programs =
  case NonEmpty.nonEmpty programs of
    Nothing -> FcProgram (FcModuleId "repl" "Interactive") []
    Just nonEmptyPrograms ->
      either
        (error . ("System FC merge error: " <>) . show)
        id
        (mergePrograms (FcModuleId "repl" "Interactive") nonEmptyPrograms)

loadExplicitStoreInterface :: Maybe FilePath -> IO (Maybe PackageInterface)
loadExplicitStoreInterface Nothing = pure Nothing
loadExplicitStoreInterface (Just storeRoot) = do
  interfacePath <- findInstalledBaseInterface storeRoot
  Just <$> loadInterface interfacePath

loadInterface :: FilePath -> IO PackageInterface
loadInterface path = do
  result <- readPackageInterface path
  case result of
    Left err -> ioError (userError (renderReplError (ReplInvalidInterface path err)))
    Right interface -> pure interface

interfaceImportedTerms :: PackageInterface -> [(Text, TypeScheme)]
interfaceImportedTerms interface =
  [ (normalizeImportedBindingName (packageInterfaceBindingName binding), tcTypeScheme (packageInterfaceBindingType binding))
  | modu <- packageInterfaceTypecheck interface,
    binding <- packageInterfaceTcModuleBindings modu
  ]

interfaceBindingTypes :: PackageInterface -> Map.Map Text (Map.Map Text TcType)
interfaceBindingTypes interface =
  Map.fromList
    [ ( packageInterfaceTcModuleName modu,
        Map.fromList
          [ (normalizeImportedBindingName (packageInterfaceBindingName binding), packageInterfaceBindingType binding)
          | binding <- packageInterfaceTcModuleBindings modu
          ]
      )
    | modu <- packageInterfaceTypecheck interface
    ]

defaultAihcBaseRoot :: IO FilePath
defaultAihcBaseRoot = defaultCoreLibraryRoot "AIHC_BASE_SRC" "aihc-base"

defaultAihcPrimRoot :: IO FilePath
defaultAihcPrimRoot = defaultCoreLibraryRoot "AIHC_PRIM_SRC" "aihc-prim"

defaultCoreLibraryRoot :: String -> FilePath -> IO FilePath
defaultCoreLibraryRoot sourceEnv packageName = do
  sourceOverride <- lookupEnv sourceEnv
  case sourceOverride of
    Just root -> pure root
    Nothing -> do
      coreLibsOverride <- lookupEnv "AIHC_CORE_LIBS_ROOT"
      case coreLibsOverride of
        Just root -> pure (root </> "core-libs" </> packageName)
        Nothing -> do
          cwd <- getCurrentDirectory
          findUp cwd
  where
    findUp dir = do
      let candidate = dir </> "core-libs" </> packageName
      exists <- doesDirectoryExist candidate
      if exists
        then pure candidate
        else do
          let parent = FilePath.takeDirectory dir
          if parent == dir
            then pure candidate
            else findUp parent

loadTransitiveModules :: [(Package, FilePath)] -> Set.Set Text -> IO (Either String [(Package, Module)])
loadTransitiveModules packageRoots initialModules =
  fmap snd <$> go Set.empty [] (Set.toAscList initialModules)
  where
    go seen loaded [] =
      pure (Right (seen, loaded))
    go seen loaded (moduleName : pending)
      | moduleName `Set.member` seen =
          go seen loaded pending
      | otherwise = do
          maybePath <- findModulePathInDependencies packageRoots moduleName
          case maybePath of
            Nothing -> do
              let dependencyNames = T.intercalate ", " (map (packageName . fst) packageRoots)
              pure (Left ("dependency module " <> T.unpack moduleName <> " not found in dependencies: " <> T.unpack dependencyNames))
            Just (package, path) -> do
              source <- Utf8.readFile path
              case parseOneModule path source of
                Left errMsg -> pure (Left ("dependency module " <> T.unpack moduleName <> " parse error: " <> errMsg))
                Right modu -> do
                  let seen' = Set.insert moduleName seen
                      newImports = Set.toAscList (importedModuleNames [modu] `Set.difference` seen')
                  depResult <- go seen' loaded newImports
                  case depResult of
                    Left errMsg -> pure (Left errMsg)
                    Right (seenWithDeps, loadedWithDeps) ->
                      go seenWithDeps (loadedWithDeps <> [(package, modu)]) pending

findModulePathInDependencies :: [(Package, FilePath)] -> Text -> IO (Maybe (Package, FilePath))
findModulePathInDependencies [] _ = pure Nothing
findModulePathInDependencies ((package, root) : rest) moduleName = do
  let path = root </> "src" </> moduleNamePath moduleName
  exists <- doesFileExist path
  if exists
    then pure (Just (package, path))
    else findModulePathInDependencies rest moduleName

moduleNamePath :: Text -> FilePath
moduleNamePath moduleName =
  FilePath.joinPath (map T.unpack (T.splitOn "." moduleName)) <> ".hs"

parseOneModule :: FilePath -> Text -> Either String Module
parseOneModule sourceName input =
  let cfg =
        defaultConfig
          { parserSourceName = sourceName
          }
      (errs, ast) = parseModule cfg input
   in if null errs
        then Right ast
        else Left ("parse module error: " <> show errs)

importedModuleNames :: [Module] -> Set.Set Text
importedModuleNames modules =
  Set.fromList [importDeclModule importDecl | modu <- modules, importDecl <- moduleImports modu]

findInstalledBaseInterface :: FilePath -> IO FilePath
findInstalledBaseInterface storeRoot = do
  storeExists <- doesDirectoryExist storeRoot
  if not storeExists
    then ioError (userError (renderReplError (ReplMissingBase storeRoot)))
    else do
      entries <- listDirectory storeRoot
      candidates <- mapM (readCandidate . (storeRoot </>)) entries
      case find isBaseManifest candidates of
        Just candidate -> pure (candidateInterfacePath candidate)
        Nothing -> ioError (userError (renderReplError (ReplMissingBase storeRoot)))
  where
    readCandidate path = do
      let manifestPath = path </> "manifest.json"
      exists <- doesFileExist manifestPath
      if exists
        then do
          bytes <- BL.readFile manifestPath
          pure $
            case Aeson.eitherDecode bytes of
              Right candidate -> candidate
              Left _ -> StoreCandidate "" "" (path </> "interfaces" </> "package-interface.json")
        else pure (StoreCandidate "" "" (path </> "interfaces" </> "package-interface.json"))

data StoreCandidate = StoreCandidate
  { candidatePackageName :: !String,
    candidatePackageVersion :: !String,
    candidateInterfacePath :: !FilePath
  }
  deriving (Show)

instance Aeson.FromJSON StoreCandidate where
  parseJSON =
    Aeson.withObject "manifest" $ \obj -> do
      pkg <- obj .: "package"
      StoreCandidate
        <$> pkg .: "name"
        <*> pkg .: "version"
        <*> obj .: "interfacePath"

isBaseManifest :: StoreCandidate -> Bool
isBaseManifest candidate =
  candidatePackageName candidate `elem` ["aihc-base", "base"]

ensurePreludeMvpScope :: ModuleExports -> ModuleExports
ensurePreludeMvpScope exports =
  Map.insert preludeKey preludeScope exports
  where
    preludeKey =
      fromMaybe
        (ModuleKey aihcBasePackage "Prelude")
        (find ((== "Prelude") . moduleKeyName) (Map.keys exports))
    existing = Map.findWithDefault emptyScope preludeKey exports
    resolvePreludeName =
      ResolvedTopLevel (packageId (moduleKeyPackage preludeKey))
        . qualifyName (Just "Prelude")
        . unqualifiedNameFromText
    preludeScope =
      existing
        { scopeTerms =
            scopeTerms existing
              <> Map.fromList
                [ ("++", resolvePreludeName "++")
                ],
          scopeTypes =
            scopeTypes existing
              <> Map.fromList
                [ ("Char", resolvePreludeName "Char"),
                  ("String", resolvePreludeName "String")
                ]
        }

ensurePreludeMvpTerms :: [(Text, TypeScheme)] -> [(Text, TypeScheme)]
ensurePreludeMvpTerms terms =
  Map.toList (Map.fromList [("++", appendScheme)] <> Map.fromList terms)

ensurePreludeMvpBindingTypes :: Map.Map Text (Map.Map Text TcType) -> Map.Map Text (Map.Map Text TcType)
ensurePreludeMvpBindingTypes =
  Map.alter (Just . addAppend . fromMaybe Map.empty) "Prelude"
  where
    addAppend = Map.insertWith (\_ old -> old) "++" (instantiateSchemeBody appendScheme)

appendScheme :: TypeScheme
appendScheme =
  ForAll [aVar] [] (TcFunTy listA (TcFunTy listA listA))
  where
    aVar = TyVarId "a" (Unique (-100))
    listA = TcTyCon (TyCon "[]" 1) [TcTyVar aVar]

normalizeImportedBindingName :: Text -> Text
normalizeImportedBindingName name =
  case T.uncons name of
    Just ('(', rest)
      | Just inner <- T.stripSuffix ")" rest ->
          inner
    _ -> name

tcTypeScheme :: TcType -> TypeScheme
tcTypeScheme ty =
  case collectForAlls ty of
    (tvs, TcQualTy preds body) -> ForAll tvs preds body
    (tvs, body) -> ForAll tvs [] body

collectForAlls :: TcType -> ([TyVarId], TcType)
collectForAlls (TcForAllTy tv body) =
  let (tvs, inner) = collectForAlls body
   in (tv : tvs, inner)
collectForAlls ty = ([], ty)

emptyScope :: Scope
emptyScope = Scope Map.empty Map.empty Map.empty Map.empty Map.empty Map.empty Map.empty

mapLeft :: (a -> b) -> Either a c -> Either b c
mapLeft f value =
  case value of
    Left err -> Left (f err)
    Right ok -> Right ok
