module Main where

import Check (checkFile, checkFileWithStore, instrumentIOContinuations)
import ContentStore
import ContentStore.Bundle
import Module.Manifest
import System.Exit (die)
import Eval (evalTricu, mainResult, result)
import FileEval
  ( ContractMode(..)
  , LoadedSource(..)
  , defaultStorePath
  , evaluateFileWithContextWithStoreAndMode
  , evaluateFileWithStore
  , loadFileWithStoreMode
  , compileFileWithStore
  )
import IODriver (IOPermissions(..), runIO)
import Parser (parseTricu)
import REPL (repl, replWithStore)
import Research (T, EvaluatedForm(..), Env, formatT, exportDag)
import Wire (encodeBundle, defaultExportNames, Bundle(..))

import Control.Monad (foldM, forM, unless, when)
import Data.Char (isAlphaNum)
import Data.List (sortOn)
import qualified Data.Text as T
import Data.Version (showVersion)
import Paths_tricu (version)
import Options.Applicative

import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.Sequence as Seq
import qualified Data.Map as Map
import System.Directory (createDirectoryIfMissing, getHomeDirectory)
import System.FilePath (takeBaseName, (</>))

-- ---------------------------------------------------------------------------
-- CLI argument types
-- ---------------------------------------------------------------------------

data AppArgs = AppArgs
  { globalStore :: Maybe FilePath
  , appCommand   :: TricuArgs
  } deriving (Show)

data TricuArgs
  = Repl
  | Check
      { checkInput :: FilePath
      , checkStore :: Maybe FilePath
      }
  | Eval
      { evalFiles         :: [FilePath]
      , evalStore         :: Maybe FilePath
      , evalFormat        :: EvaluatedForm
      , evalOutput        :: FilePath
      , evalUnchecked     :: Bool
      , evalIo            :: Bool
      , evalAllowRead     :: [FilePath]
      , evalAllowWrite    :: [FilePath]
      , evalAllowReadAll  :: Bool
      , evalAllowWriteAll :: Bool
      , evalUnsafeIo      :: Bool
      }
  | ArboricxCompile
      { compileInput   :: FilePath
      , compileStore   :: Maybe FilePath
      , compileOutput  :: FilePath
      , compileNames   :: [String]
      }
  | ArboricxImport
      { importFile   :: FilePath
      , importStore  :: Maybe FilePath
      , importModule :: Maybe String
      }
  | ArboricxExport
      { exportTargets :: [String]
      , exportModules :: [String]
      , exportOutput  :: FilePath
      , exportNames   :: [String]
      , exportStore   :: Maybe FilePath
      , exportAll     :: Bool
      , exportSplit   :: Bool
      , dag           :: Bool
      }
  | StoreAliasList
      { storeAliasKind :: AliasKind
      , storePathOpt   :: Maybe FilePath
      }
  | StoreAliasGet
      { storeAliasKind :: AliasKind
      , storeAliasName :: String
      , storePathOpt   :: Maybe FilePath
      }
  deriving (Show)

-- ---------------------------------------------------------------------------
-- optparse-applicative parsers
-- ---------------------------------------------------------------------------

readEvaluatedForm :: ReadM EvaluatedForm
readEvaluatedForm = eitherReader $ \s -> case s of
  "tree"    -> Right Tree
  "fsl"     -> Right FSL
  "ast"     -> Right AST
  "ternary" -> Right Ternary
  "ascii"   -> Right Ascii
  "decode"  -> Right Decode
  "number"  -> Right Number
  "string"  -> Right StringLit
  _         -> Left $ "Unknown format: " ++ s ++ ". Expected: tree, fsl, ast, ternary, ascii, decode, number, string"

checkParser :: Parser TricuArgs
checkParser = Check
  <$> argument str (metavar "FILE")
  <*> optional (option str
      ( long "store"
     <> short 's'
     <> metavar "PATH"
     <> help "Content-addressed store path for module import resolution"
      ))

