module Main

import DepGraph.Data
import DepGraph.Cli
import Idris.Package.Extra
import Idris.Driver.Extra

import Compiler.Common
import Core.Context
import Core.Directory
import Core.InitPrimitives
import Data.String
import Idris.Error
import Idris.Pretty
import Idris.ModTree
import Idris.SetOptions
import Idris.Syntax
import Idris.REPL
import System
import System.File
import System.Term
import Libraries.Data.StringMap
import Libraries.System.Directory.Tree
import Libraries.Utils.Path

%hide Idris.Syntax.Module -- DepGraph.Data.Module

--------------------------------------------------------------------------------
-- Common utils
--------------------------------------------------------------------------------

loadPkg :
  {auto c : Ref Ctxt Defs} ->
  {auto s : Ref Syn SyntaxInfo} ->
  {auto o : Ref ROpts REPLOpts} ->
  String ->
  Core (PkgDesc, String)
loadPkg file = do
  let Just (root, filename) = splitParent file
      | _ => throw $ InternalError "Tried to split empty string"
  let True = isSuffixOf ".ipkg" filename
      | _ => do
        coreLift $ putStrLn $
          "Packages must have an '.ipkg' extension: " ++ show file ++ "."
        coreLift $ exitWith $ ExitFailure 1

  setWorkingDir root
  pkg <- parsePkgFile True filename
  pure (pkg, root)

||| Ensure file path is absolute (uses `getWorkingDir`)
toAbsolute : String -> Core String
toAbsolute path = do
  if isAbsolute path
    then pure path
    else do
      wdir <- getWorkingDir
      pure $ wdir </> path

--------------------------------------------------------------------------------
-- Modules dependency graph
--------------------------------------------------------------------------------

processMod :
  (modFilter : List ModuleIdent -> List ModuleIdent) ->
  BuildMod ->
  Module
processMod modFilter mod =
  MkModule
    { name = show mod.buildNS
    , deps = show <$> modFilter mod.imports
    }

record RawPkgModules where
  constructor MkRawPkgModules
  name : String
  root : String
  ||| Absolute paths to modules
  mods : List String
  buildMods : List BuildMod

processPkgMods :
  ModOpts ->
  (localMods : List ModuleIdent) ->
  RawPkgModules ->
  PkgModules
processPkgMods opts localMods rawPkg = do
  let MkRawPkgModules { name, root, mods, buildMods } = rawPkg
  let modules = (processMod modFilter) <$> buildMods
  MkPkgModules { name, modules }
  where
    modFilter : List ModuleIdent -> List ModuleIdent
    modFilter =
      if opts.withExternal
       then id
       else filter $ \mod => any (== mod) localMods

getLocalMods :
  ModOpts ->
  (localFiles : List String) ->
  RawPkgModules ->
  List ModuleIdent
getLocalMods opts localFiles rawPkg =
  buildNS <$> (modFilter rawPkg.root) rawPkg.buildMods
  where
    modFilter : String -> List BuildMod -> List BuildMod
    modFilter dir =
      if opts.withExternal
       then id
       else filter $ \mod =>
         any (\local => (dir </> buildFile mod) == local) localFiles

findPkgModules :
  {auto c : Ref Ctxt Defs} ->
  {auto s : Ref Syn SyntaxInfo} ->
  {auto o : Ref ROpts REPLOpts} ->
  PkgDesc ->
  Core (List String)
findPkgModules pkg = do
  withWarnings $ addDeps pkg
  runScript (prebuild pkg)
  pure $ maybe (map snd (modules pkg))
               (\m => snd m :: map snd (modules pkg))
               (mainmod pkg)

getAllBuildMods :
  {auto c : Ref Ctxt Defs} ->
  {auto o : Ref ROpts REPLOpts} ->
  FC -> (done : List BuildMod) ->
  (allFiles : List String) ->
  Core (List BuildMod)
getAllBuildMods fc done [] = pure done
getAllBuildMods fc done (f :: fs) = do
  ms <- getBuildMods fc done f
  getAllBuildMods fc (ms ++ done) fs

