{-# LANGUAGE TemplateHaskell #-}

module Wayland.CodeGen.CodeGen
  ( GeneratorConfig(..)
  , generateInterface
  , generateInterfaceBoot
  ) where

import Wayland.CodeGen.Types
import Wayland.CodeGen.Utils
import Control.Monad.State.Strict (State, runState)
import Lens.Micro.TH (makeLenses)
import Data.List (intercalate, sort, isPrefixOf)
import Lens.Micro.Mtl ((%=))
import Data.Functor (($>))
import Control.Monad.Reader (ReaderT (..), asks)
import Data.Map ((!))
import qualified Data.Map as M

data DocComment
  = DocCommentBefore
  | DocCommentAfter
  deriving (Show, Eq)

newtype Declaration = Declaration { declLines :: [String] }
data Import
  = Import { imModule :: !String
           , imSource :: !Bool
           }
  deriving (Show, Eq)

data GeneratorState
  = GeneratorState { _stateImports       :: ![Import]
                   , _stateExports       :: ![String]
                   , _stateExtensions    :: ![String]
                   , _stateDecls         :: ![Declaration]
                   }

data GeneratorConfig
  = GeneratorConfig { genModules   :: !FlatModules
                    , genProtocol  :: !TypedProtocol
                    , genInterface :: !InterfaceDesc
                    , genTypes     :: !TypeInfo
                    }

type Generator = ReaderT GeneratorConfig (State GeneratorState)

makeLenses ''GeneratorState

generateInterface :: GeneratorConfig -> String
generateInterface = generate emit

generateInterfaceBoot :: GeneratorConfig -> String
generateInterfaceBoot = generate emitBoot

generate :: Generator () -> GeneratorConfig -> String
generate g cfg = transform $ runState (runReaderT g cfg) initialState
  where initialState = GeneratorState { _stateImports       = []
                                      , _stateExports       = []
                                      , _stateExtensions    = []
                                      , _stateDecls         = []
                                      }
        transform (_, state) = undecls $ [ extensionDecl state, moduleDecl state, importDecl state ] ++ _stateDecls state
        extensionDecl s = Declaration $ sort $ fmap formatExtension (_stateExtensions s)
        moduleDecl s = Declaration [ "module " ++ genFullModule cfg ++ " (" ++ intercalate ", " (_stateExports s) ++ ") where" ]
        importDecl s = Declaration $ sort $ fmap formatImport (_stateImports s)
        undecls decls = intercalate "\n" $ unlines . declLines <$> filter (not . null . declLines) decls
        formatExtension e = "{-# LANGUAGE " ++ e ++ " #-}"
        formatImport i = "import " ++ (if imSource i then "{-# SOURCE #-} " else "") ++ "qualified " ++ imModule i

genFullModule :: GeneratorConfig -> String
genFullModule cfg = fullModule (genProtocol cfg) (genInterface cfg)

emitDecl :: [String] -> Generator ()
emitDecl c = stateDecls %= (++ [ Declaration c ])

emitImport :: Import -> Generator ()
emitImport i = stateImports %= \is -> if i `elem` is then is else i:is

emitExport :: String -> Generator ()
emitExport e = stateExports %= \es -> es ++ [e]

emitExtension :: String -> Generator ()
emitExtension e = stateExtensions %= (e:)

resolve :: String -> String -> Generator String
resolve m e = do
  this <- asks genFullModule
  if   this == m
  then pure e
  else emitImport (Import m ("Wayland.Protocols." `isPrefixOf` m)) $> m ++ "." ++ e

resolveInterface :: Maybe String -> String -> Generator String
resolveInterface (Just name) e = do modules <- asks genModules
                                    let (p, i) = modules ! name
                                    resolve (fullModule p (fst i)) e
resolveInterface Nothing     e = pure e

wrapEnumArray :: String -> EnumDesc -> String
wrapEnumArray name e = if enumBitfield e then "[" ++ name ++ "]" else name

resolveEnum :: String -> Generator String
resolveEnum e = resolveQualified $ break (== '.') e
  where resolveQualified :: (String, String) -> Generator String
        resolveQualified (e', [])   = do types <- asks genTypes
                                         let EnumName desc name = typeEnums types ! e'
                                         pure $ wrapEnumArray name desc
        resolveQualified (i,  _:e') = do modules <- asks genModules
                                         let iface = snd (modules ! i)
                                         let EnumName desc name = typeEnums (snd iface) ! e'
                                         resolved <- resolveInterface (Just i) name
                                         pure $ wrapEnumArray resolved desc

resolveArgument :: ArgumentDesc -> Generator String
resolveArgument a = case argType a of
                      ArgumentInt    -> maybe (resolve "Foreign" "Int32")  resolveEnum (argEnum a)
                      ArgumentUInt   -> maybe (resolve "Foreign" "Word32") resolveEnum (argEnum a)
                      ArgumentFixed  -> pure "Float"
                      ArgumentArray  -> resolve "Data.ByteString" "ByteString"
                      ArgumentFD     -> resolve "System.Posix.Types" "Fd"
                      ArgumentString -> pure $ if argAllowNull a then "(Maybe String)" else "String"
                      ArgumentNewID  -> do obj <- resolveInterface (argInterface a) "Object"
                                           newObj <- resolve "Wayland.Server" "NewObject"
                                           pure $ "(" ++ newObj ++ " " ++ obj ++ ")"
                      ArgumentObject -> do obj <- resolveInterface (argInterface a) "Object"
                                           pure $ if argAllowNull a then "(Maybe " ++ obj ++ ")" else obj

emitBoot :: Generator ()
emitBoot = do
  anyObject         <- resolve "Wayland.Server" "AnyObject"
  interfaceDesc     <- resolve "Wayland.Server" "InterfaceDesc"
  objectClass       <- resolve "Wayland.Server" "Object"
  marshallableClass <- resolve "Wayland.Server" "Marshallable"
  enumTypes         <- asks $ typeEnums . genTypes
  i                 <- asks genInterface
  emitExtension "FlexibleInstances"
  emitDecl [ "newtype Object = Object " ++ anyObject ]
  emitExport "Object"
  emitDecl [ "interface :: " ++ interfaceDesc ]
  emitExport "interface"
  emitDecl [ "instance " ++ objectClass ++ " Object"
           , "instance Show Object"
           , "instance Eq Object"
           , "instance " ++ marshallableClass ++ " Object"
           , "instance " ++ marshallableClass ++ " (Maybe Object)" ]
  let emitEnumBoot e = do let EnumName _ name = enumTypes ! enumName e
                          emitDecl [ "data " ++ name
                                   , "instance Show " ++ name
                                   , "instance Eq " ++ name
                                   , "instance " ++ marshallableClass ++ " " ++ (if enumBitfield e then "[" ++ name ++ "]" else name) ]
                          emitExport name
  mapM_ emitEnumBoot (ifEnums i)

emit :: Generator ()
emit = do
  anyObject     <- resolve "Wayland.Server" "AnyObject"
  hashable      <- resolve "Data.Hashable" "Hashable"
  objectClass   <- resolve "Wayland.Server" "Object"
  throwError'   <- resolve "Control.Monad.Except" "throwError"
  i             <- asks genInterface
  types         <- asks genTypes
  emitExtension "GeneralizedNewtypeDeriving"
  emitExtension "EmptyDataDeriving"
  emitExtension "FlexibleInstances"
  emitDecl (formatDescription "" DocCommentBefore (ifDescription i) ++
            [ "newtype Object = Object " ++ anyObject
            , "  deriving (Show, Eq, Ord, " ++ hashable ++ ")" ])
  emitExport "Object"
  mapM_ (emitMessageDesc requestDescName) (ifRequests i)
  mapM_ (emitMessageDesc eventDescName)   (ifEvents i)
  emitInterfaceDesc
  emitRequests (ifRequests i)
  emitEvents   (ifEvents   i)
  mapM_ emitEnum (ifEnums i)
  mapM_ emitArguments (typeRequestArgs types)
  mapM_ emitArguments (typeEventArgs   types)
  emitExtension "TypeFamilies"
  let matchArgument m msgTypes = case msgTypes ! msgName m of
                                   ArgumentNameNone       -> ""
                                   ArgumentNameSingle _   -> " _"
                                   ArgumentNameRecord _ _ -> " _"
  let requestSignature m =
        "  requestSignature _ (Req" ++ pascalCase (msgName m) ++ " _" ++ matchArgument m (typeRequestArgs types) ++ ") = req" ++ pascalCase (msgName m)
  let eventSignature m =
        "  eventSignature _ (Ev" ++ pascalCase (msgName m) ++ matchArgument m (typeEventArgs types) ++ ") = ev" ++ pascalCase (msgName m)
  demarshalRequests <- concat <$> mapM demarshalRequest (ifRequests i)
  marshalEvents     <- concat <$> mapM marshalEvent     (ifEvents   i)
  emitDecl ([ "instance " ++ objectClass ++ " Object where"
            , "  type ObjectEvent Object = Event"
            , "  type ObjectRequest Object = Request"
            , "  toAny (Object o) = o"
            , "  fromAny = Object"
            , "  describeInterface _ = interface" ]
            ++ demarshalRequests ++
            [ "  demarshalRequest o c = " ++ throwError' ++ " $ \"Unknown opcode \" ++ show c ++ \" on \" ++ show o" ]
            ++ marshalEvents ++
            [ "  marshalEvent _ _ = undefined" ]
            ++ fmap requestSignature (ifRequests i) ++
            [ "  requestSignature _ _ = undefined" ]
            ++ fmap eventSignature (ifEvents i) ++
            [ "  eventSignature _ _ = undefined" ])
  marshallableClass <- resolve "Wayland.Server" "Marshallable"
  marshal'          <- resolve "Wayland.Server" "marshal"
  demarshal'        <- resolve "Wayland.Server" "demarshal"
  emitDecl [ "instance " ++ marshallableClass ++ " Object where"
           , "  marshal (Object o) = " ++ marshal' ++ " o"
           , "  demarshal = Object <$> " ++ demarshal' ]
  emitDecl [ "instance " ++ marshallableClass ++ " (Maybe Object) where"
           , "  marshal (Just (Object o)) = " ++ marshal' ++ " (Just o)"
           , "  marshal Nothing           = " ++ marshal' ++ " (Nothing :: Maybe " ++ anyObject ++ ")"
           , "  demarshal = fmap (Object <$>) " ++ demarshal' ]

requestDescName :: MessageDesc -> String
requestDescName m = "req" ++ pascalCase (msgName m)

eventDescName :: MessageDesc -> String
eventDescName m = "ev" ++ pascalCase (msgName m)

emitMessageDesc :: (MessageDesc -> String) -> MessageDesc -> Generator ()
emitMessageDesc name m = do
  let argumentNullability True  = resolve "Wayland.Server" "Nullable"
      argumentNullability False = resolve "Wayland.Server" "NotNull"
  let argumentType :: ArgumentDesc -> Generator String
      argumentType a = case argType a of
                         ArgumentInt    -> resolve "Wayland.Server" "ArgumentInt"
                         ArgumentUInt   -> resolve "Wayland.Server" "ArgumentUInt"
                         ArgumentFixed  -> resolve "Wayland.Server" "ArgumentFixed"
                         ArgumentArray  -> resolve "Wayland.Server" "ArgumentArray"
                         ArgumentFD     -> resolve "Wayland.Server" "ArgumentFD"
                         ArgumentString -> do aType        <- resolve "Wayland.Server" "ArgumentString"
                                              aNullability <- argumentNullability (argAllowNull a)
                                              pure $ aType ++ " " ++ aNullability
                         ArgumentObject -> do aType        <- resolve "Wayland.Server" "ArgumentObject"
                                              aIface       <- resolveInterface (argInterface a) "interface"
                                              aNullability <- argumentNullability (argAllowNull a)
                                              pure $ aType ++ " " ++ aNullability ++ " " ++ aIface
                         ArgumentNewID  -> do aType        <- resolve "Wayland.Server" "ArgumentNewID"
                                              aIface       <- resolveInterface (argInterface a) "interface"
                                              pure $ aType ++ " " ++ aIface
  messageDesc      <- resolve "Wayland.Server" "MessageDesc"
  msgName'         <- resolve "Wayland.Server" "msgName"
  msgArguments'    <- resolve "Wayland.Server" "msgArguments"
  msgSince'        <- resolve "Wayland.Server" "msgSince"
  msgIsDestructor' <- resolve "Wayland.Server" "msgIsDestructor"
  arguments        <- mapM argumentType (msgArguments m)
  emitDecl [ name m ++ " :: " ++ messageDesc
           , name m ++ " ="
           , "  " ++ messageDesc
           , "    { " ++ msgName' ++ " = " ++ escapeStr (msgName m)
           , "    , " ++ msgArguments' ++ " = [" ++ intercalate ", " arguments ++ "]"
           , "    , " ++ msgSince' ++ " = " ++ show (msgSince m)
           , "    , " ++ msgIsDestructor' ++ " = " ++ show (msgType m == MessageDestructor)
           , "    }"
           ]

emitInterfaceDesc :: Generator ()
emitInterfaceDesc = do
  interfaceDesc <- resolve "Wayland.Server" "InterfaceDesc"
  ifName'       <- resolve "Wayland.Server" "ifName"
  ifVersion'    <- resolve "Wayland.Server" "ifVersion"
  ifRequests'   <- resolve "Wayland.Server" "ifRequests"
  ifEvents'     <- resolve "Wayland.Server" "ifEvents"
  i             <- asks genInterface
  emitDecl [ "-- | Description of the interface."
           , "interface :: " ++ interfaceDesc
           , "interface ="
           , "  " ++ interfaceDesc
           , "    { " ++ ifName' ++ " = " ++ escapeStr (ifName i)
           , "    , " ++ ifVersion' ++ " = " ++ show (ifVersion i)
           , "    , " ++ ifRequests' ++ " = [" ++ intercalate ", " (requestDescName <$> ifRequests i)  ++ "]"
           , "    , " ++ ifEvents' ++ " = [" ++ intercalate ", " (eventDescName <$> ifEvents i)  ++ "]"
           , "    }" ]
  emitExport "interface"

requestConstructor :: MessageDesc -> String
requestConstructor m = "Req" ++ pascalCase (msgName m)

eventConstructor :: MessageDesc -> String
eventConstructor m = "Ev" ++ pascalCase (msgName m)

emitRequests :: [MessageDesc] -> Generator ()
emitRequests ms = emitMessages ms typeRequestArgs formatConstructor "Request" "Requests supported by the interface."
  where formatConstructor m = requestConstructor m ++ " Object"

emitEvents :: [MessageDesc] -> Generator ()
emitEvents ms = emitMessages ms typeEventArgs eventConstructor "Event" "Events supported by the interface."

emitMessages :: [MessageDesc] -> (TypeInfo -> M.Map String ArgumentName) -> (MessageDesc -> String) -> String -> String -> Generator ()
emitMessages ms argTypes formatConstructor type' comment = do
  typeable     <- resolve "Data.Typeable" "Typeable"
  messageClass <- resolve "Wayland.Server" "MessageClass"
  typeClass    <- resolve "Wayland.Server" (type' ++ "Class")
  args         <- asks $ argTypes . genTypes
  let constructorArg :: ArgumentName -> Generator String
      constructorArg ArgumentNameNone         = pure ""
      constructorArg (ArgumentNameSingle a)   = (++) " " <$> resolveArgument a
      constructorArg (ArgumentNameRecord _ r) = pure $ " " ++ r
  let constructor m = do arg <- constructorArg (args ! msgName m)
                         pure $ (formatConstructor m ++ arg) : formatDescription "    " DocCommentAfter (msgDescription m)
  constructors <- mapM constructor ms
  emitDecl $ [ "-- | " ++ comment
             , "data " ++ type' ]
           ++ dataConstructors constructors
           ++ [ "  deriving (" ++ typeable ++ ", Show, Eq)" ]
  emitExport $ type' ++ "(..)"
  emitDecl [ "instance " ++ messageClass ++ " " ++ type'
           , "instance " ++ typeClass ++ " " ++ type' ]

demarshalRequest :: MessageDesc -> Generator [String]
demarshalRequest m = do
  args <- asks $ typeRequestArgs . genTypes
  demarshal' <- resolve "Wayland.Server" "demarshal"
  case args ! msgName m of
    ArgumentNameNone        -> pure [ "  demarshalRequest o " ++ show (msgOpCode m) ++ " = pure $ " ++ requestConstructor m ++ " o" ]
    ArgumentNameSingle _    -> do pure [ "  demarshalRequest o " ++ show (msgOpCode m) ++ " = " ++ requestConstructor m ++ " o <$> " ++ demarshal' ]
    ArgumentNameRecord as r -> do let apply a = "    " ++ camelCase (argName a) ++ "' <- " ++ demarshal'
                                  let field a = recordField r (camelCase (argName a)) ++ " = " ++ camelCase (argName a) ++ "'"
                                  pure ([ "  demarshalRequest o " ++ show (msgOpCode m) ++ " = do" ]
                                        ++ fmap apply as ++
                                        [ "    pure $ " ++ requestConstructor m ++ " o $ " ++ r ]
                                        ++ formatBrackets "      {" "      ," "      }" (fmap (pure . field) as))

marshalEvent :: MessageDesc -> Generator [String]
marshalEvent m = do
  args <- asks $ typeEventArgs . genTypes
  marshal' <- resolve "Wayland.Server" "marshal"
  case args ! msgName m of
    ArgumentNameNone        -> pure [ "  marshalEvent _ " ++ eventConstructor m ++ " = pure " ++ show (msgOpCode m) ]
    ArgumentNameSingle _    -> do pure [ "  marshalEvent _ (" ++ eventConstructor m ++ " x) = do"
                                       , "    " ++ marshal' ++ " x"
                                       , "    pure " ++ show (msgOpCode m) ]
    ArgumentNameRecord as r -> do let apply a = "    " ++ marshal' ++ " $ " ++ recordField r (camelCase (argName a)) ++ " x"
                                  pure ([ "  marshalEvent _ (" ++ eventConstructor m ++ " x) = do" ]
                                        ++ fmap apply as ++
                                        [ "    pure " ++ show (msgOpCode m) ])

emitEnum :: EnumDesc -> Generator ()
emitEnum e = do
  types <- asks $ typeEnums . genTypes
  let EnumName _ name = types ! enumName e
  let constructorName entry = name ++ pascalCase (entryName entry)
  let constructorNames = fmap constructorName (enumEntries e)
  let constructorDecl (cName, entry) = cName : formatDescription "    " DocCommentAfter (entryDescription entry)
  let constructors = dataConstructors (fmap constructorDecl (zip constructorNames (enumEntries e)))
  emitDecl (formatDescription "" DocCommentBefore (enumDescription e) ++
            [ "data " ++ name ]
            ++ constructors ++
            [ "  deriving (Show, Eq)" ])
  emitExport $ name ++ "(..)"
  let values = zip constructorNames (fmap entryValue (enumEntries e))
  (if enumBitfield e then emitMarshallableBitfield else emitMarshallableEnum) name values

emitMarshallableEnum :: String -> [(String, Int)] -> Generator ()
emitMarshallableEnum name values = do
  marshallableClass <- resolve "Wayland.Server" "Marshallable"
  marshalUInt'      <- resolve "Wayland.Server" "marshalUInt"
  demarshalUInt'    <- resolve "Wayland.Server" "demarshalUInt"
  throwError'       <- resolve "Control.Monad.Except" "throwError"
  let marshal (c, i) = "  marshal " ++ c ++ " = " ++ marshalUInt' ++ " " ++ show i
  let demarshalCase (c, i) = "      " ++ show i ++ " -> pure " ++ c
  emitDecl ([ "instance " ++ marshallableClass ++ " " ++ name ++ " where" ]
            ++ fmap marshal values ++
            [ "  marshal _ = undefined"
            , "  demarshal = do"
            , "    x <- " ++ demarshalUInt'
            , "    case x of" ]
            ++ fmap demarshalCase values ++
            [ "      _ -> " ++ throwError' ++ " $ \"Unknown enum value for type " ++ name ++ ": \" ++ show x" ])

emitMarshallableBitfield :: String -> [(String, Int)] -> Generator ()
emitMarshallableBitfield name values = do
  marshallableClass <- resolve "Wayland.Server" "Marshallable"
  marshalUInt'      <- resolve "Wayland.Server" "marshalUInt"
  demarshalUInt'    <- resolve "Wayland.Server" "demarshalUInt"
  catMaybes'        <- resolve "Data.Maybe" "catMaybes"
  bitOr             <- resolve "Foreign" ".|."
  bitAnd            <- resolve "Foreign" ".&."
  let marshalVal (c, i) = "      val " ++ c ++ " = " ++ show i
  let demarshalVal (c, i) = "val " ++ show i ++ " " ++ c
  emitDecl ([ "instance " ++ marshallableClass ++ " [" ++ name ++ "] where"
            , "  marshal = " ++ marshalUInt' ++ " . foldr ((" ++ bitOr ++ ") . val) 0"
            , "    where" ]
            ++ fmap marshalVal values ++
            [ "      val _ = undefined"
            , "  demarshal = do"
            , "    x <- " ++ demarshalUInt'
            , "    let val v v' = if x " ++ bitAnd ++ " v == v then Just v' else Nothing"
            , "    pure $ " ++ catMaybes' ]
            ++ formatBrackets "      [" "      ," "      ]" (fmap (pure . demarshalVal) values))

emitArguments :: ArgumentName -> Generator ()
emitArguments  ArgumentNameNone              = pure ()
emitArguments (ArgumentNameSingle _)         = pure ()
emitArguments (ArgumentNameRecord args name) = do
  let field (a, aName) = (recordField name (camelCase (argName a)) ++ " :: " ++ aName)
                       : formatDescription "    " DocCommentAfter (argDescription a)
  resolvedArgs <- mapM resolveArgument args
  emitDecl ([ "data " ++ name
            , "  = " ++ name ]
            ++ formatBrackets "    {" "    ," "    }" (fmap field (zip args resolvedArgs)) ++
            [ "  deriving (Show, Eq)" ])
  emitExport $ name ++ "(..)"

formatBrackets :: String -> String -> String -> [[String]] -> [String]
formatBrackets open sep close ls = format ls ++ [close]
  where format []     = [open]
        format (x:xs) = addPrefix (open ++ " ") x ++ concatMap (addPrefix (sep  ++ " ")) xs

dataConstructors :: [[String]] -> [String]
dataConstructors (c:cs) = addPrefix "  = " c ++ concatMap (addPrefix "  | ") cs
dataConstructors []     = []

formatDescription :: String -> DocComment -> Maybe Description -> [String]
formatDescription _      _       Nothing     = []
formatDescription indent comment (Just desc) = formatLines $ intercalate [""] sections
  where sections = lines (descSummary desc) : [bodySection | not (null bodySection)]
        bodySection = trim [""] $ trim "\t " <$> lines (descBody desc)
        formatLines :: [String] -> [String]
        formatLines []     = []
        formatLines (x:xs) = (indent ++ "-- " ++ startComment comment ++ x) : fmap formatLine' xs
        formatLine'  x     = indent ++ "-- " ++ x
        startComment DocCommentBefore = "| "
        startComment DocCommentAfter  = "^ "

trim :: Eq a => [a] -> [a] -> [a]
trim xs = f . f
  where f = reverse . dropWhile (`elem` xs)

addPrefix :: [a] -> [[a]] -> [[a]]
addPrefix p []     = [p]
addPrefix p (x:xs) = (p ++ x) : xs

escapeStr :: String -> String
escapeStr str = "\"" ++ escapeInner str ++ "\""
  where escapeInner ('\"':xs) = "\\\"" ++ escapeInner xs
        escapeInner ('\\':xs) = "\\\\" ++ escapeInner xs
        escapeInner ('\n':xs) = "\\n"  ++ escapeInner xs
        escapeInner (x   :xs) = x      :  escapeInner xs
        escapeInner []        = []