evalParser :: Parser TricuArgs
evalParser = Eval
  <$> many (argument str (metavar "FILE..."))
  <*> optional (option str
      ( long "store"
     <> short 's'
     <> metavar "PATH"
     <> help "Content-addressed store path for module import resolution"
      ))
  <*> option readEvaluatedForm
      ( long "format"
     <> short 'f'
     <> metavar "FORM"
     <> value Tree
     <> help "Output format: tree, fsl, ast, ternary, ascii, decode, number, string"
      )
  <*> option str
      ( long "output"
     <> short 'o'
     <> metavar "FILE"
     <> value ""
     <> help "Write output to file instead of stdout"
      )
  <*> switch
      ( long "unchecked"
     <> help "Evaluate as untyped code: ignore View Contract annotations and do not publish unchecked view refs"
      )
  <*> switch
      ( long "io"
     <> help "Interpret the result as an IO action tree and execute it"
      )
  <*> many (option str
      ( long "allow-read"
     <> metavar "PATH"
     <> help "Allow reading from PATH prefix (repeatable)"
      ))
  <*> many (option str
      ( long "allow-write"
     <> metavar "PATH"
     <> help "Allow writing to PATH prefix (repeatable)"
      ))
  <*> switch
      ( long "allow-read-all"
     <> help "Allow reading from any path"
      )
  <*> switch
      ( long "allow-write-all"
     <> help "Allow writing to any path"
      )
  <*> switch
      ( long "unsafe-io"
     <> help "Allow unrestricted read and write access"
      )

compileParser :: Parser TricuArgs
compileParser = ArboricxCompile
  <$> option str
      ( long "file"
     <> short 'f'
     <> metavar "FILE"
     <> value ""
     <> help "Input .tri source file"
      )
  <*> optional (option str
      ( long "store"
     <> short 's'
     <> metavar "PATH"
     <> help "Content-addressed store path for module import resolution"
      ))
  <*> option str
      ( long "output"
     <> short 'o'
     <> metavar "FILE"
     <> value ""
     <> help "Output bundle file path (required)"
      )
  <*> many (option str
      ( long "name"
     <> short 'n'
     <> metavar "NAME"
     <> help "Definition name(s) to export as bundle roots (repeatable)"
      ))

importParser :: Parser TricuArgs
importParser = ArboricxImport
  <$> option str
      ( long "file"
     <> short 'f'
     <> metavar "FILE"
     <> value ""
     <> help "Bundle file to import"
      )
  <*> optional (option str
      ( long "store"
     <> short 's'
     <> metavar "PATH"
     <> help "Content-addressed store path"
      ))
  <*> optional (option str
      ( long "module"
     <> short 'm'
     <> metavar "NAME"
     <> help "Module alias to create for the imported bundle (defaults to bundle file basename)"
      ))

exportParser :: Parser TricuArgs
exportParser = ArboricxExport
  <$> many (option str
      ( long "target"
     <> short 't'
     <> metavar "TARGET"
     <> help "Target hash or name (repeatable)"
      ))
  <*> many (option str
      ( long "module"
     <> short 'm'
     <> metavar "MODULE"
     <> help "Module alias or manifest hash to export (repeatable; bundle export only)"
      ))
  <*> option str
      ( long "output"
     <> short 'o'
     <> metavar "FILE"
     <> value ""
     <> help "Output file path (required for bundle export)"
      )
  <*> many (option str
      ( long "name"
     <> short 'n'
     <> metavar "NAME"
     <> help "Export name(s) for the bundle manifest (repeatable)"
      ))
  <*> optional (option str
      ( long "store"
     <> short 's'
     <> metavar "PATH"
     <> help "Content-addressed store path"
      ))
  <*> switch
      ( long "all"
     <> help "Export all name aliases that point at tree-term objects"
      )
  <*> switch
      ( long "split"
     <> help "Write one single-export bundle per export; --output is treated as a directory"
      )
  <*> switch
      ( long "dag"
     <> help "Export as a topologically-sorted DAG node table instead of a bundle"
      )

aliasKindReader :: ReadM AliasKind
aliasKindReader = eitherReader $ \s -> case s of
  "names"    -> Right NameAlias
  "name"     -> Right NameAlias
  "modules"  -> Right ModuleAlias
  "module"   -> Right ModuleAlias
  "packages" -> Right PackageAlias
  "package"  -> Right PackageAlias
  _          -> Left "alias kind must be one of: names, modules, packages"