loadPkgModules :
  {auto c : Ref Ctxt Defs} ->
  {auto s : Ref Syn SyntaxInfo} ->
  {auto o : Ref ROpts REPLOpts} ->
  String ->
  Core RawPkgModules
loadPkgModules file = do
  (pkg, root) <- loadPkg file
  whenJust (builddir pkg) setBuildDir
  setOutputDir (outputdir pkg)
  mods <- traverse toAbsolute !(findPkgModules pkg)
  -- There might be duplicates, so if something is already processed, drop it
  buildMods <- dropLater <$> getAllBuildMods EmptyFC [] mods

  pure $ MkRawPkgModules { name = pkg.name, root, mods, buildMods }
  where
    dropLater : List BuildMod -> List BuildMod
    dropLater [] = []
    dropLater (b :: bs)
        = b :: dropLater (filter (\x => buildFile x /= buildFile b) bs)

modMain :
  {auto c : Ref Ctxt Defs} ->
  {auto s : Ref Syn SyntaxInfo} ->
  {auto o : Ref ROpts REPLOpts} ->
  ModOpts ->
  List String ->
  Core ()
modMain opts files = do
  pkgs <- for files loadPkgModules
  let localFiles = concat $ mods <$> pkgs
  let localMods = concat $ getLocalMods opts localFiles <$> pkgs
  let modules = (processPkgMods opts localMods) <$> pkgs
  coreLift $ putStrLn $ show modules

--------------------------------------------------------------------------------
-- Package dependency graph
--------------------------------------------------------------------------------

processPackage :
  (isLocal : Bool) ->
  (localOnly : Bool) ->
  (localPkgs : List String) ->
  PkgDesc ->
  Core Package
processPackage isLocal localOnly localPkgs pkg = do
  pure $ MkPackage pkg.name deps isLocal
  where
    depName : Depends -> String
    depName dep = dep.pkgname

    pkgFilter : List String -> List String
    pkgFilter = if localOnly
      then filter $ \pkg => any (== pkg) localPkgs
      else id

    deps : List String
    deps = pkgFilter $ depName <$> pkg.depends

record Candidate where
  constructor MkCandidate
  name      : String
  version   : Maybe PkgVersion
  directory : String

toCandidate : (name : String) -> (String,Maybe PkgVersion) -> Candidate
toCandidate name (dir,v) = MkCandidate name v dir

record ResolutionError where
  constructor MkRE
  decisions : List Candidate
  depends   : Depends
  version   : Maybe PkgVersion

prepend : Candidate -> ResolutionError -> ResolutionError
prepend p = { decisions $= (p ::)}

reason : Maybe PkgVersion -> String
reason Nothing  = "no matching version is installed."
reason (Just x) = "only found version \{show x} which is out of bounds."

printResolutionError : ResolutionError -> String
printResolutionError (MkRE ds d v) = go [<] ds
  where go : SnocList String -> List Candidate -> String
        go ss []        =
          let pre        := "Required \{d.pkgname} \{show d.pkgbounds} but"
              failure    := "\{pre} \{reason v}"
              candidates := case ss of
                               [<] => ""
                               ss  => " Resolved transitive dependencies: " ++ (fastConcat $ intersperse "; " $ cast ss) ++ "."
           in failure ++ candidates
        go ss (c :: cs) =
          let v := fromMaybe defaultVersion c.version
           in go (ss :< "\{c.name}-\{show v}") cs

data ResolutionRes : Type where
  Resolved : List String -> ResolutionRes
  Failed   : List ResolutionError -> ResolutionRes

printErrs : (pkgDirs : List String) -> PkgDesc -> List ResolutionError -> String
printErrs pkgDirs x es =
  let errors := unlines $ "Failed to resolve the dependencies for \{x.name}:"
                          :: map (indent 2 . printResolutionError) es
      dirs   := unlines $ "Searched for packages in:"
                          :: map (indent 2) pkgDirs
  in """
     \{errors}
     \{dirs}
     For more details on what packages Idris2 can locate, run `idris2 --list-packages`
     """

