Font.hs

module Codec.Ktx2.Font
  ( -- * KTX font bundles
    Bundle(..)
  , loadBundle
  , loadBundleBytes
  , loadBundleFile
  , freeBundle
  , bundleFont
    -- * Shaping context
  , StackContext(..)
  , createStackContext
  , pushBundleFont
  , destroyStackContext
    -- ** Associated data utilities
  , lookupAtlas
  , lookupBundled
  , mapWithBundle
  , mapBundled
    -- * Format internals
  , pattern KTX_KEY_atlas
  , pattern KTX_KEY_kbts
  , pattern KTX_KEY_kbts_version
  , kbtsVersion
  ) where

import Codec.Compression.Zstd qualified as Zstd
import Codec.Ktx.KeyValue qualified as KVD
import Codec.Ktx2 qualified as Ktx2
import Codec.Ktx2.Read qualified as Ktx2
import Codec.Ktx2.Write qualified as Ktx2
import Control.Concurrent (MVar, newMVar, takeMVar)
import Control.Exception (bracket)
import Control.Monad
import Data.Aeson (encode, eitherDecodeFileStrict, eitherDecodeStrict)
import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.Foldable (toList)
import Data.IntMap (IntMap)
import Data.IntMap qualified as IntMap
import Data.Map qualified as Map
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Traversable (for)
import Data.Vector.Generic qualified as Vector
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
import GHC.Generics (Generic, Generic1)
import Graphics.MSDF.Atlas.Compact (Compact, compact)
import Graphics.MSDF.Atlas.Compact qualified as Compact
import Graphics.MSDF.Atlas.Layout qualified as Layout
import KB.Text.Shape qualified as TextShape
import KB.Text.Shape.FFI.Enums qualified as KBTS
import KB.Text.Shape.FFI.Handles (Font(..), intHandle)
import KB.Text.Shape.Font (FontData(..), createFont, destroyFont, withFontData)
import KB.Text.Shape.Font qualified as Font

pattern KTX_KEY_atlas :: Text
pattern KTX_KEY_atlas = "msdf-atlas"

pattern KTX_KEY_kbts :: Text
pattern KTX_KEY_kbts = "kbts-blob"

pattern KTX_KEY_kbts_version :: Text
pattern KTX_KEY_kbts_version = "kbts-version"

kbtsVersion :: Text
kbtsVersion = Text.pack $ show KBTS.VERSION_CURRENT <> "," <> show KBTS.BLOB_VERSION_CURRENT

bundleFont
  :: FilePath -- ^ Source font (ttf etc)
  -> FilePath -- ^ Atlas JSON
  -> FilePath -- ^ KTX2 texture
  -> FilePath -- ^ Output
  -> IO ()
bundleFont pathTtf pathJson pathKtx2 pathKtxf = do
  layout <- eitherDecodeFileStrict pathJson >>= either fail pure

  ttfData <- ByteString.readFile pathTtf
  kbtsData <- Font.extractBlob ttfData 0

  font <- createFont kbtsData 0
  rescale <-
    case layout.atlas.sizeUnit of
      Just Layout.CapsHeight -> pure id
      Nothing -> do
        factor <- withFontData font $ pure . Font.emToCaps
        pure $ scalePlanes factor
  destroyFont font

  sourceKtx <- Ktx2.fromFile pathKtx2
  let
    atlas = rescale $ compact layout
    atlasData = ByteString.toStrict $ encode atlas
    kvd =
      KVD.insertBytes KTX_KEY_atlas (Zstd.compress 19 atlasData) $
      KVD.insertBytes KTX_KEY_kbts (Zstd.compress 19 kbtsData) $
      KVD.insertText KTX_KEY_kbts_version kbtsVersion $
      KVD.setWriterWith ("ktx-font 0.1.0.1 / " <>) sourceKtx.kvd
  Ktx2.toFile pathKtxf sourceKtx{Ktx2.kvd}

-- | Rescale planes to caps height instead of ems.
--
-- With 1.4 factor, 1em targeting 10px would become 14px.
-- That means you need smaller font sizes in caps that in ems.
scalePlanes :: Float -> Compact -> Compact
scalePlanes s a = a
  { Compact._size = a._size * s
  , Compact.planes = Vector.map (Compact.scaleBox s) a.planes
  }

{- | Put already-loaded font to context.

The context will NOT keep a fontData reference.
-}
pushBundleFont :: TextShape.Context -> Bundle -> IO Int
pushBundleFont ctx Bundle{fontData} =
  withFontData fontData $ TextShape.pushFont ctx

