{-# LANGUAGE ApplicativeDo         #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE NamedFieldPuns        #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE RecordWildCards       #-}
{-# LANGUAGE ScopedTypeVariables   #-}

module Main (main) where

import           Control.Arrow (second)
import           Control.Applicative             (Alternative(..), (<|>), (<**>))
import           Control.Applicative.MultiExcept
import           Control.Exception               (bracket)
import           Control.Monad                   (unless)
import qualified Data.ByteString                 as B
import qualified Data.ByteString.UTF8            as BSU
import qualified Data.DList.NonEmpty             as DNE
import           Data.Foldable                   (traverse_)
import           Data.List                       (scanl', foldl', intercalate)
import           Data.List.NonEmpty              (NonEmpty(..))
import           Data.Maybe                      (catMaybes, isJust)
import           Language.C
import           Language.C.Data.Ident
import           Options.Applicative             (Parser)
import qualified Options.Applicative             as O
import           System.Directory                (getCurrentDirectory)
import           System.Environment
import           System.FilePath                 ((</>))
import           System.IO
import           System.Process
import qualified Text.PrettyPrint                as PP

data Options
  = Options
  { sizeof           :: Bool
  , uses             :: Bool
  , padding          :: Bool
  , noColor          :: Bool
  , inFile           :: FilePath
  , nonZeroOnPadding :: Bool
  } deriving Show

data Handles
  = Handles
  { input  :: Handle
  , output :: Handle
  }

data Env
  = Env
  { options :: Options
  , handles :: Handles
  }

optionsParser :: Parser Options
optionsParser = Options
  <$> O.switch
    ( O.long "sizeof"
    <> O.short 's'
    <> O.help "Output size of fields (excluding padding)")
  <*> O.switch
    ( O.long "uses"
    <> O.short 'u'
    <> O.help "Output bytes used by field (including padding)")
  <*> O.switch
    ( O.long "padding"
    <> O.short 'p'
    <> O.help "Output bytes of padding that follow a field")
  <*> O.switch
    ( O.long "no-color"
    <> O.short 'n'
    <> O.help "Don't output color")
  <*> O.strOption
    ( O.long "input"
    <> O.short 'i'
    <> O.metavar "FILE"
    <> O.help "Input file (default is stdin)"
    <> O.value "-"
    )
  <*> O.switch
    ( O.long "nonzero-on-padding"
    <> O.short 'z'
    )

getOptions :: IO Options
getOptions = O.execParser optInfo

optInfo :: O.ParserInfo Options
optInfo = O.info (optionsParser <**> O.helper)
  ( O.fullDesc
  <> O.progDesc "Print structure padding and field size info"
  <> O.header "Struct Inspector" )

failWithHelp :: B.ByteString -> IO a
failWithHelp str = do
  B.putStr "Error: "
  B.putStr str
  B.putStr "\n\n"
  O.handleParseResult . O.Failure $ O.parserFailure pprefs optInfo (O.ShowHelpText Nothing) mempty

anyInfo :: Options -> Bool
anyInfo Options{..} = or debug
  where
    debug :: [Bool]
    debug = [padding, uses, sizeof]

pprefs :: O.ParserPrefs
pprefs = O.prefs mempty

environmentFromOptions :: Options -> IO Env
environmentFromOptions options@Options{inFile} = do
  unless (anyInfo options) $
    failWithHelp "Please specify at least one of '--padding', '--uses', or '--sizeof'"
  let output = stdout
  input <- case inFile of
    "-" -> pure stdin
    "" -> failWithHelp "Empty input file name"
    _ -> openFile inFile ReadMode
  let handles = Handles{..}
  pure $ Env{..}

createEnvironment :: IO Env
createEnvironment = getOptions >>= environmentFromOptions

data Field
  = FieldNamed
    { name :: String
    , declr :: CDeclr
    , typeSpec :: CTypeSpec
    }
  | FieldAnonymousStruct
    { inner :: [Field]
    }
  | FieldNamedStruct
    { name :: String
    , inner :: [Field]
    }
  | FieldAnonymousUnion
    { inner :: [Field]
    }
  | FieldNamedUnion
    { name :: String
    , inner :: [Field]
    }
  deriving Show

data Struct
  = Struct
    { name :: String
    , fields :: [Field]
    }
  | TypedefStruct
    { name :: String
    , fields :: [Field]
    }
  deriving Show

data QueryErr
  = FieldWithoutType String
  | NonStructFieldWithoutName
  deriving Show

type QM = MultiExcept QueryErr

main :: IO ()
main = createEnvironment >>= run

cFile :: FilePath
cFile = ".struct.debug.c"

ifThenElse :: Bool -> a -> a -> a
ifThenElse c a b = if c then a else b

cPrelude :: Env -> B.ByteString
cPrelude Env{options = Options{..}} = B.intercalate "\n"
  $ colors <>
  [ "#include <stdio.h>"
  , "// End of prelude\n\n"
  ]
  where
    colors :: [B.ByteString]
    colors
      = ("#define " <>)
      . (uncurry (<>))
      . second (" " <>)
      . (second $ ifThenElse noColor "\"\"")
      <$>
      [ ("RESET", "\"\\033[0m\"")
      , ("BLACK", "\"\\033[30m\"")
      , ("RED", "\"\\033[31m\"")
      , ("GREEN", "\"\\033[32m\"")
      , ("YELLOW", "\"\\033[33m\"")
      , ("BLUE", "\"\\033[34m\"")
      , ("MAGENTA", "\"\\033[35m\"")
      , ("CYAN", "\"\\033[36m\"")
      , ("WHITE", "\"\\033[37m\"")
      , ("BOLDBLACK", "\"\\033[1m\\033[30m\"")
      , ("BOLDRED", "\"\\033[1m\\033[31m\"")
      , ("BOLDGREEN", "\"\\033[1m\\033[32m\"")
      , ("BOLDYELLOW", "\"\\033[1m\\033[33m\"")
      , ("BOLDBLUE", "\"\\033[1m\\033[34m\"")
      , ("BOLDMAGENTA", "\"\\033[1m\\033[35m\"")
      , ("BOLDCYAN", "\"\\033[1m\\033[36m\"")
      , ("BOLDWHITE", "\"\\033[1m\\033[37m\"")
      ]

getEnvDef :: String -> String -> IO String
getEnvDef def env = flip fmap (lookupEnv env) $ \case
  Nothing -> def
  Just a -> a

run :: Env -> IO ()
run env@Env{handles = Handles{..}} = do
  source <- hGetContents input
  cc <- getEnv "CC"
  dir <- getCurrentDirectory
  -- putStrLn $ "Using C compiler ($CC): " <> cc
  cpath <- getEnvDef "" "CPATH"
  includePath <- getEnvDef "" "C_INCLUDE_PATH"
  preprocessed <- readProcess "cc" ["-std=c11", "-E", "-", "-I" <> cpath, "-I" <> includePath] source
  case parseC (BSU.fromString preprocessed) $ initPos "stdin" of
    Left err -> hPrint stderr err
    Right ast -> case runMultiExcept $ getStructsFromTranslUnit ast of
      Left errs -> traverse_ (hPrint stderr) errs
      Right structs -> do
        bracket (openFile cFile WriteMode) hClose $ \handle -> do
          B.hPut handle $ cPrelude env
          hPutStrLn handle source
          hPutStrLn handle $ mkDebugger env structs
        callCommand $ intercalate " "
          [ cc
          , "-O0"
          , "-s"
          , cFile
          ]
        callProcess (dir </> "a.out") []


mkDebugger :: Env -> [Struct] -> String
mkDebugger env@Env{options = Options{..}} structs
  =  "int main(void) {\n"
  <> strif nonZeroOnPadding "  int __exit_code = 0;\n"
  <> out
  <> concatMap (debugStruct env) structs
  <> strif nonZeroOnPadding "  return __exit_code;\n"
  <> "}"
  where
    out :: String
    out 
      =  strif (indirect env)
      $  "  FILE *__out = stdout;\n"
      <> "  char *__struct_out_string;\n"
      <> "  size_t __struct_out_size;\n"

structName :: String
structName = "__s"

debugPrint :: Bool
debugPrint = False

strif :: Bool -> String -> String
strif b s = if b then s else ""

indirect :: Env -> Bool
indirect Env{options = Options{..}} = padding

{- The best Haskell printf implementation isn't some type level
 - tomfoolery, it's to shell out to C. Thanks for coming to my
 - Ted Talk.
 -}
printf :: Env -> Int -> String -> [String] -> String
printf env n s args = replicate n ' ' <> start <> argsStr <> ");\n"
  where
    argsStr :: String
    argsStr = show s <> ", " <> intercalate ", " args

    start :: String
    start = if indirect env
      then "fprintf(__out, "
      else "printf("

puts :: Env -> Int -> String -> String
puts env n s = replicate n ' ' <> start <> show s <> ");\n"
  where
    start :: String
    start = if indirect env
      then "fprintf(__out, \"%s\\n\", "
      else "puts("

printSizePaths :: Env -> Int -> String -> String -> String
printSizePaths env@Env{options = Options{..}} n path nextPath
   =  "    {\n"
   <> "      size_t __size_of = sizeof(" <> path <> ");\n"
   <> "      size_t __uses    = ((size_t)&" <> nextPath <> ") - ((size_t)&" <> path <> ");\n"
   <> "      size_t __padding = __uses - __size_of;\n"
   <> strif padding "      __struct_total_padding += __padding;\n"
   <> exitPadding
   <> printSizeof
   <> printUses
   <> printPadding
   <> db
   <> "    }\n"
  where
    db :: String
    db = if debugPrint
      then printIndent env 4 n
        <> puts env 4 ("// path: " <> path <> "\n")
        <> printIndent env 4 n
        <> puts env 4 ("// nextPath: " <> nextPath <> "\n")
      else ""

    printSizeof :: String
    printSizeof
      =  strif sizeof
      $  printIndent env 6 n
      <> printf env 6 "// sizeof: %zu\n" ["__size_of"]

    printUses :: String
    printUses
      =  strif uses
      $  printIndent env 6 n
      <> printf env 6 "// uses: %zu\n" ["__uses"]

    printPadding :: String
    printPadding
      = strif padding
      $  printIndent env 6 n
      <> printf env 6 "// padding: %s%zu%s\n" ["__padding > 0 ? RED : GREEN", "__padding", "RESET"]

    exitPadding :: String
    exitPadding
      = strif nonZeroOnPadding
      $ "      if (__padding > 0) { __exit_code = 1; }\n"

debugStruct' :: Env -> Bool -> String -> [Field] -> String
debugStruct' env@Env{options = Options{..}} istypedef s fields
  =  "  {\n"
  <> strif padding "    size_t __struct_total_padding = 0;\n"
  <> out
  <> decl
  <> "    printf(\"// sizeof: %zu\\n\", sizeof(" <> path <> "));\n"
  <> start
  <> printSubs env 0 path nextPath fields
  <> end
  <> total
  <> strif (indirect env) "    putc(0, __out);\n"
  <> strif (indirect env) "    fclose(__out);\n"
  <> strif (indirect env) "    printf(\"%s\", __struct_out_string);\n"
  <> "    putc('\\n', stdout);\n"
  <> "  }\n"
  where
    out :: String
    out
      =  strif padding
      $  "    __out = open_memstream(&__struct_out_string, &__struct_out_size);\n"

    path :: String
    path = structName <> "[0]"

    nextPath :: String
    nextPath = structName <> "[1]"

    decl :: String
    decl = if istypedef
      then ("    " <> s <> " " <> structName <> "[2];\n")
      else ("    struct " <> s <> " " <> structName <> "[2];\n")

    start :: String
    start = puts env 4 $ if istypedef
      then "typedef struct {"
      else "struct " <> s <> " {"

    end :: String
    end = puts env 4 $ if istypedef
      then "} " <> s <> ";"
      else "};"

    total :: String
    total = strif padding
      $  "    printf(\"// total struct padding: %s%zu%s\\n\""
      <> ", __struct_total_padding > 0 ? RED : GREEN"
      <> ", __struct_total_padding"
      <> ", RESET);\n"

debugStruct :: Env -> Struct -> String
debugStruct env (Struct s fs) = debugStruct' env False s fs
debugStruct env (TypedefStruct s fs) = debugStruct' env True s fs

printIndent :: Env -> Int -> Int -> String
printIndent _ _ 0 = ""
printIndent env n m = replicate n ' ' <>
  if indirect env
  then "fprintf(__out, \"%" <> show m <> "s\", \"\");\n"
  else "printf(\"%" <> show m <> "s\", \"\");\n"

altList :: (Foldable f, Alternative a) => f (a b) -> a b
altList = foldl' (<|>) empty

altMap :: (Functor f, Foldable f, Alternative a) => (b -> a c) -> f b -> a c
altMap f = altList . fmap f

firstPath :: Field -> Maybe String
firstPath = \case
  FieldNamed{..} -> pure name
  FieldAnonymousStruct{..} -> altMap firstPath inner
  FieldAnonymousUnion{..} -> altMap firstPath inner
  FieldNamedStruct{..} -> pure name
  FieldNamedUnion{..} -> pure name

(<.>) :: String -> String -> String
(<.>) a b = a <> "." <> b

printSubs :: Env -> Int -> String -> String -> [Field] -> String
printSubs env n path nextPath fields =
  concat $ zipWith (flip (debugField env (n + 2) path)) fields nexts
  where
    nexts :: [String]
    nexts = reverse $ scanl' f nextPath $ reverse $ tail $ fieldFirsts

    f :: String -> Maybe String -> String
    f _ (Just a) = a
    f a Nothing = a

    fieldFirsts :: [Maybe String]
    fieldFirsts = fmap (path <.>) . firstPath <$> fields

printSubsUnion :: Env -> Int -> String -> String -> [Field] -> String
printSubsUnion env n path nextPath fields =
  concatMap (debugField env (n + 2) path nextPath) fields

debugField :: Env -> Int -> String -> String -> Field -> String
debugField env n path nextPath = \case
  (FieldNamed name declr ts) ->
    let newPath = path <.> name
     in printSizePaths env n newPath nextPath
       <> printIndent env 4 n
       <> puts env 4 (PP.render (pretty ts) <> " " <> PP.render (pretty declr) <> ";")
  FieldAnonymousStruct{..} ->
    printIndent env 4 n
    <> puts env 4 "struct {"
    <> printSubs env n path nextPath inner
    <> printIndent env 4 n
    <> puts env 4 "};"
  FieldNamedStruct{..} ->
    printIndent env 4 n
    <> puts env 4 "struct {"
    <> printSubs env n (path <.> name) nextPath inner
    <> printIndent env 4 n
    <> puts env 4 ("} " <> name <> ";")
  FieldAnonymousUnion{..} ->
    printIndent env 4 n
    <> puts env 4 "union {"
    <> printSubsUnion env n path nextPath inner
    <> printIndent env 4 n
    <> puts env 4 "};"
  FieldNamedUnion{..} ->
    printIndent env 4 n
    <> puts env 4 "union {"
    <> printSubsUnion env n (path <.> name) nextPath inner
    <> printIndent env 4 n
    <> puts env 4 ("} " <> name <> ";")

getStructsFromTranslUnit :: CTranslUnit -> QM [Struct]
getStructsFromTranslUnit (CTranslUnit decls _) = concat <$> traverse getStructsFromExtDecl decls

getStructsFromExtDecl :: CExtDecl -> QM [Struct]
getStructsFromExtDecl (CDeclExt a) = getStructsFromDecl a
getStructsFromExtDecl _ = pure []

getStructsFromDecl :: CDeclaration NodeInfo -> QM [Struct]
getStructsFromDecl (CDecl specs decs _) =
  concat <$> traverse (flip getStructsFromDeclSpec $ getTypedef specs decs) specs
getStructsFromDecl _ = pure []

getTypedef :: [CDeclSpec] -> [(Maybe CDeclr, b, c)] -> Maybe String
getTypedef specs decls
  | any isTd specs = altMap getName $ catMaybes $ fst3 <$> decls
  | otherwise = Nothing
  where
    isTd :: CDeclSpec -> Bool
    isTd (CStorageSpec (CTypedef _)) = True
    isTd _ = False

    getName :: CDeclr -> Maybe String
    getName (CDeclr (Just (Ident s _ _)) _ _ _ _) = Just s
    getName _ = Nothing

getStructsFromDeclSpec :: CDeclSpec -> Maybe String -> QM [Struct]
getStructsFromDeclSpec (CTypeSpec a) = getStructsFromCTypeSpec a
getStructsFromDeclSpec _ = const $ pure []

getStructsFromCTypeSpec :: CTypeSpec -> Maybe String -> QM [Struct]
getStructsFromCTypeSpec (CSUType su _) = getStructsFromCSU su
getStructsFromCTypeSpec _ = const $ pure []

hasParent :: NodeInfo -> Bool
hasParent = \case
  OnlyPos pos _ -> f pos
  NodeInfo pos _ _ -> f pos
  where
    f :: Position -> Bool
    f pos = isJust $ posParent pos

getStructsFromCSU :: CStructUnion -> Maybe String -> QM [Struct]
getStructsFromCSU (CStruct _ _ _ _ ni) _
  | hasParent ni = pure []
getStructsFromCSU (CStruct CStructTag _ Nothing _ _) (Just s) = pure $ [TypedefStruct s []]
getStructsFromCSU (CStruct CStructTag (Just (Ident s _ _)) Nothing _ _) Nothing = pure $ [Struct s []]
getStructsFromCSU (CStruct CStructTag (Just (Ident s _ _)) (Just decls) _ _) Nothing
  = pure . Struct s <$> concat <$> traverse getFieldsFromDecl decls
getStructsFromCSU (CStruct CStructTag _ (Just decls) _ _) (Just s)
  = pure . TypedefStruct s <$> concat <$> traverse getFieldsFromDecl decls
getStructsFromCSU _ _ = pure []

fst3 :: (a, b, c) -> a
fst3 (a, _, _) = a

getFieldsFromDecl :: CDeclaration NodeInfo -> QM [Field]
getFieldsFromDecl (CDecl specs trips _) =
  let pairs = concatMap (getFieldsFromDeclarators . fst3) trips
      (names, _) = unzip pairs
      typeNode = getTypeFromDeclSpecs specs
  in case typeNode of
    Nothing -> case names of
    -- I really hope this never happens
      [] -> throwError $ FieldWithoutType []
      x : xs -> throwErrors
        $ DNE.fromNonEmpty
        $ FieldWithoutType <$> (x :| xs)
    -- TODO test whether we need struct name here
    Just (CSUType (CStruct CStructTag _ Nothing _ _) _) -> pure $ case names of
      [] -> [FieldAnonymousStruct []]
      ns -> flip FieldNamedStruct [] <$> ns
    Just (CSUType (CStruct CStructTag _ (Just decls) _ _) _) -> do
      fields <- concat <$> traverse getFieldsFromDecl decls
      pure $ case names of
        [] -> [FieldAnonymousStruct fields]
        ns -> flip FieldNamedStruct fields <$> ns
    Just (CSUType (CStruct CUnionTag _ Nothing _ _) _) -> pure $ case names of
      [] -> [FieldAnonymousUnion []]
      ns -> flip FieldNamedUnion [] <$> ns
    Just (CSUType (CStruct CUnionTag _ (Just decls) _ _) _) -> do
      fields <- concat <$> traverse getFieldsFromDecl decls
      pure $ case names of
        [] -> [FieldAnonymousUnion fields]
        ns -> flip FieldNamedUnion fields <$> ns
    Just t -> case pairs of
      [] -> throwError $ NonStructFieldWithoutName
      ns -> pure $ uncurry FieldNamed <$> ns <*> pure t
getFieldsFromDecl _ = pure []

getTypeFromDeclSpecs :: [CDeclarationSpecifier NodeInfo] -> Maybe CTypeSpec
getTypeFromDeclSpecs = foldl' (<|>) Nothing . fmap getTypeFromDeclSpec

getTypeFromDeclSpec :: CDeclarationSpecifier NodeInfo -> Maybe CTypeSpec
getTypeFromDeclSpec (CTypeSpec a) = Just a
getTypeFromDeclSpec _ = Nothing

getFieldsFromDeclarators :: Maybe (CDeclarator NodeInfo) -> [(String, CDeclr)]
getFieldsFromDeclarators (Just declr@(CDeclr (Just (Ident s _ _)) _ _ _ _)) = [(s, declr)]
getFieldsFromDeclarators _ = []