module FileEval where

import Eval
import Lexer
import Parser
import Research

import Data.List       (partition)
import Data.Maybe      (mapMaybe)
import Control.Monad   (foldM)
import System.IO
import System.FilePath (takeDirectory, normalise, (</>))

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

extractMain :: Env -> Either String T
extractMain env =
  case Map.lookup "main" env of
    Just result -> Right result
    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 err)
    Right ast -> do
      processedAst <- preprocessFile filePath
      let finalEnv = evalTricu Map.empty processedAst
      case extractMain finalEnv of
        Right result -> return result
        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 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 err)
    Right ast -> do
      ast <- preprocessFile filePath
      pure $ evalTricu 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 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

nsDefinitions :: String -> [TricuAST] -> [TricuAST]
nsDefinitions moduleName = map (nsDefinition moduleName)

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)
  | isPrefixed name = SVar name
  | otherwise = SVar (nsVariable moduleName name)
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)
  | isPrefixed name = SDef name args (nsBody moduleName body)
  | otherwise = SDef (nsVariable moduleName name)
                        args (nsBody moduleName body)
nsBody _ other = other

nsBodyScoped :: String -> [String] -> TricuAST -> TricuAST
nsBodyScoped moduleName args body = case body of
  SVar name ->
    if name `elem` args
      then SVar name
      else nsBody moduleName (SVar name)
  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