{-# OPTIONS -Wall #-}
{-# LANGUAGE DeriveDataTypeable #-}

-- The grm grammar generator
-- Copyright 2011-2012, Brett Letner

module Main where

import Control.Monad
import Data.Char
import Data.List
import Data.Maybe
import Distribution.Text
import Grm.Prims
import Paths_grm
import System.Console.CmdArgs
import System.Directory
import System.FilePath.Posix
import Text.ParserCombinators.Parsec
import Text.ParserCombinators.Parsec.Language
import Text.PrettyPrint.Leijen (braces,brackets,text,hsep,(<+>),empty,vcat,Doc,indent,(<>),parens,int,list)
import qualified Text.ParserCombinators.Parsec.Token as P

type P a = Parser a

data Args = Args
  { files :: [FilePath]
  , locations :: Bool
  } deriving (Show, Data, Typeable)

argsDesc :: Args
argsDesc = Args
  { files = def &= args
  , locations = def &= help "generate location information"
  } &= summary summ &= program prog
  where
  summ = prog ++ " v" ++ display version ++ ", " ++ copyright
  prog = "grm"
  copyright = "(C) Brett Letner 2011-2012"

main :: IO ()
main = do
  a <- cmdArgs argsDesc
  mapM_ (doit (locations a)) $ files a

doit :: Bool -> String -> IO ()
doit gen_locs fn = do
  eea <- parseFromFile grm fn
  case eea of
    Left e -> error $ show e
    Right xs -> do
      let bn = takeBaseName fn
      let dir = combine "Language" bn
      createDirectoryIfMissing True dir
      absGrm gen_locs dir bn xs
      parGrm gen_locs dir bn xs

absGrm :: Bool -> FilePath -> String -> [Decl] -> IO ()
absGrm gen_locs dir bn xs = writeFile (combine dir "Abs.hs") $ show d
  where
  d = vcat $
    [ text "{-# LANGUAGE DeriveDataTypeable #-}"
    , hsep [ text "module"
           , text $ "Language." ++ bn ++ ".Abs"
           , text "where"]
    , vcat [ text "import" <+> text imp | imp <- imps ]
    , lexDecls xs
    ] ++
    map (dataDecl gen_locs) xs ++ 
    (if gen_locs then map metaDecl xs else []) ++
    map (ppDecl gen_locs) xs
  imps =
    [ "Control.DeepSeq"
    , "Data.Generics"
    , "Grm.Prims"
    , "Grm.Lex"
    , "Text.PrettyPrint.Leijen"
    ]

lexDecls :: [Decl] -> Doc
lexDecls xs = vcat
  [ hsep [text "myLexemes =", text $ show $ resDecls xs]
  , text "grmLexFilePath = lexFilePath myLexemes"
  , text "grmLexContents = lexContents myLexemes"
  ]

absName :: FilePath -> String
absName fn = takeBaseName fn

hdr :: String
hdr = unlines
  [ "{-# OPTIONS -w #-}"
  , "-- Haskell module generated by grm *** DO NOT EDIT ***"
  ]

parGrm :: Bool -> FilePath -> String -> [Decl] -> IO ()
parGrm gen_locs dir bn xs = do
  let fn1 = dir </> "Par.y"
  writeFile fn1 $ show d
  where
  d = vcat $
    [ parHdr
    ] ++
    map (parDecl gen_locs [ s | List s _ _ _ _ _ <- xs ]) xs
  parHdr = vcat $ map text
    [ "{"
    , hdr
    , ""
    , "module Language." ++ bn ++ ".Par where"
    , "import Grm.Prims"
    , "import Grm.Lex"
    , "import Language." ++ bn ++ ".Abs"
    , "}"
    , "%tokentype { (Token Point) }"
    , "%name grmParse"
    , "%token"
    , unlines [ "  " ++ show sym ++ " " ++ "{ TSymbol _ " ++ show sym ++ " }"
                | sym <- resDecls xs ]
    , "  uident { TUident _ _ }"
    , "  usym { TUsym _ _ }"
    , "  lident { TLident _ _ }"
    , "  string { TString _ _ }"
    , "  char { TChar _ _ }"
    , "  number { TNumber _ _ }"
    , "%%"
    ]

capitalize :: String -> String
capitalize "" = ""
capitalize (c:cs) = toUpper c : cs

grm :: P [Decl]
grm = do
  whiteSpace
  ds <- many decl
  eof
  return ds

pGroup :: P Decl
pGroup = do
  reserved "group"
  liftM Group $ pBraces $ many1 pData

pData :: P DataD
pData = do
  reserved "data"
  n <- identifier
  xs <- many1 alt
  return $ DataD n xs

pList :: P Decl
pList = do
  reserved "list"
  liftM6 List identifier identifier pEmpty terminator stringLiteral horiz

decl :: P Decl
decl = choice [pGroup, liftM (Group . singleton) pData, pType, pList]

vcatStr :: String -> [Doc] -> Doc
vcatStr s xs = vcat $ intersperse (text s) xs

altList :: [Doc] -> Doc
altList (x0:xs) = vcat $
  [text "=" <+> x0] ++ [ text "|" <+> x | x <- xs ] ++
  [text "deriving (Show,Eq,Ord,Data,Typeable)"]
altList [] = unreachable

dataDecl :: Bool -> Decl -> Doc
dataDecl gen_locs x = case x of
  Group [] -> unreachable
  Group xs@(DataD n0 _ : _) -> vcat $
    [text "data" <+> text n0 <+> locs, indent 2 $ altList $ concatMap (dataData gen_locs) xs]
    ++ types (dataDecl gen_locs) xs
  Type a b -> dataType gen_locs a $ text (capitalize b) <+> locs
  List a b _ _ _ _ -> dataType gen_locs a $ brackets $ text b <+> locs
  where
    locs = if gen_locs then text "a" else empty
    
types :: (Decl -> Doc) -> [DataD] -> [Doc]
types f (DataD b _ : xs) = [ f $ Type a b | DataD a _ <- xs ]
types _ [] = unreachable

dataType :: Bool -> String -> Doc -> Doc
dataType gen_locs a b = hsep [text "type", text a, locs, text "=", b]
  where locs = if gen_locs then text "a" else empty
        
dataData :: Bool -> DataD -> [Doc]
dataData gen_locs (DataD _ xs) =
  [ hsep $ map text $ s : locs : (catMaybes $ map (nameAltTok gen_locs) ts)
    | Alt s ts <- xs ]
  where locs = if gen_locs then "a" else ""
        
nameAltTok :: Bool -> AltTok -> Maybe String
nameAltTok gen_locs x = case x of
  StringT{} -> Nothing
  IdentT s -> Just $ if gen_locs then "(" ++ s ++ " a)" else s
  PrimT s -> Just $ capitalize s

metaDecl :: Decl -> Doc
metaDecl x = case x of
  Group [] -> unreachable
  Group xs@(DataD n0 _ : _) -> vcat $
    [ hsep [text "instance HasMeta", text n0, text "where"]
    , indent 2 $ text $ "meta x = case x of"
    , indent 4 $ vcat $ concatMap metaData xs
    ]
  Type{} -> empty
  List{} -> empty

metaData :: DataD -> [Doc]
metaData (DataD _ xs) = [ f s ts | Alt s ts <- xs ]
  where
  f s ts =
    hsep [text s, text "a", hsep $ replicate (length ys) (text "_"), text "-> a" ]
    where
    ys = filter (not . isStringT) ts

ppDecl :: Bool -> Decl -> Doc
ppDecl gen_locs x = case x of
  Group [] -> unreachable
  Group xs@(DataD n0 _ : _) -> vcat $
    [ hsep [ text "instance Pretty"
           , parens (text n0 <+> if gen_locs then text "a" else empty)
           , text "where" ]
    , indent 2 $ text "pretty = pp" <> text n0
    , text $ "pp" ++ n0 ++ " x = case x of"
    , indent 2 $ vcat $ concatMap (ppData gen_locs) xs
    ] ++ types (ppDecl gen_locs) xs
  Type a b -> hsep $ map text ["pp" ++ a, "=", "pp" ++ capitalize b]
  List a b _ d e f ->
    hsep $ map text ["pp" ++ a, "=", "ppList", "pp" ++ b, show d, show e, show f]

ppData :: Bool -> DataD -> [Doc]
ppData gen_locs (DataD _ xs) = [ f s ts | Alt s ts <- xs ]
  where
  f s ts =
    hsep [ text s
         , if gen_locs then text "_" else empty
         , hsep $ map text [ v | (Just v, _) <- ys ], text "->", ppAltToks ys ]
    where
    ys = nameToks "v" $ numberToks (not . isStringT) ts

isStringT :: AltTok -> Bool
isStringT (StringT{}) = True
isStringT _ = False

nameToks :: String -> [(Maybe Int, AltTok)] -> [(Maybe String, AltTok)]
nameToks s xs = [ (fmap (\i -> s ++ show i) mi, t) | (mi,t) <- xs ]

resDecls :: [Decl] -> [String]
resDecls xs = filter (not . null) $ sort $ nub $ concatMap resDecl xs

numberToks :: (AltTok -> Bool) -> [AltTok] -> [(Maybe Int, AltTok)]
numberToks f = loop 1
  where
  loop _ [] = []
  loop i (x:xs)
    | f x = (Just i, x) : loop (succ i) xs
    | otherwise = (Nothing, x) : loop i xs

ppAltToks :: [(Maybe String, AltTok)] -> Doc
ppAltToks [] = text "Text.PrettyPrint.Leijen.empty"
ppAltToks [x] = ppAltTok x
ppAltToks (x:y:xs) = ppAltTok x <+> rest
  where
  rest = case y of
    (_, StringT "") -> text "<>" <+> ppAltToks xs
    _ -> text "<+>" <+> ppAltToks (y:xs)

ppAltTok :: (Maybe String, AltTok) -> Doc
ppAltTok x = case x of
  (_, StringT "") -> error "unexpected empty string"
  (_, StringT s) -> hsep [text "text", text $ show s]
  (Just v, IdentT s) -> text $ "pp" ++ s ++ " " ++ v
  (mv, PrimT s) -> ppAltTok (mv, IdentT $ capitalize s)
  _ -> unreachable

pEmpty :: P Empty
pEmpty = choice
  [ reserved "empty" >> return Empty
  , reserved "nonempty" >> return NonEmpty
  ]

terminator :: P Terminator
terminator = choice
  [ reserved "separator" >> return Separator
  , reserved "terminator" >> return Terminator
  ]

horiz :: P Horiz
horiz = choice
  [ reserved "vert" >> return Vert
  , reserved "horiz" >> return Horiz
  ]

alt :: P Alt
alt = do
  reservedOp "|"
  choice [pAlt,defAlt]

pAlt :: P Alt
pAlt = do
  c <- identifier
  ts <- many altTok
  return $ Alt c ts

defAlt :: P Alt
defAlt = do
  reserved "_"
  liftM DefAlt identifier

pType :: P Decl
pType = do
  reserved "type"
  c <- identifier
  reservedOp "="
  p <- identifier
  return $ Type c p

altTok :: P AltTok
altTok = choice
  [ liftM StringT stringLiteral
  , liftM IdentT identifier
  , liftM PrimT prim
  ]

prim :: P String
prim = choice $ map res primNames

res :: String -> P String
res s = reserved s >> return s

primNames :: [String]
primNames = ["string","number","char","uident","lident","usym"]

data Decl
  = Group [DataD]
  | Type String String
  | List String String Empty Terminator String Horiz
  deriving Show

resDecl :: Decl -> [String]
resDecl x = case x of
  Group xs -> concatMap resData xs
  Type{} -> []
  List _ _ _ _ s _ -> [s]

resData :: DataD -> [String]
resData (DataD _ xs) = concatMap resAlt xs

resAlt :: Alt -> [String]
resAlt x = case x of
  Alt _ ys -> concatMap resAltTok ys
  DefAlt{} -> []

data DataD = DataD String [Alt] deriving Show

data Alt
  = Alt String [AltTok]
  | DefAlt String
  deriving Show

resAltTok :: AltTok -> [String]
resAltTok x = case x of
  StringT s -> [s]
  _ -> []

data AltTok
  = StringT String
  | IdentT String
  | PrimT String
  deriving (Show, Eq)

parData :: Bool -> [String] -> DataD -> Doc
parData gen_locs ss (DataD c xs0) = case filter f xs0 of
  [] -> unreachable
  x0:xs -> vcat
    [ text c
    , indent 2 $ vcat
        [ hsep [ text ":", parAlt gen_locs ss x0 ]
        , vcat [ text "|" <+> parAlt gen_locs ss x | x <- xs ]
        ]
    ]
  where
  f x = case x of
    Alt a _ -> last a /= '_'
    DefAlt{} -> True

parAlt :: Bool -> [String] -> Alt -> Doc
parAlt gen_locs ss x = case x of
  Alt c ts -> hsep $
    map parAltTokL ts ++
    [braces $ hsep $ text c : (if gen_locs then loc else empty) : map parAltTokR ns ]
    where
    ns = nameToks "$" $ numberToks ((/=) (StringT "")) ts
    loc = case filter ((/=) (StringT "")) ts of
      [] -> text "noPoint"
      [_] -> parens (text "point $1")
      zs -> parens $ text "lrPoint" <+> list (map (pointAltTok ss) (zip zs [ 1 .. ]))
  DefAlt c -> hsep [text c, text "{ $1 }"]

pointAltTok :: [String] -> (AltTok,Int) -> Doc
pointAltTok ss (x,i) = case x of
  IdentT s | s `elem` ss -> text "lrPointList" <+> text "$" <> int i
  _ -> text "point" <+> text "$" <> int i

parAltTokR :: (Maybe String, AltTok) -> Doc
parAltTokR x = case x of
  (Just i, IdentT _) -> text i
  (Just i, PrimT s) -> parens (text ("unT" ++ capitalize s) <+> text i)
  _ -> empty

parAltTokL :: AltTok -> Doc
parAltTokL x = case x of
  StringT "" -> empty
  StringT s -> text $ show s
  IdentT s -> text s
  PrimT s -> text s

parDecl :: Bool -> [String] -> Decl -> Doc
parDecl gen_locs ss x = case x of
  Type a b -> hsep [text a, text ":", text b, braces $ text "$1"]
  Group xs -> vcat $ map (parData gen_locs ss) xs
  List prodsStr prodStr empt term sepStr _ -> vcat
    [ prods
    , indent 2 $ vcat
      [ hsep [text ":", rev_prods, braces $ text "reverse $1"]
      , case empt of
          Empty -> text "| {- empty -} { [] }"
          NonEmpty -> empty
      ]
    , rev_prods
    , indent 2 $ vcat
      [ hsep [text ":", a, braces $ brackets $ text "$1"]
      , hsep [text "|", rev_prods, b, braces $ hsep [c, text ": $1"]]
      ]
    ]
    where
    a = case term of
      Terminator | sepStr /= "" -> prod <+> sp
      _ -> prod
    b = case (sepStr, term) of
      ("", _) -> prod
      (_, Terminator) -> prod <+> sp
      (_, Separator) -> sp <+> prod
    c = case term of
      Separator | sepStr /= "" -> text "$3"
      _ -> text "$2"
    sp = text $ show sepStr
    rev_prods = text $ "REV_" ++ prodsStr
    prod = text prodStr
    prods = text prodsStr


lexer :: P.TokenParser a
lexer = P.makeTokenParser grmStyle

reserved :: String -> P ()
reserved = P.reserved lexer

pBraces :: P a -> P a
pBraces = P.braces lexer

identifier :: P String
identifier = P.identifier lexer

reservedOp :: String -> P ()
reservedOp = P.reservedOp lexer

stringLiteral :: P String
stringLiteral = P.stringLiteral lexer

whiteSpace :: P ()
whiteSpace = P.whiteSpace lexer

grmStyle :: LanguageDef a
grmStyle = haskellStyle
  { P.identStart = upper
  }

liftM6 :: Monad m =>
  (t -> t1 -> t2 -> t3 -> t4 -> t5 -> b) ->
  m t -> m t1 -> m t2 -> m t3 -> m t4 -> m t5 -> m b
liftM6 z ma mb mc md me mf = do
  a <- ma
  b <- mb
  c <- mc
  d <- md
  e <- me
  f <- mf
  return $ z a b c d e f