module Main where

import ContentStore (initContentStoreWithPath, loadEnvironment, loadTerm, loadTree, resolveExportTarget)
import System.Exit (die)
import Server (runServerWithPath)
import Eval (evalTricu, evalTricuWithStore, mainResult, result)
import FileEval (evaluateFileWithContext, evaluateFileWithStore, compileFile)
import IODriver (IOPermissions(..), checkIOSentinel, runIO)
import Parser (parseTricu)
import REPL (repl)
import Research (T, EvaluatedForm(..), Env, formatT, exportDag)
import Wire (buildBundle, encodeBundle, importBundle, defaultExportNames, Bundle(..))

import Control.Monad (foldM, unless, when)
import Data.Text (unpack, pack)
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 Database.SQLite.Simple (Connection, close)

import qualified Data.Map as Map
import System.Environment (lookupEnv)

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

data TricuArgs
  = Repl
  | Eval
      { evalFiles      :: [FilePath]
      , evalFormat     :: EvaluatedForm
      , evalOutput     :: FilePath
      , evalDb         :: Maybe FilePath
      , evalIo         :: Bool
      , evalAllowRead  :: [FilePath]
      , evalAllowWrite :: [FilePath]
      }
  | ArboricxCompile
      { compileInput   :: FilePath
      , compileOutput  :: FilePath
      , compileNames   :: [String]
      , compileDb      :: Maybe FilePath
      }
  | ArboricxImport
      { importFile :: FilePath
      , importDb   :: Maybe FilePath
      }
  | ArboricxExport
      { exportTargets :: [String]
      , exportOutput  :: FilePath
      , exportNames   :: [String]
      , exportDb      :: Maybe FilePath
      , dag           :: Bool
      }
  | ArboricxServe
      { serveHost :: String
      , servePort :: Int
      , serveDb   :: 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
  _         -> Left $ "Unknown format: " ++ s ++ ". Expected: tree, fsl, ast, ternary, ascii, decode"

evalParser :: Parser TricuArgs
evalParser = Eval
  <$> many (argument str (metavar "FILE..."))
  <*> option readEvaluatedForm
      ( long "format"
     <> short 'f'
     <> metavar "FORM"
     <> value Tree
     <> help "Output format: tree, fsl, ast, ternary, ascii, decode"
      )
  <*> option str
      ( long "output"
     <> short 'o'
     <> metavar "FILE"
     <> value ""
     <> help "Write output to file instead of stdout"
      )
  <*> optional (option str
      ( long "db"
     <> short 'd'
     <> metavar "PATH"
     <> help "Content store database path"
      ))
  <*> 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)"
      ))

compileParser :: Parser TricuArgs
compileParser = ArboricxCompile
  <$> option str
      ( long "file"
     <> short 'f'
     <> metavar "FILE"
     <> value ""
     <> help "Input .tri source file"
      )
  <*> 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)"
      ))
  <*> optional (option str
      ( long "db"
     <> short 'd'
     <> metavar "PATH"
     <> help "Content store database path"
      ))

importParser :: Parser TricuArgs
importParser = ArboricxImport
  <$> option str
      ( long "file"
     <> short 'f'
     <> metavar "FILE"
     <> value ""
     <> help "Bundle file to import"
      )
  <*> optional (option str
      ( long "db"
     <> short 'd'
     <> metavar "PATH"
     <> help "Content store database path"
      ))

exportParser :: Parser TricuArgs
exportParser = ArboricxExport
  <$> many (option str
      ( long "target"
     <> short 't'
     <> metavar "TARGET"
     <> help "Target hash or name (repeatable)"
      ))
  <*> 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 "db"
     <> short 'd'
     <> metavar "PATH"
     <> help "Content store database path"
      ))
  <*> switch
      ( long "dag"
     <> help "Export as a topologically-sorted DAG node table instead of a bundle"
      )

serveParser :: Parser TricuArgs
serveParser = ArboricxServe
  <$> option str
      ( long "host"
     <> metavar "HOST"
     <> value "127.0.0.1"
     <> help "Host to bind the server to"
      )
  <*> option auto
      ( long "port"
     <> short 'p'
     <> metavar "PORT"
     <> value 8787
     <> help "HTTP port to listen on"
      )
  <*> optional (option str
      ( long "db"
     <> short 'd'
     <> metavar "PATH"
     <> help "Content store database path"
      ))

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

tricuParser :: Parser TricuArgs
tricuParser = (subparser topCommands <|> pure Repl)
  <**> infoOption versionStr (long "version" <> help "Show version")
  where
    topCommands = mconcat
      [ 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"))
      ]

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"))
  , command "serve" (info (serveParser <**> helper)
      (progDesc "Start a read-only HTTP server for Arboricx bundles"))
  ]

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

