module Main where

import Prelude hiding (catch)
import Control.Applicative
import Control.Exception
import Control.Monad.Error.Class
import Control.Monad.Reader
import Control.Monad.Trans
import Data.List
import Data.Maybe
import Language.Haskell.Interpreter.GHC
import Language.Haskell.Parser
import Language.Haskell.Syntax
import System.Directory
import System.Environment

type InterleavableReader a = Reader (String, [String]) a
type InterleavableReaderT a = ReaderT (String, [String]) IO a
type TypeToDefinition = Reader (String, String, HsType) String

addTypes :: InterleavableReaderT String
addTypes
  = do
    reader <- ask
    liftIO
      $ do
        session <- newSession
        unlines <$> zipWith (flip (++)) (runReader newFunctionsWithoutType reader) <$> withSession session
          (
            catchInterpreterError
            $ setUseLanguageExtensions True
            >> loadModules [runReader fileName reader]
            >> setTopLevelModules [runReader newModule reader]
            >> mapM typeOfHaskell (runReader newFunctionNames reader)
          )

argumentList :: Char -> Int -> String
argumentList char number = unwords_ (map ((char :) . show) $ tail [0 .. number])

arguments :: HsType -> [HsType]
arguments (HsTyFun type_ rest) = type_ : arguments rest
arguments _ = []

bug :: String
bug = "\nThis shouldn't happen.  Please report the bug to marcot@riseup.net."

catchInterpreterError :: Interpreter a -> Interpreter a
catchInterpreterError function = catchError function treatInterpreterError

createModule :: InterleavableReader String -> InterleavableReaderT ()
createModule code
  = do
    reader <- ask
    liftIO
      $ createDirectoryIfMissing True (runReader directoryName reader)
      >> writeFile (runReader fileName reader) (runReader code reader)

class Decls a where
  getImports :: a -> [String]
  isIO :: a -> Bool

instance Decls HsDecl where
  getImports (HsTypeSig _ _ (HsQualType context type_))
    = union (foldr union [] $ map (getImports . fst) context) $ getImports type_
  getImports _ = error $ "getImports (_ :: HsDecl)." ++ bug

  isIO (HsTypeSig _ _ (HsQualType _ type_)) = isIO type_
  isIO _ = error $ "isIO (_ :: HsDecl)." ++ bug

instance Decls HsType where
  getImports (HsTyFun type_ rest) = union (getImports type_) (getImports rest)
  getImports (HsTyTuple list) = foldr union [] $ map getImports list
  getImports (HsTyApp type_ rest) = union (getImports type_) (getImports rest)
  getImports (HsTyCon hsName) = getImports hsName
  getImports _ = []

  isIO (HsTyFun type_ rest) = isIO type_ || isIO rest
  isIO (HsTyApp (HsTyCon hsName) _) = isIO hsName
  isIO _ = False

instance Decls HsQName where
  getImports (Qual (Module module_) _) = [module_]
  getImports _ = []

  isIO (Qual _ (HsIdent "IO")) = True
  isIO (UnQual (HsIdent "IO")) = True
  isIO _ = False

definitionBody :: TypeToDefinition
definitionBody
  = do
    suffix <- definitionSuffix
    (_, _, type_) <- ask
    case returnsFunction type_ of
      Nothing
        -> do
          prefix <- definitionPrefix " "
          return $ " =" ++ prefix ++ suffix
      Just result
        -> do
          prefix <- definitionPrefix "\n  "
          return
            $ "\n  =" ++ prefix ++ "do\n    function <- " ++ suffix
            ++ "\n    return " ++ treatArgument "liftIO" "function" result

definitionPrefix :: String -> TypeToDefinition
definitionPrefix separator
  = do
    (_, _, type_) <- ask
    return
      $ " "
      ++ (if needsInterleavable type_ then "embed $" ++ separator ++ "\\ buffer ->" else "liftIO" ++ separator ++ "$")
      ++ " "

definitionSuffix :: TypeToDefinition
definitionSuffix
  = do
    (module_, hsName, type_) <- ask
    argumentNames <- getArgumentNames
    return
      $ module_ ++ "." ++ hsName
      ++ if needsInterleavable type_
      then unwords_ $ zipWith (treatArgument "callback buffer") (words argumentNames) $ arguments type_
      else argumentNames

