import Data.Char (ord, toUpper, toLower)
import Data.DList (DList, fromList, toList)
import Data.List (intercalate, isSuffixOf)
import Data.Map (Map)
import qualified Data.Map as M
import Data.Maybe (listToMaybe)
import Data.Monoid (Monoid, mempty, mappend, mconcat)
import System.Directory (doesDirectoryExist, doesFileExist, getDirectoryContents)
import System.Environment (getArgs)
import System.FilePath ((</>), (<.>), joinPath, splitDirectories, splitExtension)

import Paths_binembed (getLibDir)

{- file path -}

type Name = String
type Dir  = Map Name Node
data Node = File | Dir Dir

listRec :: FilePath -> FilePath -> IO (Name, Node)
listRec t f = do
  let tf = t </> f
  isDir  <- doesDirectoryExist tf
  isFile <- doesFileExist tf
  case (isDir, isFile) of
    (True, False) -> do
      cs <- mapM (listRec tf) =<< filter (not . hidden) `fmap` getDirectoryContents tf
      return $ (f, Dir (M.fromList cs))
    (False, True) -> do
      return $ (f, File)
  where
    hidden ('.':_) = True
    hidden _ = False

depthFirstPostOrder :: Monad m => ([Name] -> m a) -> ([Name] -> [a] -> m a) -> [Name] -> Node -> m a
depthFirstPostOrder fact _dact path File = fact path
depthFirstPostOrder fact dact path (Dir dir) =
  mapM (\(name, node) -> depthFirstPostOrder fact dact (path ++ [name]) node) (M.toAscList dir) >>= dact path

textMode :: [String] -> String -> Bool
textMode exts file = any (`isSuffixOf` file) exts

{- code output monoid -}

type Code = DList Char
data Output = Output{ oAS, oCH, oCC, oHS1, oHS2 :: Code }