main :: IO ()
main = do
  args <- execParser $ info (tricuParser <**> helper)
    ( fullDesc
   <> progDesc "Exploring Tree Calculus"
   <> header versionStr
    )
  case args of
    Repl               -> runRepl
    Eval {}            -> runEval args
    ArboricxCompile {} -> runCompile args
    ArboricxImport {}  -> runImport args
    ArboricxExport {}  -> runExport args
    ArboricxServe {}   -> runServe args

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

runRepl :: IO ()
runRepl = do
  putStrLn "Welcome to the tricu REPL"
  putStrLn "You may exit with `CTRL+D` or the `!exit` command."
  repl

runEval :: TricuArgs -> IO ()
runEval opts = do
  let files = evalFiles opts
      form  = evalFormat opts
      out   = evalOutput opts
  mconn <- case evalDb opts of
    Just dbPath -> Just <$> initContentStoreWithPath (Just dbPath)
    Nothing     -> do
      mDbPath <- lookupEnv "TRICU_DB_PATH"
      case mDbPath of
        Just _  -> Just <$> initContentStoreWithPath Nothing
        Nothing -> return Nothing
  resultT <- case files of
    [] -> do
      input <- getContents
      env <- evalTricuWithStore mconn Map.empty (parseTricu input)
      return $ result env
    _  -> do
      finalEnv <- foldM (evaluateFileWithStore mconn) Map.empty files
      return $ mainResult finalEnv
  finalT <- if evalIo opts
    then case checkIOSentinel resultT of
      Right (1, action) -> do
        let perms = IOPermissions (evalAllowRead opts) (evalAllowWrite opts)
        runIO perms action
      Right (v, _) -> die $ "Unsupported IO ABI version: " ++ show v
      Left err -> die $ "IO mode requested but " ++ err
    else return resultT
  case mconn of
    Just conn -> close conn
    Nothing   -> return ()
  writeOutput out (formatT form finalT)

runCompile :: TricuArgs -> IO ()
runCompile opts = do
  let input = compileInput opts
      out   = compileOutput opts
      names = compileNames 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
  compileFile input out nameTexts

runImport :: TricuArgs -> IO ()
runImport opts = do
  let file = importFile opts
  when (null file) $ die "tricu arboricx import: input file is required"
  withContentStore (importDb opts) $ \conn -> do
    bundleData <- BL.readFile file
    roots <- map T.unpack <$> importBundle conn (BL.toStrict bundleData)
    putStrLn $ "Imported " ++ show (length roots) ++ " root(s):"
    mapM_ (\r -> putStrLn $ "  " ++ r) roots

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

runExportBundle :: TricuArgs -> IO ()
runExportBundle opts = do
  let targets = exportTargets opts
      out     = exportOutput opts
      names   = exportNames opts
  when (null out)     $ die "tricu arboricx export: --output is required"
  when (null targets) $ die "tricu arboricx export: at least one --target is required"
  withContentStore (exportDb opts) $ \conn -> do
    terms <- mapM (\t -> do
      (h, _) <- resolveExportTarget conn t
      maybeTree <- loadTree conn h
      case maybeTree of
        Nothing -> die $ "Term not found in store: " ++ t
        Just tree -> return tree) targets
    let expNames = if null names
          then defaultExportNames (length terms)
          else map T.pack names
    when (length expNames /= length terms) $
      die "tricu arboricx export: number of --name values must match number of TARGETs"
    let namedTerms = zip expNames terms
        bundle = buildBundle namedTerms
        bundleData = encodeBundle bundle
    BL.writeFile out (BL.fromStrict bundleData)
    putStrLn $ "Exported bundle with " ++ show (length namedTerms) ++ " export(s) to " ++ out
    putStrLn $ "  nodes: " ++ show (Seq.length (bundleNodes bundle))
    putStrLn $ "  size: " ++ show (BS.length bundleData) ++ " bytes"

runExportDag :: TricuArgs -> IO ()
runExportDag opts = do
  let targets = exportTargets opts
      out     = exportOutput opts
  case targets of
    [target] -> withContentStore (exportDb opts) $ \conn -> do
      maybeTerm <- loadTerm conn target
      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"

runServe :: TricuArgs -> IO ()
runServe opts = do
  let hostStr = serveHost opts
      portNum = servePort opts
  putStrLn $ "Starting Arboricx bundle server on " ++ hostStr ++ ":" ++ show portNum
  putStrLn $ "  GET /bundle/hash/:hash   -- primary endpoint"
  putStrLn $ "  GET /bundle/name/:name   -- convenience endpoint"
  putStrLn $ "  Content-Type: application/vnd.arboricx.bundle"
  runServerWithPath (serveDb opts) hostStr portNum

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

withContentStore :: Maybe FilePath -> (Connection -> IO a) -> IO a
withContentStore mPath act = do
  conn <- initContentStoreWithPath mPath
  result <- act conn
  close conn
  return result

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