-- try all possible resolution paths, keeping the first
-- that works
tryAll :  List Candidate
       -> (Candidate -> Core ResolutionRes)
       -> Core ResolutionRes
tryAll ps f = go [<] ps
  where go :  SnocList ResolutionError
           -> List Candidate
           -> Core ResolutionRes
        go se []        = pure (Failed $ se <>> [])
        go se (x :: xs) = do
          Failed errs <- f x | Resolved res => pure (Resolved res)
          go (se <>< map (prepend x) errs) xs

pkgDirs :
    {auto c : Ref Ctxt Defs} ->
    Core (List String)
pkgDirs = do
  localdir <- pkgLocalDirectory
  d <- getDirs
  pure (localdir ::  (show <$> d.package_search_paths))

||| Find all dependencies (transitively) from the given package file.
findDeps :
    {auto c : Ref Ctxt Defs} ->
    {auto s : Ref Syn SyntaxInfo} ->
    {auto o : Ref ROpts REPLOpts} ->
    PkgDesc ->
    Core (List String)
findDeps pkg = do
  Resolved allPkgs <- getTransitiveDeps pkg.depends empty
    | Failed errs => throw $ GenericMsg EmptyFC (printErrs !pkgDirs pkg errs)
  log "package.depends" 10 $ "all depends: \{show allPkgs}"
  pure allPkgs
  where
    -- Note: findPkgDir throws an error if a package is not found
    -- *unless* --ignore-missing-ipkg is enabled
    -- therefore, if findPkgDir returns Nothing, skip the package
    --
    -- We use a backtracking algorithm here: If several versions of
    -- a package are installed, we must try all, which are are
    -- potentially in bounds, because their set of dependencies
    -- might be different across versions and not all of them
    -- might lead to a resolvable situation.
    getTransitiveDeps :
        List Depends ->
        (done : StringMap (Maybe PkgVersion, String)) ->
        Core ResolutionRes
    getTransitiveDeps [] done = do
      ms <- for (StringMap.toList done) $
        \(pkg, (mv, pkgFile)) => do pure $ (</> pkgFile) <$> !(findPkgDir pkg (exactBounds mv))
      pure . Resolved $ catMaybes ms

    getTransitiveDeps (dep :: deps) done =
      case lookup dep.pkgname done of
        Just (mv, pkgFile) =>
          if inBounds mv dep.pkgbounds
            -- already resolved dependency is in bounds
            -- so we keep it and resolve remaining deps
            then getTransitiveDeps deps done
            -- the resolved dependency does not satisfy the
            -- current bounds. we return an error and backtrack
            else pure (Failed [MkRE [] dep $ mv <|> Just defaultVersion])
        Nothing => do
          log "package.depends" 50 "adding new dependency: \{dep.pkgname} (\{show dep.pkgbounds})"
          pkgDirs <- findPkgDirs dep.pkgname dep.pkgbounds
          let candidates := toCandidate dep.pkgname <$> pkgDirs

          case candidates of
            [] => do
              defs <- get Ctxt
              if defs.options.session.ignoreMissingPkg
                -- this corresponds to what `findPkgDir` does in
                -- case of `ignoreMissingPkg` being set to `True`
                then getTransitiveDeps deps done
                else pure (Failed [MkRE [] dep Nothing])

            _  => tryAll candidates $ \(MkCandidate name mv pkgDir) => do
              let pkgFile = pkgDir </> name <.> "ipkg"
              True <- coreLift $ exists pkgFile
                | False => getTransitiveDeps deps (insert name (mv, pkgFile) done)
              pkg <- parsePkgFile False pkgFile
              getTransitiveDeps
                (pkg.depends ++ deps)
                (insert pkg.name (pkg.version, pkgFile) done)

pkgMain :
  {auto c : Ref Ctxt Defs} ->
  {auto s : Ref Syn SyntaxInfo} ->
  {auto o : Ref ROpts REPLOpts} ->
  PkgOpts ->
  List String ->
  Core ()
