-- Copyright 2012 gatoatigrado (nicholas tung) [ntung at ntung]
-- Licensed under the Apache License, Version 2.0 (the "License"); you may
-- not use this file except in compliance with the License. You may obtain a
-- copy of the License at http://www.apache.org/licenses/LICENSE-2.0 .

-- enable to remove excess imports {-# OPTIONS_GHC -fwarn-unused-imports #-}
{-# LANGUAGE Arrows,
             BangPatterns,
             DeriveDataTypeable,
             EmptyDataDecls,
             FlexibleContexts,
             FlexibleInstances,
             FunctionalDependencies,
             GADTs,
             GeneralizedNewtypeDeriving,
             ImpredicativeTypes,
             MultiParamTypeClasses,
             NamedFieldPuns,
             NoMonomorphismRestriction,
             RankNTypes,
             ScopedTypeVariables,
             StandaloneDeriving,
             TemplateHaskell,
             TypeFamilies,
             TypeOperators,
             TypeSynonymInstances,
             UndecidableInstances,
             ViewPatterns #-}

module Hz2.Preproc where

import Prelude hiding (id, (.))
import Control.Arrow
import Control.Category
import Control.Monad
import Control.Monad.Trans.Class

import Control.Applicative hiding (many, (<|>))

import qualified Data.Map as Map
import qualified Data.List as List
import qualified Data.Text as Text

import System.Environment
import System.Exit
import System.IO

import Text.Parsec hiding (spaces)
import Text.Parsec.Combinator

import qualified Text.Parsec.Token as P
import Text.Parsec.Language (haskellDef)

import Text.Printf

infixl 4 <++>
(<++>) = liftM2 (++)

haskell = P.makeTokenParser haskellDef
ident = P.identifier haskell
parens = P.parens haskell
commaSep = P.commaSep haskell

relist = List.intercalate ", " . map ((++ " α") . strip)
strip = Text.unpack . Text.strip . Text.pack
instance_trans_conds [] = ""
instance_trans_conds lst = "(" ++ List.intercalate ", " lst ++ ") => "

mk_suitable (strip -> nme) (strip -> cls) (relist -> conds) = "\n\
\data instance Constraints (" ++ cls ++ ") α =\n\
\    (" ++ conds ++ ") => " ++ nme ++ "\n\
\instance (" ++ conds ++ ") => Suitable (" ++ cls ++ ") α where\n\
\    constraints = " ++ nme ++ "\n"

mk_e_trans (strip -> name) (strip -> base) conds monad_conds (strip -> newtyp) =
    mk_suitable constr_name ("ExprTyp (" ++ name ++ ")") conds ++ "\n\
\instance " ++ mc ++ "HzExprTrans (ExprTyp (" ++ name ++ ")) where\n\
\    type PrimExpr (ExprTyp (" ++ name ++ ")) = " ++ base ++ "\n\
\    e_lift = fmap " ++ newtyp ++ "\n\
\    e_lower = fmap (\\(" ++ newtyp ++ " x) -> x)\n\
\    e_lower_suitably x@(ExprT v) f =\n\
\        withConstraintsOf v $ \\" ++ constr_name ++ " -> f (e_lower x)\n\
\instance " ++ mc ++ "BinaryCmpExpr (ExprTyp (" ++ name ++ "))\n\
\instance " ++ mc ++ "BinaryOpExpr (ExprTyp (" ++ name ++ "))\n\
\instance " ++ mc ++ "SketchConstrExpr (ExprTyp (" ++ name ++ "))\n\
\instance " ++ mc ++ "Arithmetic (ExprTyp (" ++ name ++ "))\n\
\instance " ++ mc ++ "VariablesExpr (ExprTyp (" ++ name ++ "))\n"

    where constr_name = newtyp ++ "_Constr"
          mc = instance_trans_conds monad_conds

mk_monad_trans (strip -> nme) (strip -> base) monad_conds (strip -> newtyp) = "\n\
\instance " ++ mc ++ " MonadTrans (" ++ nme ++ ") where\n\
\    type BaseMonad (" ++ nme ++ ") = " ++ base ++ "\n\
\    lift = " ++ newtyp ++ "\n\
\instance " ++ mc ++ " HzMonadUnliftTrans (" ++ nme ++ ") where\n\
\    unlift (" ++ newtyp ++ " x) = x\n"
    where
        mc = instance_trans_conds monad_conds

mk_array_trans
    (strip -> monadtyp) -- e.g. StateT 𝔪 α
    (strip -> base_array) -- e.g. ArrayTyp 𝔪 α
    (strip -> base_alloc) -- e.g. AllocTyp 𝔪 α
    conds -- e.g. Suitable (ArrayTyp 𝔪) α
    monad_conds -- e.g. HzArrayMonad 𝔪
    (strip -> newtyp_base) -- e.g. StateT
    = mk_suitable constr_name ("ArrayTyp (" ++ monadtyp ++ ")") conds ++
    "\n\
\instance " ++ mc ++ "HzArrayMonad (" ++ monadtyp ++ ") where\n\
\    data AllocTyp (" ++ monadtyp ++ ") α = " ++ alloc_newtyp ++ " ((" ++ base_alloc ++ ") α)\n\
\    data ArrayTyp (" ++ monadtyp ++ ") α = " ++ array_newtyp ++ " ((" ++ base_array ++ ") α)\n\
\    alloc_lower = fmap $ \\(" ++ alloc_newtyp ++ " x) -> x\n\
\    alloc_lift = fmap " ++ alloc_newtyp ++ "\n\
\    array_lower = fmap $ \\(" ++ array_newtyp ++ " x) -> x\n\
\    array_lift = fmap " ++ array_newtyp ++ "\n\
\    array_suitably m f = withConstraintsOf m $ \\" ++ constr_name ++ " -> f\n"
    where constr_name = newtyp_base ++ "_Array_Constr"
          array_newtyp = newtyp_base ++ "_Array"
          alloc_newtyp = newtyp_base ++ "_Alloc"
          mc = instance_trans_conds monad_conds





-- grab any expression within parenthesis
any_parens = pure "(" <++>
    parens (nonparen <++> option "" any_parens <++> nonparen) <++> pure ")"
    where nonparen = many (noneOf "()")

-- way to get spaces
spaces m = many space *> m <* many space

-- an entry of a comma-separated list
-- fails for a whitespace-only string
no_comma_ent = do
    lookAhead (many space >> notFollowedBy (space <|> char ')'))
    List.concat <$> many1 (many1 (noneOf ",()") <|> any_parens)

paren_list = parens (commaSep no_comma_ent)

-- a comma-separated list of named items
{---------------------------------------------------
-- optlist f = commaSep $ do
--     ident <- spaces ident
--     spaces (char '=')
--     f ident
----------------------------------------------------}
optlistent isEnd nme f = spaces (string nme) >> spaces (char '=') >> f'
    where f' | isEnd = f <* many space
             | otherwise = f <* spaces (char ',')

suitable_macro = do
    name <- ident
    spaces (char '=') >> string "SUITABLE" >> many space
    -- parens (many space >> return "")
    parens $ do
        lst <- map List.concat <$>
                commaSep (many (many1 (noneOf ",()") <|> any_parens))
        return (mk_suitable name (last lst) (init lst))

named_macro n f = do
    try (spaces (string n))
    parens f

expr_trans_macro = named_macro "EXPR_TRANS" $ do
    name <- optlistent False "name" no_comma_ent
    base <- optlistent False "base" no_comma_ent
    monad_conds <- optlistent False "monad_conds" paren_list
    newtyp <- optlistent True "newtyp" no_comma_ent
    return $ mk_e_trans name base
                 -- replace "extra_conds" with default value
                 [ ("Suitable (" ++ base ++ ")") ]
                 monad_conds
                 newtyp

array_trans_macro = named_macro "ARRAY_TRANS" $ do
    monadtyp <- optlistent False "monadtyp" no_comma_ent
    base_array <- optlistent False "base_array" no_comma_ent
    base_alloc <- optlistent False "base_alloc" no_comma_ent
    monad_conds <- optlistent False "monad_conds" paren_list
    newtyp_base <- optlistent True "newtyp_base" no_comma_ent
    return $ mk_array_trans monadtyp base_array base_alloc
                 -- replace "extra_conds" with default value
                 [ ("Suitable (" ++ base_array ++ ")") ]
                 monad_conds
                 newtyp_base

monad_trans_macro = named_macro "MONAD_TRANS" $ do
    name <- optlistent False "name" no_comma_ent
    base <- optlistent False "base" no_comma_ent
    monad_conds <- optlistent False "monad_conds" paren_list
    newtyp <- optlistent True "newtyp" no_comma_ent
    return $ mk_monad_trans name base monad_conds newtyp

hz_trans_macro = named_macro "HZ_TRANS" $ do
    name <- optlistent False "name" no_comma_ent
    monad_base <- optlistent False "monad_base" no_comma_ent
    let expr_base = "ExprTyp (" ++ monad_base ++ ")"
        array_base = "ArrayTyp (" ++ monad_base ++ ")"
        alloc_base = "AllocTyp (" ++ monad_base ++ ")"
    monad_conds <- optlistent False "monad_conds" paren_list
    array_monad_cond <- optlistent False "array_monad_cond" no_comma_ent
    monad_newtyp <- optlistent False "monad_newtyp" no_comma_ent
    expr_newtyp <- optlistent True "expr_newtyp" no_comma_ent
    return $
        (mk_monad_trans name monad_base monad_conds monad_newtyp) ++
        (mk_e_trans name expr_base
            [ ("Suitable (" ++ expr_base ++ ")") ]
            monad_conds expr_newtyp) ++
        (mk_array_trans name array_base alloc_base
            [ ("Suitable (" ++ array_base ++ ")") ]
            (array_monad_cond : monad_conds) (expr_newtyp ++ "A"))

parseMacro = try suitable_macro <|>
    expr_trans_macro <|>
    array_trans_macro <|>
    monad_trans_macro <|>
    hz_trans_macro

parseMacro' = many space *> parseMacro <* (many space >> string "}}")

parseAll = non_lbrace <++>
    ( ("" <$ eof) <|>
      (try (string "{" <++> non_lbrace1) <++> parseAll) <|>
      (string "{{" >> parseMacro' <++> parseAll))
    where non_lbrace = many (noneOf "{")
          non_lbrace1 = many1 (noneOf "{")

preprocessor :: String -> String -> IO ()
preprocessor n = handleErr . parse (parseAll <* eof) n
    where handleErr (Right x) = putStrLn x
          handleErr (Left y) = do
              hPutStrLn stderr "parse error!"
              hPutStrLn stderr (show y)
              exitWith (ExitFailure 1)

main = do
    (n:_) <- getArgs
    c <- getContents
    preprocessor n c