storePathParser :: Parser (Maybe FilePath)
storePathParser = optional (option str
  ( long "store"
 <> short 's'
 <> metavar "PATH"
 <> help "Content-addressed store path"
  ))

aliasKindParser :: Parser AliasKind
aliasKindParser = option aliasKindReader
  ( long "kind"
 <> short 'k'
 <> metavar "KIND"
 <> value NameAlias
 <> help "Alias kind: names, modules, packages (default: names)"
  )

storeAliasListParser :: Parser TricuArgs
storeAliasListParser = StoreAliasList
  <$> aliasKindParser
  <*> storePathParser

storeAliasGetParser :: Parser TricuArgs
storeAliasGetParser = StoreAliasGet
  <$> aliasKindParser
  <*> argument str (metavar "NAME")
  <*> storePathParser

versionStr :: String
versionStr = "tricu " ++ showVersion version

tricuParser :: Parser AppArgs
tricuParser = AppArgs
  <$> optional (option str
      ( long "store"
     <> metavar "PATH"
     <> help "Global content-addressed store path used by commands and the REPL unless a subcommand overrides it"
      ))
  <*> ((subparser topCommands <|> pure Repl)
      <**> infoOption versionStr (long "version" <> help "Show version"))
  where
    topCommands = mconcat
      [ command "check" (info (checkParser <**> helper)
          (progDesc "Check View Contract annotations and report ok or diagnostics"))
      , command "eval" (info (evalParser <**> helper)
          (progDesc "Evaluate tricu source and print the result of the final expression"))
      , command "arboricx" (info (arboricxParser <**> helper)
          (progDesc "Arboricx bundle operations"))
      , command "store" (info (storeParser <**> helper)
          (progDesc "Inspect and manage the content-addressed store"))
      ]

arboricxParser :: Parser TricuArgs
arboricxParser = subparser $ mconcat
  [ command "compile" (info (compileParser <**> helper)
      (progDesc "Compile a .tri file into a standalone Arboricx bundle"))
  , command "import" (info (importParser <**> helper)
      (progDesc "Import an Arboricx bundle into the content store"))
  , command "export" (info (exportParser <**> helper)
      (progDesc "Export one or more terms from the content store"))
  ]

storeParser :: Parser TricuArgs
storeParser = subparser $ mconcat
  [ command "alias" (info (storeAliasParser <**> helper)
      (progDesc "Inspect workspace aliases"))
  ]

storeAliasParser :: Parser TricuArgs
storeAliasParser = subparser $ mconcat
  [ command "list" (info (storeAliasListParser <**> helper)
      (progDesc "List aliases by kind"))
  , command "get" (info (storeAliasGetParser <**> helper)
      (progDesc "Resolve an alias by kind and name"))
  ]

-- ---------------------------------------------------------------------------
-- Entry point
-- ---------------------------------------------------------------------------

main :: IO ()
main = do
  appArgs <- execParser $ info (tricuParser <**> helper)
    ( fullDesc
   <> progDesc "Exploring Tree Calculus"
   <> header versionStr
    )
  let mGlobalStore = globalStore appArgs
      args = applyGlobalStore mGlobalStore (appCommand appArgs)
  case args of
    Repl               -> runReplWithStore mGlobalStore
    Check {}           -> runCheck args
    Eval {}            -> runEval args
    ArboricxCompile {} -> runCompile args
    ArboricxImport {}  -> runImport args
    ArboricxExport {}  -> runExport args
    StoreAliasList {}  -> runStoreAliasList args
    StoreAliasGet {}   -> runStoreAliasGet args


-- ---------------------------------------------------------------------------
-- Command runners
-- ---------------------------------------------------------------------------

applyGlobalStore :: Maybe FilePath -> TricuArgs -> TricuArgs
applyGlobalStore mGlobal args = case args of
  Repl -> Repl
  Check {} -> args { checkStore = preferLocal (checkStore args) }
  Eval {} -> args { evalStore = preferLocal (evalStore args) }
  ArboricxCompile {} -> args { compileStore = preferLocal (compileStore args) }
  ArboricxImport {} -> args { importStore = preferLocal (importStore args) }
  ArboricxExport {} -> args { exportStore = preferLocal (exportStore args) }
  StoreAliasList {} -> args { storePathOpt = preferLocal (storePathOpt args) }
  StoreAliasGet {} -> args { storePathOpt = preferLocal (storePathOpt args) }
  where
    preferLocal local = case local of
      Just _  -> local
      Nothing -> mGlobal

