Preprocessor.hs

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

import Prelude

import Control.Applicative ((<**>))
import Data.Aeson (FromJSON)
import Data.Kind (Type)
import Data.String.Interpolate (__i, i)
import GHC.Generics (Generic)
import GHC.Records (HasField(getField))
import ModuleMunging
  ( DeclBody(..), DeclName(..), ModuleDeclaration(..), ModuleFragment(..), ModuleImport(..)
  , ModuleImportStyle(..), ModuleName(..), buildModule, displayModule
  )

import Control.Monad qualified as Monad
import Data.Aeson qualified as Aeson
import Data.Aeson.Types qualified as Aeson.Types
import Data.Attoparsec.Text qualified as Attoparsec
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

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 [i|aseprite2haskell: Aseprite atlas JSON does not exist at "#{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 = [i|#{builderName}_slice|] :: 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 ->
                          [__i|
                            #{builderName} :: AnimBuilder
                            #{builderName} = #{builderFn dir} #{sliceName}
                          |]
                        Just (TagRepeat n) ->
                          [__i|
                            #{builderName} :: AnimBuilder
                            #{builderName} = repeatAnim (AnimRepeatCount #{n}) $ #{builderFn dir} #{sliceName}
                          |]
                  , ModuleDeclaration True (DeclName sliceName) $ DeclBody
                      [__i|
                        #{sliceName} :: AnimSlice
                        #{sliceName} =
                          AnimSlice
                            { animSliceDir = #{toAnimDir dir}
                            , animSliceFrameDurs = U.slice #{frameIdx} #{1 + length xs} #{durationsVecName}
                            , animSliceFrames = U.slice #{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"
              [ [i|#{framesVecName} :: U.Vector AnimFrame|]
              , [i|#{framesVecName} = U.fromListN #{length xs}|]
              , "  [ " <> List.intercalate "\n  , " (sourceRect . frame <$> xs)
              , "  ]"
              ]
          ]
      }
  where
  sourceRect :: Frame -> String
  sourceRect Frame { x, y, w, h } =
    [i|AnimFrame { animFrameX = #{x}, animFrameY = #{y}, animFrameW = #{w}, animFrameH = #{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"
              [ [i|#{durationsVecName} :: U.Vector Double|]
              , [i|#{durationsVecName} = U.fromListN #{length xs}|]
              , "  [ " <> 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
  }

instance FromJSON TagNameField where
  parseJSON :: Aeson.Value -> Aeson.Types.Parser TagNameField
  parseJSON = Aeson.withText "FromJSON TagNameField" \t -> do
    case Attoparsec.parseOnly tagnameFieldParser t of
      Left err -> fail err
      Right x -> pure x

tagnameFieldParser :: Attoparsec.Parser TagNameField
tagnameFieldParser = do
  file <- Attoparsec.takeWhile1 (/= '|')
  _ <- Attoparsec.char '|'
  tag <- Attoparsec.takeWhile1 (/= '|')
  pure TagNameField { file, tag }

type TagRepeat :: Type
newtype TagRepeat = TagRepeat Int

instance FromJSON TagRepeat where
  parseJSON :: Aeson.Value -> Aeson.Types.Parser TagRepeat
  parseJSON = Aeson.withText "FromJSON TagRepeat" \t -> do
    case Attoparsec.parseOnly tagRepeatParser t of
      Left err -> fail err
      Right x -> pure x

tagRepeatParser :: Attoparsec.Parser TagRepeat
tagRepeatParser = TagRepeat <$> Attoparsec.decimal

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 } = [i|#{file}_#{tag}|]

instance FromJSON FilenameField where
  parseJSON :: Aeson.Value -> Aeson.Types.Parser FilenameField
  parseJSON = Aeson.withText "FromJSON FilenameField" \t -> do
    case Attoparsec.parseOnly filenameFieldParser t of
      Left err -> fail err
      Right x -> pure x

filenameFieldParser :: Attoparsec.Parser FilenameField
filenameFieldParser = do
  file <- Text.map sanitize <$> Attoparsec.takeWhile1 (/= '|')
  _ <- Attoparsec.char '|'
  tag <- Text.map sanitize <$> Attoparsec.takeWhile1 (/= '|')
  _ <- Attoparsec.char '|'
  frameIndex <- Attoparsec.decimal
  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
    }