module FileEval where

import Eval
import Lexer
import Parser
import Research

import Data.List       (partition)
import Control.Monad   (foldM)
import System.IO

import qualified Data.Map as Map
import qualified Data.Set as Set

evaluateFileResult :: FilePath -> IO T
evaluateFileResult filePath = do
  contents <- readFile filePath
  let tokens = lexTricu contents
  let moduleName = case parseProgram tokens of
        Right ((SModule name) : _) -> name
        _                          -> ""
  case parseProgram tokens of
    Left err -> errorWithoutStackTrace (handleParseError err)
    Right _ -> do
      ast <- preprocessFile filePath
      let finalEnv = mainAlias moduleName $ evalTricu Map.empty ast
      case Map.lookup "main" finalEnv of
        Just finalResult -> return finalResult
        Nothing -> errorWithoutStackTrace "No `main` function detected"

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

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

mainAlias :: String -> Env -> Env
mainAlias "" env = env
mainAlias moduleName env =
  case Map.lookup (moduleName ++ ".main") env of
    Just value -> Map.insert "main" value env
    Nothing    -> env

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

preprocessFile' :: Set.Set FilePath -> FilePath -> IO [TricuAST]
preprocessFile' inProgress filePath
  | filePath `Set.member` inProgress = 
      errorWithoutStackTrace $ "Encountered cyclic import: " ++ filePath
  | otherwise = do
      contents <- readFile filePath
      let tokens = lexTricu contents
      case parseProgram tokens of
        Left err -> errorWithoutStackTrace (handleParseError err)
        Right asts -> do
          let (moduleName, restAST) = extractModule asts
          let (imports, nonImports) = partition isImport restAST
          let newInProgress = Set.insert filePath inProgress
          importedASTs <- concat <$> mapM (processImport newInProgress) imports
          let namespacedAST = namespaceDefinitions moduleName nonImports
          pure $ importedASTs ++ namespacedAST
  where
    extractModule :: [TricuAST] -> (String, [TricuAST])
    extractModule ((SModule name) : xs) = (name, xs)
    extractModule xs                    = ("", xs)

    isImport :: TricuAST -> Bool
    isImport (SImport _ _) = True
    isImport _             = False

    processImport :: Set.Set FilePath -> TricuAST -> IO [TricuAST]
    processImport inProgress (SImport filePath moduleName) = do
      importedAST <- preprocessFile' inProgress filePath
      pure $ namespaceDefinitions moduleName importedAST
    processImport _ _ = error "Unexpected non-import in processImport"

namespaceDefinitions :: String -> [TricuAST] -> [TricuAST]
namespaceDefinitions moduleName = map (namespaceDefinition moduleName)

namespaceDefinition :: String -> TricuAST -> TricuAST
namespaceDefinition "" def = def
namespaceDefinition moduleName (SDef name args body)
  | isPrefixed name = SDef name args (namespaceBody moduleName body)
  | otherwise = SDef (namespaceVariable moduleName name) 
                        args (namespaceBody moduleName body)
namespaceDefinition moduleName other =
  namespaceBody moduleName other

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

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

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

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