module FileEval
  ( preprocessFile
  , evaluateFile
  , evaluateFileWithContext
  , evaluateFileWithStore
  , evaluateFileResult
  , compileFile
  ) where

import Eval (evalTricu, evalTricuWithStore)
import Lexer
import Parser
import Research
import Wire (buildBundle, encodeBundle, decodeBundle, verifyBundle, Bundle(..))
import Database.SQLite.Simple (Connection)

import Data.List       (partition)
import Data.Maybe      (mapMaybe)
import System.FilePath (takeDirectory, normalise, (</>))
import System.Exit (die)

import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Sequence as Seq
import qualified Data.Text as T

extractMain :: Env -> Either String T
extractMain env =
  case Map.lookup "main" env of
    Just evalResult -> Right evalResult
    Nothing -> Left "No `main` function detected"

processImports :: Set.Set FilePath -> FilePath -> FilePath -> [TricuAST]
                 -> Either String ([TricuAST], [(FilePath, String, FilePath)])
processImports seen _base currentPath asts =
  let (imports, nonImports) = partition isImp asts
      importPaths = mapMaybe getImportInfo imports
  in if currentPath `Set.member` seen
     then Left $ "Encountered cyclic import: " ++ currentPath
     else Right (nonImports, importPaths)
  where
    isImp (SImport _ _) = True
    isImp _ = False
    getImportInfo (SImport p n) = Just (p, n, makeRelativeTo currentPath p)
    getImportInfo _ = Nothing

evaluateFileResult :: FilePath -> IO T
evaluateFileResult filePath = do
  contents <- readFile filePath
  let tokens = lexTricu contents
  case parseProgram tokens of
    Left err -> errorWithoutStackTrace (handleParseError tokens err)
    Right _ast -> do
      processedAst <- preprocessFile filePath
      let finalEnv = evalTricu Map.empty processedAst
      case extractMain finalEnv of
        Right evalResult -> return evalResult
        Left err -> errorWithoutStackTrace err

evaluateFile :: FilePath -> IO Env
evaluateFile filePath = do
  contents <- readFile filePath
  let tokens = lexTricu contents
  case parseProgram tokens of
    Left err -> errorWithoutStackTrace (handleParseError tokens err)
    Right _ast -> do
      ast <- preprocessFile filePath
      pure $ evalTricu Map.empty ast

evaluateFileWithContext :: Env -> FilePath -> IO Env
evaluateFileWithContext env filePath = do
  contents <- readFile filePath
  let tokens = lexTricu contents
  case parseProgram tokens of
    Left err -> errorWithoutStackTrace (handleParseError tokens err)
    Right _ast -> do
      ast <- preprocessFile filePath
      pure $ evalTricu env ast

-- | File evaluation that lazily resolves missing names from the
--   content store instead of pre-loading the entire store into memory.
evaluateFileWithStore :: Maybe Connection -> Env -> FilePath -> IO Env
evaluateFileWithStore mconn env filePath = do
  contents <- readFile filePath
  let tokens = lexTricu contents
  case parseProgram tokens of
    Left err -> errorWithoutStackTrace (handleParseError tokens err)
    Right _ast -> do
      ast <- preprocessFile filePath
      evalTricuWithStore mconn env ast

preprocessFile :: FilePath -> IO [TricuAST]
preprocessFile p = preprocessFile' Set.empty p p