runRepl :: IO ()
runRepl = runReplWithStore Nothing

runReplWithStore :: Maybe FilePath -> IO ()
runReplWithStore mStore = do
  putStrLn "Welcome to the tricu REPL"
  putStrLn "You may exit with `CTRL+D` or the `!exit` command."
  case mStore of
    Nothing -> repl
    Just store -> replWithStore (StorePath store)

runCheck :: TricuArgs -> IO ()
runCheck opts = do
  output <- case checkStore opts of
    Nothing -> checkFile (checkInput opts)
    Just storePath -> checkFileWithStore (StorePath storePath) (checkInput opts)
  putStrLn output

evaluateCheckedIOFile :: StorePath -> ContractMode -> Env -> FilePath -> IO Env
evaluateCheckedIOFile store mode env filePath = do
  loaded <- loadFileWithStoreMode mode store filePath
  checkedAst <- case instrumentIOContinuations (loadedAst loaded) of
    Left err -> die err
    Right asts -> pure asts
  viewEnv <- evaluateFileWithStore (Just store) "./lib/view.tri"
  pure $ evalTricu (Map.unions [viewEnv, loadedImports loaded, env]) checkedAst

runEval :: TricuArgs -> IO ()
runEval opts = do
  let files = evalFiles opts
      form  = evalFormat opts
      out   = evalOutput opts
  resultT <- case files of
    [] -> do
      input <- getContents
      let env = evalTricu Map.empty (parseTricu input)
      return $ result env
    _  -> do
      mStoreOpt <- traverse (pure . StorePath) (evalStore opts)
      let contractMode = if evalUnchecked opts then IgnoreContracts else EnforceContracts
      finalEnv <- if evalIo opts && contractMode == EnforceContracts
        then do
          store <- maybe defaultStorePath pure mStoreOpt
          foldM (evaluateCheckedIOFile store contractMode) Map.empty files
        else foldM (evaluateFileWithContextWithStoreAndMode contractMode mStoreOpt) Map.empty files
      return $ mainResult finalEnv
  finalT <- if evalIo opts
    then do
      let perms = IOPermissions
            { allowRead = evalAllowRead opts
            , allowWrite = evalAllowWrite opts
            , allowReadAll = evalUnsafeIo opts || evalAllowReadAll opts
            , allowWriteAll = evalUnsafeIo opts || evalAllowWriteAll opts
            }
      result <- runIO perms resultT
      case result of
        Left err  -> die $ "IO error: " ++ err
        Right val -> pure val
    else return resultT
  writeOutput out (formatT form finalT)

runCompile :: TricuArgs -> IO ()
runCompile opts = do
  let input = compileInput opts
      out   = compileOutput opts
      names = compileNames opts
      mStore = StorePath <$> compileStore opts
  when (null out)   $ die "tricu arboricx compile: --output is required"
  when (null input) $ die "tricu arboricx compile: input file is required"
  let nameTexts = if null names then [] else map T.pack names
  compileFileWithStore mStore input out nameTexts

runImport :: TricuArgs -> IO ()
runImport opts = do
  let file = importFile opts
  when (null file) $ die "tricu arboricx import: input file is required"
  store <- resolveStorePath (importStore opts)
  bundleData <- BL.readFile file
  roots <- unpackBundleToStore store (BL.toStrict bundleData)
  mapM_ (\(name, root) ->
    writeAlias store NameAlias name (treeTermRef root)) roots
  let manifest = ModuleManifest []
        [ ModuleExport
            name
            (treeTermRef root)
            "arboricx.abi.tree.v1"
            Nothing
            Nothing
        | (name, root) <- roots
        ]
      moduleName = T.pack $ maybe (takeBaseName file) id (importModule opts)
  manifestHash <- putManifest store manifest
  writeAlias store ModuleAlias moduleName (ObjectRef (unDomain manifestDomain) manifestHash)
  putStrLn $ "Imported " ++ show (length roots) ++ " root(s):"
  mapM_ (\(name, root) -> putStrLn $ "  " ++ T.unpack name ++ " -> " ++ T.unpack root) roots
  putStrLn $ "Created module alias " ++ T.unpack moduleName ++ " -> " ++ T.unpack manifestHash