pkgMain opts files = do
  let MkPkgOpts {localOnly} = opts
  pkgsWithRoots <- for files loadPkg
  let pkgs = fst <$> pkgsWithRoots
  let localPkgs = (.name) <$> pkgs
  deps <- if localOnly
    then pure []
    else do
      -- Remove duplicates
      deps <- dropLater <$> concat <$> for pkgs findDeps
      pkgsWithRoots <- for deps loadPkg
      let deps = fst <$> pkgsWithRoots
      -- There might also be local pkgs in deps, remove those too
      pure $ filter (\d => all (\local => local /= d.name) localPkgs) deps
  local <- for pkgs $ processPackage True localOnly localPkgs
  deps <- for deps $ processPackage False localOnly localPkgs
  coreLift $ putStrLn $ show (local ++ deps)
  where
    dropLater : List String -> List String
    dropLater [] = []
    dropLater (b :: bs)
        = b :: dropLater (filter (/= b) bs)

--------------------------------------------------------------------------------
-- Common
--------------------------------------------------------------------------------

setupPackPkgPath :
  {auto c : Ref Ctxt Defs} ->
  Core Bool
setupPackPkgPath = do
  (res, code) <- coreLift $ run "pack package-path"
  let success = code >= 0
  when success $
    traverseList1_ addPackageSearchPath $ splitPaths res
  pure success

findIpkgs : {root : Path} -> Tree root -> IO (List String)
findIpkgs t = depthFirst check t (pure [])
  where
    check : {root : Path} -> FileName root ->
            Lazy (IO (List String)) -> IO (List String)
    check fn next = do
      if ".ipkg" `isSuffixOf` fileName fn
        then do pure $ (toFilePath fn) :: !next
        else next

||| If the user did not provide any package file we can look in the working
||| directory.
localPackageFiles : List String -> Core (List String)
localPackageFiles pkgs@(_ :: _) = for pkgs toAbsolute
localPackageFiles _ = do
  wdir <- getWorkingDir
  tree <-  coreLift $ filter filePred dirPred <$> (explore $ parse wdir)
  candidates <- coreLift $ findIpkgs tree
  let pkgs = filter (".ipkg" `isSuffixOf`) candidates
  for pkgs toAbsolute
    where
      filePred : {root : _} -> FileName root -> Bool
      filePred name = True

      dirPred : {root : _} -> FileName root -> Bool
      dirPred name = all (/= fileName name) ignore
        where
          ignore : List String
          ignore = [".git", ".pijul"]

handleSubCmd :
  {auto c : Ref Ctxt Defs} ->
  {auto s : Ref Syn SyntaxInfo} ->
  {auto o : Ref ROpts REPLOpts} ->
  List String ->
  List CLOpt ->
  Core ()
handleSubCmd pathargs opts = do
  files <- localPackageFiles pathargs
  case subCmd opts of
    SubCmdMod opts => modMain opts files
    SubCmdPkg opts => pkgMain opts files

quitWithError :
  {auto c : Ref Ctxt Defs} ->
  {auto s : Ref Syn SyntaxInfo} ->
  {auto o : Ref ROpts REPLOpts} ->
  Error ->
  Core a
quitWithError err = do
  doc <- display err
  msg <- render doc
  coreLift (die msg)

auxMain : List CLOpt -> Core ()
auxMain opts = do
  defs <- initDefs
  c <- newRef Ctxt defs
  s <- newRef Syn initSyntax
  setCG {c} Chez
  addPrimitives
  setWorkingDir "."
  let outmode = REPL InfoLvl
  o <- newRef ROpts (REPL.Opts.defaultOpts Nothing outmode [])
  True <- setupPackPkgPath
    | False => updateEnv
  let ipkgs = findInputs opts
  flip catch quitWithError $ handleSubCmd ipkgs opts

main : IO ()
main = do
  Right opts <- getCmdOpts
    | Left err => do ignore $ fPutStrLn stderr $ "Error: " ++ err
                     exitWith (ExitFailure 1)
  continue <- quitOpts opts
  when continue $ do
    setupTerm
    coreRun (auxMain opts)
      (\err : Error => do ignore $ fPutStrLn stderr $ "Uncaught error: " ++ show err
                          exitWith (ExitFailure 1))
      (\res => pure ())