instance Monoid Output where
  mempty = Output mempty mempty mempty mempty mempty
  Output as ch cc hs1 hs2 `mappend` Output as' ch' cc' hs1' hs2' =
    Output (as `mappend` as') (ch `mappend` ch') (cc `mappend` cc') (hs1 `mappend` hs1') (hs2 `mappend` hs2')

{- code output monad -}

data Env = Env{ eCHName, eCCName, eCInclude, eHSName :: String, eAlignment :: Int, eTextExts :: [String] }
data CodeGen a = CodeGen{ runCodeGen :: Env -> (a, Output) }

instance Monad CodeGen where
  return a = CodeGen $ \_e -> (a, mempty)
  CodeGen a >>= b = CodeGen $ \e -> let (aa, ao) = a e
                                        (bb, bo) = runCodeGen (b aa) e
                                    in  (bb, ao `mappend` bo)

execCodeGen :: Env -> CodeGen () -> Output
execCodeGen e (CodeGen f) = snd (f e)

tell :: Output -> CodeGen ()
tell out = CodeGen $ \_e -> ((), out)

asks :: (Env -> a) -> CodeGen a
asks x = CodeGen $ \e -> (x e, mempty)

{- symbol munging -}

symbols :: String -> String
symbols (s:ss) = symbol0 s ++ concatMap symbol ss

isSymbol0, isSymbol :: Char -> Bool
isSymbol0 c = ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z')
isSymbol c = isSymbol0 c || ('0' <= c && c <= '9')

escape, symbol0, symbol :: Char -> String
escape c = "_u" ++ show (ord c) ++ "_"
symbol0 c | isSymbol0 c = [c]
          | otherwise = escape c
symbol  c | isSymbol  c = [c]
          | otherwise = escape c

{- code generation -}

initialEnv :: [Char] -> Int -> [String] -> String -> Env
initialEnv name@(n:ame) alignment textExts include =
  Env
    { eCHName = map toUpper name
    , eCCName = name
    , eCInclude = include
    , eHSName = toUpper n : ame
    , eAlignment = alignment
    , eTextExts = textExts
    }

initialCode :: CodeGen ()
initialCode = do
  chName <- asks eCHName
  ccName <- asks eCCName
  cInc   <- asks eCInclude
  hsName <- asks eHSName
  tell Output
    { oAS = fromList $  "/*" ++ warning ++ "*/\n.section .rodata\n"
    , oCH = fromList $  "/*" ++ warning ++ "*/\n\
                        \#ifndef " ++ chName ++ "_H\n\
                        \#define " ++ chName ++ "_H 1\n\
                        \#include <" ++ cInc ++ ">\n"
    , oCC = fromList $  "/*" ++ warning ++ "*/\n\
                        \#include \"" ++ ccName <.> "h" ++ "\"\n"
    , oHS1 = fromList $ "{-# LANGUAGE ForeignFunctionInterface #-}\n\
                        \--" ++ warning ++ "\n\
                        \module " ++ hsName ++ "(\n  module Data.BinEmbed\n"
    , oHS2 = fromList $ ") where\n\
                        \import Foreign.Ptr (Ptr)\n\
                        \import Data.ByteString (ByteString)\n\
                        \import Data.Map (fromList)\n\
                        \import Data.BinEmbed\n"
    }
  where
    warning = " autogenerated file; do not edit "

finalCode :: CodeGen ()
finalCode = do
  tell Output
    { oAS = fromList $  ""
    , oCH = fromList $  "#endif\n"
    , oCC = fromList $  ""
    , oHS1 = fromList $ ""
    , oHS2 = fromList $ ""
    }

fileCode :: [String] -> CodeGen (Either (String, String) (String, String))
fileCode path = do
  align <- asks eAlignment
  exts <- asks eTextExts
  let f = joinPath . tail . splitDirectories . joinPath $ path
      s@(s0:ss)    = symbols . intercalate "/" . splitDirectories . joinPath $ path
      e@(e0:es)    = s ++ "_end"
      z            = s ++ "_size"
      fs@(fs0:fss) = s ++ "_file"
      hs  = toLower s0  : ss
      he  = toLower e0  : es
      hfs = toLower fs0 : fss
      text = textMode exts (last path)
  tell Output
    { oAS  = fromList . unlines . map concat $
                          [ [".align ", show align]
                          , [".global ", s], [s, ":"]
                          , [".incbin ", show f]
                          , if text then [".byte 0"] else []
                          , [".global ", e], [e, ":"]
                          , [".global ", z], [".set ", z, ", ", e, " - ", s]
                          ]
    , oCH  = fromList $   "extern char const " ++ s ++ "[];\n\
                          \extern void const " ++ z ++ ";\n"
    , oCC  = fromList $   "static struct binembed_file const " ++ fs ++ " = {\n\
                          \  " ++ s ++ ", (intptr_t) &" ++ z ++ "\n};\n"
    , oHS1 = fromList $   "  , " ++ hfs ++ "\n"
    , oHS2 = fromList $   "foreign import stdcall \"&" ++ s ++ "\" " ++ hs ++ " :: Ptr ()\n\
                          \foreign import stdcall \"&" ++ e ++ "\" " ++ he ++ " :: Ptr ()\n" ++
                          hfs ++ " :: IO ByteString\n" ++
                          hfs ++ " = unBinEmbedFile " ++ hs ++ " " ++ he ++ "\n"
    }
  return $ Left (last path, fs)

dirCode :: Bool -> [String] -> [Either (String, String) (String, String)] -> CodeGen (Either (String, String) (String, String))
dirCode extern path contents = do
  let dc@(dc0:dcs) = symbols . intercalate "/" . splitDirectories . joinPath $ path
      hdc = toLower dc0 : dcs
  tell Output
    { oAS  = fromList $ ""
    , oCH  = fromList $  if extern then "extern\             \ struct binembed_node const " ++ dc ++ "[];\n" else ""
    , oCC  = fromList ( (if extern then "" else "static") ++ " struct binembed_node const " ++ dc ++ "[] = {\n" ) `mappend`
             mconcat  ( map (\edf -> fromList $ case edf of
                          Left  (f, fs) -> "  { " ++ show f ++ ", binembed_file, { .file = &" ++ fs ++ " } },\n"
                          Right (d, ds) -> "  { " ++ show d ++ ", binembed_dir,  { .dir  = " ++ ds ++ " } },\n"
                      ) contents )
            `mappend` fromList "  { 0, binembed_none, { .none = 0 } }\n};\n"
    , oHS1 = fromList $ if extern then "  , " ++ hdc ++ "\n" else ""
    , oHS2 = fromList ( hdc ++ " :: Node (IO ByteString)\n" ++ hdc ++ " = Dir . fromList . tail $\n  [ undefined\n" ) `mappend`
             mconcat  ( map (\edf -> fromList $ case edf of
                          Left  (f, h:fs) -> "  , (" ++ show f ++ ", File " ++ (toLower h : fs) ++ ")\n"
                          Right (d, h:ds) -> "  , (" ++ show d ++ ",      " ++ (toLower h : ds) ++ ")\n"
                      ) contents )
            `mappend` fromList "  ]\n"
    }
  return $ Right (last path, dc)

{- main program -}

main :: IO ()
main = do
  args <- getArgs
  main' args

data Args = Args{ outputS, outputH, outputC, outputHS, infile :: [FilePath] }

instance Monoid Args where
  mempty = Args mempty mempty mempty mempty mempty
  Args a b c d e `mappend` Args x y z u v = Args (a`mappend`x) (b`mappend`y) (c`mappend`z) (d`mappend`u) (e`mappend`v)

toArg :: String -> Args
toArg ('-':'-':'o':'u':'t':'p':'u':'t':'-':'s':'='    :xs) = mempty{ outputS   = [xs] }
toArg ('-':'-':'o':'u':'t':'p':'u':'t':'-':'h':'='    :xs) = mempty{ outputH   = [xs] }
toArg ('-':'-':'o':'u':'t':'p':'u':'t':'-':'c':'='    :xs) = mempty{ outputC   = [xs] }
toArg ('-':'-':'o':'u':'t':'p':'u':'t':'-':'h':'s':'=':xs) = mempty{ outputHS  = [xs] }
toArg a@('-':_) = error $ "binembed: bad argument: " ++ a
toArg xs = mempty{ infile = [xs] }

main' :: [String] -> IO ()
main' args = do
  let args' = mconcat . map toArg $ args
      infile' = listToMaybe . reverse . infile $ args'
      outs  = listToMaybe . reverse . outputS  $ args'
      outh  = listToMaybe . reverse . outputH  $ args'
      outc  = listToMaybe . reverse . outputC  $ args'
      ouths = listToMaybe . reverse . outputHS $ args'
      Just infile'' = infile'
      dirs = splitDirectories infile''
      top  = init dirs
      name = last dirs
      (stem, _ext) = splitExtension name
      textExts = []
  libDir <- getLibDir
  let include = libDir </> "include" </> "binembed.h"
  ins <- mapM (listRec (joinPath top)) =<< lines `fmap` readFile infile''
  let out = execCodeGen (initialEnv stem 64 textExts include) $ do
                initialCode
                tops <- mapM (\(name', node) -> depthFirstPostOrder fileCode (dirCode False) [stem, name'] node) ins
                _ <- dirCode True [stem] tops
                finalCode
  case outs of
    Just outs'  -> writeFile outs'  (toList $ oAS out)
    Nothing     -> return ()
  case ouths of
    Just ouths' -> writeFile ouths' (toList $ oHS1 out `mappend` oHS2 out)
    Nothing     -> return ()
  case outc of
    Just outc'  -> writeFile outc'  (toList $ oCC out)
    Nothing     -> return ()
  case outh of
    Just outh'  -> writeFile outh'  (toList $ oCH out)
    Nothing     -> return ()