runExport :: TricuArgs -> IO ()
runExport opts =
  if dag opts
    then runExportDag opts
    else runExportBundle opts

runExportBundle :: TricuArgs -> IO ()
runExportBundle opts = do
  let targets = exportTargets opts
      modules = exportModules opts
      out     = exportOutput opts
      names   = exportNames opts
      allFlag = exportAll opts
      splitFlag = exportSplit opts
  when (null out) $ die "tricu arboricx export: --output is required"
  when (null targets && null modules && not allFlag) $
    die "tricu arboricx export: at least one --target, --module, or --all is required"
  when (splitFlag && not (null names)) $
    die "tricu arboricx export --split: --name is not supported; split bundles use their export names"
  store <- resolveStorePath (exportStore opts)
  allEntries <- if allFlag then resolveAllNameExports store else pure []
  targetRoots <- mapM (resolveStoreTarget store) targets
  moduleRoots <- concat <$> mapM (resolveModuleExports store) modules
  let targetEntries = zip (defaultExportNames (length targetRoots)) targetRoots
      entries = allEntries ++ targetEntries ++ moduleRoots
      expNames = if null names then map fst entries else map T.pack names
  when (null entries) $
    die "tricu arboricx export: no tree-term exports found"
  when (length expNames /= length entries) $
    die "tricu arboricx export: number of --name values must match number of exported roots"
  if splitFlag
    then runExportBundleSplit store out (zip expNames (map snd entries))
    else do
      bundle <- packBundleFromStore store (zip expNames (map snd entries))
      let bundleData = encodeBundle bundle
      BL.writeFile out (BL.fromStrict bundleData)
      putStrLn $ "Exported bundle with " ++ show (length entries) ++ " export(s) to " ++ out
      putStrLn $ "  nodes: " ++ show (Seq.length (bundleNodes bundle))
      putStrLn $ "  size: " ++ show (BS.length bundleData) ++ " bytes"

runExportBundleSplit :: StorePath -> FilePath -> [(T.Text, ObjectHash)] -> IO ()
runExportBundleSplit store outDir entries = do
  createDirectoryIfMissing True outDir
  written <- forM (zip [0 :: Int ..] entries) $ \(i, (name, root)) -> do
    bundle <- packBundleFromStore store [(name, root)]
    let bundleData = encodeBundle bundle
        path = outDir </> splitBundleFileName i name
    BL.writeFile path (BL.fromStrict bundleData)
    pure (path, Seq.length (bundleNodes bundle), BS.length bundleData)
  putStrLn $ "Exported " ++ show (length written) ++ " split bundle(s) to " ++ outDir
  mapM_ (\(path, nodeCount, byteCount) ->
    putStrLn $ "  " ++ path ++ " (nodes: " ++ show nodeCount ++ ", size: " ++ show byteCount ++ " bytes)") written

runStoreAliasList :: TricuArgs -> IO ()
runStoreAliasList opts = do
  store <- resolveStorePath (storePathOpt opts)
  aliases <- listAliases store (storeAliasKind opts)
  mapM_ (\(name, ref) -> putStrLn $ T.unpack name ++ " -> " ++ formatObjectRef ref) aliases

runStoreAliasGet :: TricuArgs -> IO ()
runStoreAliasGet opts = do
  store <- resolveStorePath (storePathOpt opts)
  mRef <- readAlias store (storeAliasKind opts) (T.pack $ storeAliasName opts)
  case mRef of
    Nothing  -> die $ "alias not found: " ++ storeAliasName opts
    Just ref -> putStrLn $ storeAliasName opts ++ " -> " ++ formatObjectRef ref

runExportDag :: TricuArgs -> IO ()
runExportDag opts = do
  let targets = exportTargets opts
      modules = exportModules opts
      out     = exportOutput opts
  unless (null modules) $
    die "tricu arboricx export --dag: --module is only supported for bundle export"
  case targets of
    [target] -> do
      store <- resolveStorePath (exportStore opts)
      root <- resolveStoreTarget store target
      maybeTerm <- getTreeTerm store root
      case maybeTerm of
        Nothing -> die $ "Term not found: " ++ target
        Just term -> do
          let (rootIdx, nodes) = Research.exportDag term
              output = unlines $
                show rootIdx :
                map (\(tag, refs) -> unwords (tag : map show refs)) nodes
          writeOutput out output
    [] -> die "tricu arboricx export --dag: exactly one --target is required"
    _  -> die "tricu arboricx export --dag: exactly one --target is required"