directoryName :: InterleavableReader String
directoryName = reverse <$> dropWhile (/= '/') <$> reverse <$> slash

draftBody :: InterleavableReader String
draftBody = unlines <$> newFunctionsWithoutType

extraImports :: InterleavableReader [String]
extraImports = foldr union [] <$> map getImports <$> newFunctionOldDecls

extraImportsWithPrelude :: InterleavableReader [String]
extraImportsWithPrelude = union ["Prelude"] <$> extraImports

fileName :: InterleavableReader String
fileName = (++ ".hs") <$> slash

getArgumentNames :: TypeToDefinition
getArgumentNames
  = do
    (_, _, type_) <- ask
    return $ argumentList 'a' $ numArguments type_

getExports :: String -> IO [String]
getExports module_
  = do
    session <- newSession
    withSession session $ catchInterpreterError
      $ setImports [module_] >> getModuleFunctions module_ >>= mapM typeOfHaskell

getExtraImportsHaskell :: InterleavableReaderT String
getExtraImportsHaskell
  = do
    reader <- ask
    liftIO
      $ do
        session <- newSession
        unlines
          <$> zipWith (\ import_ exports -> runReader (getImportHaskell import_ exports) reader)
          (runReader extraImportsWithPrelude reader)
          <$> withSession session
          (catchInterpreterError $ mapM getModuleFunctions $ runReader extraImportsWithPrelude reader)

getHeader :: InterleavableReader String
getHeader
  = do
    module_ <- asks fst
    newModule_ <- newModule
    newFunctionNames_ <- newFunctionNames
    return
      $ "{-# OPTIONS_GHC -fno-monomorphism-restriction #-}\n{-# LANGUAGE FlexibleContexts #-}\n"
      ++ "-- This code was generated by interlavableGen.  Comments: Marco Túlio Gontijo e Silva <marcot@riseup.net>\n"
      ++ "module " ++ newModule_ ++ parenthesisList (("module " ++ module_) : newFunctionNames_) ++ "\n  where\n\n"

getImportHaskell :: String -> [String] -> InterleavableReader String
getImportHaskell import_ exports
  = do
    importHidings_ <- importHidings exports
    return $ "import " ++ import_ ++ if null importHidings_ then "" else " hiding" ++ parenthesisList importHidings_

getImportsHaskell :: String -> InterleavableReader String
getImportsHaskell extraImportsHaskell
  = do
    module_ <- asks fst
    newFunctionNames_ <- newFunctionNames
    return
      $ "import Control.Monad.Trans\nimport Control.Monad.Trans.InterleavableIO\n"
      ++ "import " ++ module_ ++ " hiding" ++ parenthesisList newFunctionNames_ ++ "\nimport qualified " ++ module_ ++ "\n"
      ++ extraImportsHaskell

getModuleFunctions :: String -> Interpreter [String]
getModuleFunctions module_ = catMaybes <$> map isFunction <$> getModuleExports module_

importHidings :: [String] -> InterleavableReader [String]
importHidings exports = intersect exports <$> union ["liftIO", "embed", "callback"] <$> newFunctionNames

interleavableModule :: (Bool, String) -> IO ()
interleavableModule parameters @ (_, module_) = catch (interleavable parameters) $ treatInterleavable module_

interleavable :: (Bool, String) -> IO ()
interleavable (withType, module_)
  = do
    moduleExports <- getExports module_
    liftIO $ flip runReaderT (module_, moduleExports)
      $ do
        extraImportsHaskell <- getExtraImportsHaskell
        createModule $ moduleCode extraImportsHaskell (runReader draftBody (module_, moduleExports))
        when withType $ moduleCode extraImportsHaskell <$> addTypes >>= createModule

isFunction :: ModuleElem -> Maybe String
isFunction (Fun function) = Just function
isFunction _ = Nothing

main :: IO ()
main = getArgs >>= parseOptions

moduleCode :: String -> String -> InterleavableReader String
moduleCode extraImportsHaskell body
  = do
    imports <- getImportsHaskell extraImportsHaskell
    header <- getHeader
    return $ header ++ imports ++ body

nameWithType :: HsDecl -> (String, HsType)
nameWithType (HsTypeSig _ [(HsIdent hsName)] (HsQualType _ type_)) = (hsName, type_)
nameWithType _ = error $ "nameWithType _" ++ bug