preprocessFile' :: Set.Set FilePath -> FilePath -> FilePath -> IO [TricuAST]
preprocessFile' seen base currentPath = do
  contents <- readFile currentPath
  let tokens = lexTricu contents
  case parseProgram tokens of
    Left err -> errorWithoutStackTrace (handleParseError tokens err)
    Right ast ->
      case processImports seen base currentPath ast of
        Left err -> errorWithoutStackTrace err
        Right (nonImports, importPaths) -> do
          let seen' = Set.insert currentPath seen
          imported <- concat <$> mapM (processImportPath seen' base) importPaths
          pure $ imported ++ nonImports
  where
    processImportPath _seen _base (_path, name, importPath) = do
      ast <- preprocessFile' _seen _base importPath
      pure $ map (nsDefinition (if name == "!Local" then "" else name))
           $ filter (not . isImp) ast
    isImp (SImport _ _) = True
    isImp _ = False

makeRelativeTo :: FilePath -> FilePath -> FilePath
makeRelativeTo f i =
    let d = takeDirectory f
    in normalise $ d </> i

nsDefinition :: String -> TricuAST -> TricuAST
nsDefinition "" def = def
nsDefinition moduleName (SDef name args body)
  | isPrefixed name = SDef name args (nsBody moduleName body)
  | otherwise = SDef (nsVariable moduleName name)
                        args (nsBody moduleName body)
nsDefinition moduleName other =
  nsBody moduleName other

nsBody :: String -> TricuAST -> TricuAST
nsBody moduleName (SVar name mhash)
  | isPrefixed name = SVar name mhash
  | otherwise = SVar (nsVariable moduleName name) mhash
nsBody moduleName (SApp func arg) =
  SApp (nsBody moduleName func) (nsBody moduleName arg)
nsBody moduleName (SLambda args body) =
  SLambda args (nsBodyScoped moduleName args body)
nsBody moduleName (SList items) =
  SList (map (nsBody moduleName) items)
nsBody moduleName (TFork left right) =
  TFork (nsBody moduleName left) (nsBody moduleName right)
nsBody moduleName (TStem subtree) =
  TStem (nsBody moduleName subtree)
nsBody moduleName (SDef name args body) =
  SDef (nsVariable moduleName name) args (nsBodyScoped moduleName args body)
nsBody _ other = other

nsBodyScoped :: String -> [String] -> TricuAST -> TricuAST
nsBodyScoped moduleName args body = case body of
  SVar name mhash ->
    if name `elem` args
      then SVar name mhash
      else nsBody moduleName (SVar name mhash)
  SApp func arg ->
    SApp (nsBodyScoped moduleName args func) (nsBodyScoped moduleName args arg)
  SLambda innerArgs innerBody ->
    SLambda innerArgs (nsBodyScoped moduleName (args ++ innerArgs) innerBody)
  SList items ->
    SList (map (nsBodyScoped moduleName args) items)
  TFork left right ->
    TFork (nsBodyScoped moduleName args left) (nsBodyScoped moduleName args right)
  TStem subtree ->
    TStem (nsBodyScoped moduleName args subtree)
  SDef name innerArgs innerBody ->
    SDef (nsVariable moduleName name) innerArgs (nsBodyScoped moduleName (args ++ innerArgs) innerBody)
  other -> other

isPrefixed :: String -> Bool
isPrefixed name = '.' `elem` name

nsVariable :: String -> String -> String
nsVariable "" name = name
nsVariable moduleName name = moduleName ++ "." ++ name

-- | Compile a tricu source file to a standalone Arboricx bundle.
-- Emits a canonical indexed bundle with no SHA-256 hashing.
compileFile :: FilePath -> FilePath -> [T.Text] -> IO ()
compileFile inputPath outputPath maybeNames = do
  env <- evaluateFile inputPath
  let defaultNames = ["main"]
      wantedNames  = if null maybeNames then defaultNames else maybeNames
      wantedNamesUnpacked = map T.unpack wantedNames
  compiledTerms <- mapM (\n -> case Map.lookup n env of
    Nothing -> die $ "No definition '" ++ n ++ "' found in " ++ inputPath
    Just t  -> return (T.pack n, t)) wantedNamesUnpacked
  let bundle = buildBundle compiledTerms
      bundleData = encodeBundle bundle
      nodeCount = Seq.length (bundleNodes bundle)
      bundleSize = BS.length bundleData
  BL.writeFile outputPath (BL.fromStrict bundleData)
  putStrLn $ "Compiled " ++ inputPath ++ " -> " ++ outputPath
  putStrLn $ "  exports: " ++ T.unpack (T.intercalate ", " (map fst compiledTerms))
  putStrLn $ "  nodes: " ++ show nodeCount
  putStrLn $ "  size: " ++ show bundleSize ++ " bytes"
  case decodeBundle bundleData of
    Left err -> putStrLn $ "  round-trip decode failed: " ++ err
    Right decoded -> case verifyBundle decoded of
      Left err -> putStrLn $ "  round-trip verify failed: " ++ err
      Right () -> putStrLn $ "  round-trip: OK"