Preprocessor.hs

{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE StrictData #-}
module Anitomata.Aseprite.Preprocessor
  ( preprocessor
  , preprocessorWith
  , PreprocessorOpts(..)
  ) where

import Prelude

import Control.Applicative ((<**>))
import Data.Aeson (FromJSON)
import Data.Kind (Type)
import GHC.Generics (Generic)
import GHC.Records (HasField(getField))
import ModuleMunging
  ( DeclBody(..), DeclName(..), ModuleDeclaration(..), ModuleFragment(..), ModuleImport(..)
  , ModuleImportStyle(..), ModuleName(..), buildModule, displayModule
  )
import Text.ParserCombinators.ReadP (ReadP)
import Text.Printf (printf)
import Type.Reflection (Typeable, typeRep)

import Control.Monad qualified as Monad
import Data.Aeson qualified as Aeson
import Data.Aeson.Types qualified as Aeson.Types
import Data.Char qualified as Char
import Data.List qualified as List
import Data.Maybe qualified as Maybe
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.IO qualified as Text.IO
import Options.Applicative qualified as Opt
import System.Directory qualified as Directory
import System.Exit qualified as Exit
import System.FilePath qualified as FilePath
import Text.Read qualified as Read
import Text.ParserCombinators.ReadP qualified as ReadP

preprocessor :: IO ()
preprocessor = do
  opts <- Opt.execParser optsParser
  preprocessorWith opts

preprocessorWith :: PreprocessorOpts -> IO ()
preprocessorWith opts = do
  sourceFile' <- Directory.makeRelativeToCurrentDirectory sourceFile
  atlasPath <- do
    let dir = FilePath.dropFileName sourceFile'
    let baseName = FilePath.takeBaseName sourceFile'
    let atlasPath = dir ++ FilePath.addExtension baseName ".json"
    atlasExists <- Directory.doesFileExist atlasPath
    Monad.unless atlasExists do
      Exit.die $ printf "aseprite2haskell: Aseprite atlas JSON does not exist at %s" atlasPath
    pure atlasPath
  atlas <- Aeson.eitherDecodeFileStrict atlasPath >>= \case
    Left decodeErr -> Exit.die $ "aseprite2haskell: failed to parse JSON - " <> show decodeErr
    Right x -> pure x
  Text.IO.writeFile outputFile $ Text.pack $ displayModule $ buildModule (ModuleNameFromFilePath sourceFile') $ mconcat
    [ mkBuildersAndSlices atlas
    , mkFrames $ frames atlas
    , mkDurations $ frames atlas
    ]
  where
  PreprocessorOpts
    { preprocessorOptsOrigSourceFile = sourceFile
    , preprocessorOptsOutputFile = outputFile
    } = opts

mkBuildersAndSlices :: Atlas -> ModuleFragment
mkBuildersAndSlices Atlas { frames, meta = Meta { frameTags = tags } }
  | null frames = mempty
  | otherwise = mconcat $ flattenFrameInfo $ List.groupBy groupFrameInfo $ zip frames [0 :: Int ..]
  where
  flattenFrameInfo :: [[(FrameInfo, Int)]] -> [ModuleFragment]
  flattenFrameInfo = foldMap \case
    [] -> []
    (FrameInfo { filename }, frameIdx) : xs ->
      let builderName = mkBuilderName filename
          sliceName = printf "%s_slice" builderName :: String
          (dir, mTagRepeat) = dirAndRepeatFromTags filename
       in [ ModuleFragment
              { moduleFragmentImports =
                  [ ModuleImport
                      { moduleImportName = "Prelude"
                      , moduleImportStyle = ModuleImportStyleOpen
                      }
                  , ModuleImport
                      { moduleImportName = "Anitomata"
                      , moduleImportStyle = ModuleImportStyleExplicit
                          $ "AnimBuilder"
                          : "AnimDir(..)"
                          : "AnimSlice"
                          : "AnimSlice_(..)"
                          : builderFn dir
                          : [x | Maybe.isJust mTagRepeat, x <- ["AnimRepeat(..)", "repeatAnim"]]
                      }
                  , ModuleImport
                      { moduleImportName = "Data.Vector.Unboxed"
                      , moduleImportStyle = ModuleImportStyleQualified $ Just "U"
                      }
                  ]
              , moduleFragmentDeclarations =
                  [ ModuleDeclaration True (DeclName builderName) $ DeclBody
                      case mTagRepeat of
                        Nothing ->
                          List.intercalate "\n"
                            [ printf "%s :: AnimBuilder" builderName
                            , printf "%s = %s %s" builderName (builderFn dir) sliceName
                            ]
                        Just (TagRepeat n) ->
                          List.intercalate "\n"
                            [ printf "%s :: AnimBuilder" builderName
                            , printf "%s = repeatAnim (AnimRepeatCount %d) $ %s %s" builderName n (builderFn dir) sliceName
                            ]
                  , ModuleDeclaration True (DeclName sliceName) $ DeclBody $ List.intercalate "\n"
                      [ printf "%s :: AnimSlice" sliceName
                      , printf "%s =" sliceName
                      ,        "  AnimSlice"
                      , printf "    { animSliceDir = %s" (toAnimDir dir)
                      , printf "    , animSliceFrameDurs = U.slice %d %d %s" frameIdx (1 + length xs) durationsVecName
                      , printf "    , animSliceFrames = U.slice %d %d %s" frameIdx (1 + length xs) framesVecName
                      ,        "    }"
                      ]
                  ]
              }
          ]

  groupFrameInfo :: (FrameInfo, Int) -> (FrameInfo, Int) -> Bool
  groupFrameInfo (FrameInfo { filename = x }, _) (FrameInfo { filename = y }, _) =
    getField @"file" x == getField @"file" y && getField @"tag" x == getField @"tag" y

  dirAndRepeatFromTags :: FilenameField -> (Direction, Maybe TagRepeat)
  dirAndRepeatFromTags FilenameField { file, tag } =
    maybe (Forward, Nothing) ((,) <$> direction <*> getField @"repeat")
      $ flip List.find tags \case
          FrameTag { name = nameField } ->
            file == getField @"file" nameField && tag == getField @"tag" nameField

  toAnimDir :: Direction -> String
  toAnimDir = \case
    Forward -> "AnimDirForward"
    Reverse -> "AnimDirBackward"
    Pingpong -> "AnimDirForward"
    PingpongReverse -> "AnimDirBackward"

  builderFn :: Direction -> String
  builderFn = \case
    Forward -> "fromAnimSlice"
    Reverse -> "fromAnimSlice"
    Pingpong -> "pingpongAnimSlice"
    PingpongReverse -> "pingpongAnimSlice"

mkFrames :: [FrameInfo] -> ModuleFragment
mkFrames = \case
  [] -> mempty
  xs ->
    ModuleFragment
      { moduleFragmentImports =
          [ ModuleImport
              { moduleImportName = "Prelude"
              , moduleImportStyle = ModuleImportStyleOpen
              }
          , ModuleImport
              { moduleImportName = "Data.Vector.Unboxed"
              , moduleImportStyle = ModuleImportStyleQualified $ Just "U"
              }
          , ModuleImport
              { moduleImportName = "Anitomata"
              , moduleImportStyle = ModuleImportStyleExplicit ["AnimFrame", "AnimFrame_(..)"]
              }
          ]
      , moduleFragmentDeclarations =
          [ ModuleDeclaration False (DeclName framesVecName) $ DeclBody $ List.intercalate "\n"
              [ printf "%s :: U.Vector AnimFrame" framesVecName
              , printf "%s = U.fromListN %d" framesVecName $ length xs
              , printf "  [ %s" $ List.intercalate "\n  , " (sourceRect . frame <$> xs)
              ,        "  ]"
              ]
          ]
      }
  where
  sourceRect :: Frame -> String
  sourceRect Frame { x, y, w, h } =
    printf "AnimFrame { animFrameX = %d, animFrameY = %d, animFrameW = %d, animFrameH = %d }" x y w h

mkDurations :: [FrameInfo] -> ModuleFragment
mkDurations = \case
  [] -> mempty
  xs ->
    ModuleFragment
      { moduleFragmentImports =
          [ ModuleImport
              { moduleImportName = "Prelude"
              , moduleImportStyle = ModuleImportStyleOpen
              }
          , ModuleImport
              { moduleImportName = "Data.Vector.Unboxed"
              , moduleImportStyle = ModuleImportStyleQualified $ Just "U"
              }
          ]
      , moduleFragmentDeclarations =
          [ ModuleDeclaration False (DeclName durationsVecName) $ DeclBody $ List.intercalate "\n"
              [ printf "%s :: U.Vector Double" durationsVecName
              , printf "%s = U.fromListN %d" durationsVecName $ length xs
              , printf "  [ %s" $ List.intercalate "\n  , " (show . toSeconds . duration <$> xs)
              ,        "  ]"
              ]
          ]
      }
  where
  toSeconds :: Int -> Double
  toSeconds ms = fromIntegral ms / 1000

type Atlas :: Type
data Atlas = Atlas
  { frames :: [FrameInfo]
  , meta :: Meta
  } deriving stock (Generic)
    deriving anyclass (FromJSON)

type Meta :: Type
newtype Meta = Meta
  { frameTags :: [FrameTag]
  } deriving stock (Generic)
    deriving anyclass (FromJSON)

type FrameTag :: Type
data FrameTag = FrameTag
  { name :: TagNameField
  , from :: Int
  , to :: Int
  , direction :: Direction
  , repeat :: Maybe TagRepeat
  } deriving stock (Generic)
    deriving anyclass (FromJSON)

type TagNameField :: Type
data TagNameField = TagNameField
  { file :: Text
  , tag :: Text
  } deriving stock (Show)

instance FromJSON TagNameField where
  parseJSON :: Aeson.Value -> Aeson.Types.Parser TagNameField
  parseJSON = parseJSONViaReadP tagnameFieldParser

tagnameFieldParser :: ReadP TagNameField
tagnameFieldParser = do
  file <- Text.pack <$> ReadP.munch1 (/= '|')
  _ <- ReadP.char '|'
  tag <- Text.pack <$> ReadP.munch1 (/= '|')
  ReadP.eof
  pure TagNameField { file, tag }

type TagRepeat :: Type
newtype TagRepeat = TagRepeat Int
  deriving stock (Show)

instance FromJSON TagRepeat where
  parseJSON :: Aeson.Value -> Aeson.Types.Parser TagRepeat
  parseJSON = parseJSONViaReadP tagRepeatParser

tagRepeatParser :: ReadP TagRepeat
tagRepeatParser = do
  count <- intParser
  ReadP.eof
  pure $ TagRepeat count

type Direction :: Type
data Direction
  = Forward
  | Reverse
  | Pingpong
  | PingpongReverse

instance FromJSON Direction where
  parseJSON :: Aeson.Value -> Aeson.Types.Parser Direction
  parseJSON = Aeson.withText "FromJSON Direction" \case
    "forward" -> pure Forward
    "reverse" -> pure Reverse
    "pingpong" -> pure Pingpong
    "pingpong_reverse" -> pure PingpongReverse
    other -> fail $ "Invalid direction: " <> show other

type FrameInfo :: Type
data FrameInfo = FrameInfo
  { filename :: FilenameField
  , frame :: Frame
  , duration :: Int
  } deriving stock (Generic, Show)
    deriving anyclass (FromJSON)

framesVecName :: String
framesVecName = "frames"

durationsVecName :: String
durationsVecName = "durations"

type FilenameField :: Type
data FilenameField = FilenameField
  { file :: Text
  , tag :: Text
    -- | This is the frame index in the specific .aseprite file, NOT the frame
    -- index in the overall texture atlas frames. It is only parsed here for
    -- debugging's sake to cross-reference individual .aseprite files.
  , frameIndex :: Int
  } deriving stock (Show)

mkBuilderName :: FilenameField -> String
mkBuilderName FilenameField { file, tag } = printf "%s_%s" file tag

instance FromJSON FilenameField where
  parseJSON :: Aeson.Value -> Aeson.Types.Parser FilenameField
  parseJSON = parseJSONViaReadP filenameFieldParser

filenameFieldParser :: ReadP FilenameField
filenameFieldParser = do
  file <- Text.pack . map sanitize <$> ReadP.munch1 (/= '|')
  _ <- ReadP.char '|'
  tag <- Text.pack . map sanitize <$> ReadP.munch1 (/= '|')
  _ <- ReadP.char '|'
  frameIndex <- intParser
  ReadP.eof
  pure FilenameField { file, tag, frameIndex }
  where
  sanitize :: Char -> Char
  sanitize = \case
    '-' -> '_'
    c -> c

type Frame :: Type
data Frame = Frame
  { x :: Int
  , y :: Int
  , w :: Int
  , h :: Int
  } deriving stock (Generic, Show)
    deriving anyclass (FromJSON)

optsParser :: Opt.ParserInfo PreprocessorOpts
optsParser = Opt.info (preprocessorOptsParser <**> Opt.helper) $ mconcat
  [ Opt.fullDesc
  , Opt.progDesc "Convert an Aseprite texture atlas JSON into Haskell code"
  , Opt.header "aseprite2haskell - Aseprite -> Haskell preprocessor"
  ]

type PreprocessorOpts :: Type
data PreprocessorOpts = PreprocessorOpts
  { preprocessorOptsOrigSourceFile :: FilePath
  , preprocessorOptsOutputFile :: FilePath
  } deriving stock (Show)

preprocessorOptsParser :: Opt.Parser PreprocessorOpts
preprocessorOptsParser = do
  preprocessorOptsOrigSourceFile <- Opt.argument Opt.str $ mconcat
    [ Opt.metavar "SOURCE"
    , Opt.help "Original source filepath (passed by GHC)"
    ]
  -- It seems this positional argument only matters when multiple preprocessors
  -- are stacked (e.g. CPP and this custom one). It is ignored here.
  _ <- Opt.argument @String Opt.str $ mconcat
    [ Opt.metavar "INPUT"
    , Opt.help "Input filepath (passed by GHC)"
    ]
  preprocessorOptsOutputFile <- Opt.argument Opt.str $ mconcat
    [ Opt.metavar "OUTPUT"
    , Opt.help "Output filepath (passed by GHC)"
    ]
  pure PreprocessorOpts
    { preprocessorOptsOrigSourceFile
    , preprocessorOptsOutputFile
    }

parseJSONViaReadP
  :: forall a
   . (Show a, Typeable a)
  => ReadP a
  -> Aeson.Types.Value
  -> Aeson.Types.Parser a
parseJSONViaReadP parser =
  Aeson.withText (printf "FromJSON %s" $ show $ typeRep @a) \case
    t -> case ReadP.readP_to_S parser s of
      [] -> fail $ printf "Parse failed on input: %s" s
      [(x, _)] -> pure x
      xs -> fail $ printf "Produced too many parses - input: %s, parses: %s" s $ show xs
      where
      s = Text.unpack t

intParser :: ReadP Int
intParser = do
  digitChars <- ReadP.many1 digitParser
  case Read.readMaybe digitChars of
    Nothing -> fail $ printf "Failed to read digits as Int: %s" digitChars
    Just x -> pure x

digitParser :: ReadP Char
digitParser = ReadP.satisfy Char.isDigit