needsInterleavable :: HsType -> Bool
needsInterleavable (HsTyFun type_ rest) = isIO type_ || needsInterleavable rest
needsInterleavable _ = False

newFunctionNames :: InterleavableReader [String]
newFunctionNames = map fst <$> newFunctionNamesWithOldTypes

newFunctionNamesWithOldTypes :: InterleavableReader [(String, HsType)]
newFunctionNamesWithOldTypes =  map nameWithType <$> newFunctionOldDecls

newFunctionOldDecls :: InterleavableReader [HsDecl]
newFunctionOldDecls = filter isIO <$> parsedBody

newFunctionsWithoutType :: InterleavableReader [String]
newFunctionsWithoutType = newFunctionNamesWithOldTypes >>= mapM typeToDefinition

newModule :: InterleavableReader String
newModule = ("Control.Monad.Trans.InterleavableIO." ++) <$> asks fst

numArguments :: HsType -> Int
numArguments = length . arguments

parenthesisList :: [String] -> String
parenthesisList list = "\n  ( " ++ intercalate "\n  , " list ++ "\n  )"

parsedBody :: InterleavableReader [HsDecl]
parsedBody
  = do
    moduleExports <- asks snd
    case parseModule $ unlines moduleExports of
      ParseOk (HsModule _ _ _ _ parsed) -> return parsed
      ParseFailed (SrcLoc _ line _) message
        -> error
        $ "parsed: " ++ moduleExports !! pred line ++ ": " ++ message ++ "." ++ bug

parseOptions :: [String] -> IO ()
parseOptions options
  | elem "--no-type-signatures" options
    = mapM_ interleavableModule $ map ((,) False) $ filter (/= "--no-type-signatures") options
  | elem "--help" options || elem "-h" options || null options
    = do
      progName <- getProgName
      putStrLn $ "Usage:\n  " ++ progName ++ " [--no-type-signatures] <Module>...\n  " ++ progName ++ " -h|--help"
  | otherwise = mapM_ interleavableModule $ map ((,) True) options

replace :: Eq a => a -> a -> [a] -> [a]
replace from to = map (\ x -> if x == from then to else x)

returnsFunction :: HsType -> Maybe HsType
returnsFunction (HsTyFun _ rest) = returnsFunction rest
returnsFunction (HsTyApp (HsTyCon (Qual _ (HsIdent "IO"))) result @ (HsTyFun _ _)) = Just result
returnsFunction _ = Nothing

slash :: InterleavableReader String
slash = replace '.' '/' <$> newModule

treatArgument :: String -> String -> HsType -> String
treatArgument caller argumentName argument
  | isIO argument && numArguments argument > 0
    = "(\\" ++ (argumentList 'b' $ numArguments argument) ++ " -> "
    ++ caller ++ " $ " ++ argumentName ++ " " ++ (argumentList 'b' $ numArguments argument) ++ ")"
  | isIO argument = "(" ++ caller ++ " " ++ argumentName ++ ")"
  | otherwise = argumentName

treatInterleavable :: String -> Exception -> IO ()
treatInterleavable module_ exception
  = do
    progName <- getProgName
    putStrLn $ progName ++ ": Error: " ++ module_ ++ ": " ++ show exception

treatInterpreterError :: InterpreterError -> Interpreter a
treatInterpreterError (UnknownError string)= error $ "InterpreterError.UnknownError: " ++ string
treatInterpreterError (WontCompile list) = error $ "InterpreterError.WontCompile: List of GHC Errors: " ++ show list
treatInterpreterError (NotAllowed string) = error $ "InterpreterError.NotAllowed: " ++ string
treatInterpreterError (GhcException exception) = error $ "InterpreterError.GhcException: " ++ show exception

typeOfHaskell :: String -> Interpreter String
typeOfHaskell function
  = do
    type_ <- typeOf function
    return $ "\n" ++ function ++ " :: " ++ (replace '\n' ' ' type_)

typeToDefinition :: (String, HsType) -> InterleavableReader String
typeToDefinition (hsName, type_)
  =  do
     module_ <- asks fst
     return $ flip runReader (module_, hsName, type_)
       $ do
         argumentNames <- getArgumentNames
         definition <- definitionBody
         return $ '\n' : hsName ++ argumentNames ++ definition

unwords_ :: [String] -> String
unwords_ [] = ""
unwords_ list = " " ++ unwords list