-- ---------------------------------------------------------------------------
-- Helpers
-- ---------------------------------------------------------------------------

resolveStorePath :: Maybe FilePath -> IO StorePath
resolveStorePath (Just path) = return (StorePath path)
resolveStorePath Nothing = do
  home <- getHomeDirectory
  return (StorePath (home </> ".tricu" </> "store"))

treeTermRef :: ObjectHash -> ObjectRef
treeTermRef = ObjectRef (unDomain treeTermDomain)

resolveStoreTarget :: StorePath -> String -> IO ObjectHash
resolveStoreTarget store target = do
  mAlias <- readAlias store NameAlias (T.pack target)
  let root = maybe (T.pack target) objectRefHash mAlias
  mTree <- getTreeTerm store root
  case mTree of
    Just _ -> return root
    Nothing -> die $ "Term not found in store: " ++ target

resolveAllNameExports :: StorePath -> IO [(T.Text, ObjectHash)]
resolveAllNameExports store = do
  aliases <- sortOn fst <$> listAliases store NameAlias
  fmap concat $ mapM exportAlias aliases
  where
    exportAlias (name, ref)
      | objectRefKind ref /= unDomain treeTermDomain = pure []
      | otherwise = do
          mTree <- getTreeTerm store (objectRefHash ref)
          case mTree of
            Nothing -> die $ "Name alias tree term not found: " ++ T.unpack name
            Just _  -> pure [(name, objectRefHash ref)]

resolveModuleExports :: StorePath -> String -> IO [(T.Text, ObjectHash)]
resolveModuleExports store moduleTarget = do
  manifestHash <- resolveModuleManifestHash store moduleTarget
  mManifest <- getManifest store manifestHash
  manifest <- case mManifest of
    Nothing -> die $ "Module manifest not found in store: " ++ moduleTarget
    Just value -> return value
  mapM exportEntry (moduleManifestExports manifest)
  where
    exportEntry ex = do
      let ref = moduleExportObject ex
      unless (objectRefKind ref == unDomain treeTermDomain) $
        die $ "Unsupported module export object kind for " ++ T.unpack (moduleExportName ex) ++ ": " ++ T.unpack (objectRefKind ref)
      mTree <- getTreeTerm store (objectRefHash ref)
      case mTree of
        Nothing -> die $ "Module export tree term not found: " ++ T.unpack (moduleExportName ex)
        Just _  -> return (moduleExportName ex, objectRefHash ref)

resolveModuleManifestHash :: StorePath -> String -> IO ObjectHash
resolveModuleManifestHash store moduleTarget = do
  mAlias <- readAlias store ModuleAlias (T.pack moduleTarget)
  case mAlias of
    Just ref -> do
      unless (objectRefKind ref == unDomain manifestDomain) $
        die $ "Module alias does not point at a module manifest: " ++ moduleTarget
      return (objectRefHash ref)
    Nothing -> return (T.pack moduleTarget)

formatObjectRef :: ObjectRef -> String
formatObjectRef ref = T.unpack (objectRefKind ref) ++ " " ++ T.unpack (objectRefHash ref)

splitBundleFileName :: Int -> T.Text -> FilePath
splitBundleFileName i name = show i ++ "-" ++ sanitize (T.unpack name) ++ ".arboricx"
  where
    sanitize [] = "export"
    sanitize xs = case map safeChar xs of
      [] -> "export"
      ys -> ys
    safeChar c
      | isAlphaNum c || c == '-' || c == '_' || c == '.' = c
      | otherwise = '_'

writeOutput :: FilePath -> String -> IO ()
writeOutput path content
  | null path = putStr content
  | otherwise = writeFile path content

runTricuTEnv :: Env -> String -> T
runTricuTEnv env input =
  let asts     = parseTricu input
      finalEnv = evalTricu env asts
   in result finalEnv