data Bundle = Bundle
  { fontData :: FontData -- ^ The bundle holds the font memory too.
  , atlas :: Compact -- ^ Glyph atlas data for the texture.
  }
  deriving (Eq)

loadBundleFile :: FilePath -> IO Bundle
loadBundleFile fontFile =
  bracket (Ktx2.open fontFile) Ktx2.close loadBundle

loadBundleBytes :: ByteString -> IO Bundle
loadBundleBytes bytes =
  Ktx2.bytes bytes >>= loadBundle

-- | Load bundle data from already opened KTX2 reader context
loadBundle :: Ktx2.ReadChunk a => Ktx2.Context a -> IO Bundle
loadBundle ktx = do
  kvd <- Ktx2.keyValueData ktx
  let
    bundled = (,,)
      <$> (Map.lookup KTX_KEY_atlas kvd >>= KVD.fromValue)
      <*> (Map.lookup KTX_KEY_kbts kvd >>= KVD.fromValue)
      <*> (Map.lookup KTX_KEY_kbts_version kvd >>= KVD.fromValue)

  (mtsdfZstd, kbtsZstd, kbtsVer) <-
    case bundled of
      Nothing ->
        error "Missing KTXF metadata"
      Just compressed ->
        pure compressed

  unless (kbtsVer == kbtsVersion) $
    error $ "KBTS blob version mismatch: " <> show kbtsVer <> ", expected " <> show kbtsVersion

  atlas <-
    case Zstd.decompress mtsdfZstd of
      Zstd.Decompress json ->
        either error pure $ eitherDecodeStrict json
      Zstd.Skip ->
        error "Empty atlas data"
      Zstd.Error err ->
        error $ "Atlas decompression error: " <> err

  fontData <-
    case Zstd.decompress kbtsZstd of
      Zstd.Decompress fontData ->
        createFont fontData 0
      Zstd.Skip ->
        error $ "Empty font data"
      Zstd.Error err ->
        error $ "Font decompression error: " <> err

  pure Bundle{..}

freeBundle :: Bundle -> IO ()
freeBundle Bundle{fontData} = destroyFont fontData

-- * Shaping contexts

data StackContext a = StackContext
  { shapeContext :: MVar TextShape.Context
  , bundled :: IntMap a
  , atlases :: IntMap Compact -- XXX: Font.Handle ~ Ptr Font.Handle ~ Int
  }
  deriving stock (Functor, Foldable, Traversable, Generic, Generic1)

{- | Create shaping context and push fonts from all the bundles.
-}
createStackContext :: Foldable t => t Bundle -> IO (StackContext ())
createStackContext bundles = do
  ctx <- TextShape.createContext
  locals <- for (toList bundles) \Bundle{..} ->
    withFontData fontData \font -> do
      _refs <- TextShape.pushFont ctx font
      pure (intHandle font, atlas)
  let atlases = IntMap.fromList locals
  let bundled = IntMap.map (const ()) atlases
  shapeContext <- newMVar ctx
  pure StackContext{..}

{- | Update font annotations with the bundle collection annotated with update functions.

This may be used to attach more information, e.g. after all fonts were assigned texture slots.

The update collection should be a superset of what was initially bundled.
Otherwise you may see missing things downstream.
-}
mapWithBundle :: Foldable t => t (Bundle, Maybe a -> b) -> StackContext a -> StackContext b
mapWithBundle bundles StackContext{bundled = old, ..} = StackContext{bundled = new, ..}
  where
    new = IntMap.fromList do
      (Bundle{fontData=Font.FontData{fontData}}, f) <- toList bundles
      let key = intHandle $ unsafeForeignPtrToPtr fontData
      pure (key, f $ IntMap.lookup key old)

{- | Destroy shaping context.

NB: Does NOT free the fonts as they may be shared with other stacks. Use `freeBundle` for that.
-}
destroyStackContext :: StackContext a -> IO ()
destroyStackContext StackContext{shapeContext} = takeMVar shapeContext >>= TextShape.destroyContext

{-# INLINE lookupBundled #-}
lookupBundled :: Font -> StackContext a -> Maybe a
lookupBundled font StackContext{..} = IntMap.lookup (intHandle font) bundled

{-# INLINE lookupAtlas #-}
lookupAtlas :: Font -> StackContext a -> Maybe Compact
lookupAtlas font StackContext{atlases} = IntMap.lookup (intHandle font) atlases

mapBundled :: (a -> b) -> StackContext a -> StackContext b
mapBundled f ctx@StackContext{bundled} = ctx
  { bundled = IntMap.map f bundled
  }