{- | Kitty graphics protocol encoder.

Serializes an RGBA image into the APC escape sequences the Kitty graphics
protocol uses (supported by Kitty, WezTerm, and Ghostty). Raw 32-bit RGBA is
transmitted (@f=32@) and displayed at the cursor (@a=T@); the base64 payload is
split into <=4096-byte chunks framed with @m=1@/@m=0@ per the protocol.

This module is pure: it produces the exact bytes to write to the terminal. It
performs no IO and does not know about vty — the caller decides where the
cursor is and when to emit.
-}
module Termtab.Graphics.Kitty (
    encodeImage,
    encodeRGBA,
    encodeVirtual,
    encodeDirect,
    deleteImage,
) where

import Codec.Picture (Image (..), PixelRGBA8, imageData)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Base64 qualified as B64
import Data.ByteString.Char8 qualified as BC
import Data.ByteString.Internal qualified as BSI
import Data.Vector.Storable qualified as VS
import Data.Word (Word8)
import Foreign.ForeignPtr (castForeignPtr)

{- | Encode an image for direct placement at the cursor (@a=T@) — used by the
one-shot demo.
-}
encodeImage :: Image PixelRGBA8 -> ByteString
encodeImage :: Image PixelRGBA8 -> ByteString
encodeImage Image PixelRGBA8
img =
    Int -> Int -> ByteString -> ByteString
encodeRGBA (Image PixelRGBA8 -> Int
forall a. Image a -> Int
imageWidth Image PixelRGBA8
img) (Image PixelRGBA8 -> Int
forall a. Image a -> Int
imageHeight Image PixelRGBA8
img) (Vector Word8 -> ByteString
vectorToBS (Image PixelRGBA8 -> Vector (PixelBaseComponent PixelRGBA8)
forall a. Image a -> Vector (PixelBaseComponent a)
imageData Image PixelRGBA8
img))

{- | Encode raw row-major RGBA bytes (@w*h*4@ bytes) for direct placement at the
cursor.
-}
encodeRGBA :: Int -> Int -> ByteString -> ByteString
encodeRGBA :: Int -> Int -> ByteString -> ByteString
encodeRGBA Int
w Int
h ByteString
rgba = ByteString -> ByteString -> ByteString
frameChunks (Int -> Int -> ByteString
baseCtrl Int
w Int
h) ByteString
rgba

{- | Encode an image as a *virtual* placement (@U=1@) for display via Unicode
placeholders: transmit it under image id @imgId@ occupying @cols@×@rows@ terminal
cells. The image is not shown until placeholder cells referencing @imgId@ are
drawn (see "Termtab.Graphics.KittyPlacement"). @q=2@ suppresses the terminal's
acknowledgement so it cannot corrupt input.
-}
encodeVirtual :: Int -> Int -> Int -> Image PixelRGBA8 -> ByteString
encodeVirtual :: Int -> Int -> Int -> Image PixelRGBA8 -> ByteString
encodeVirtual Int
imgId Int
cols Int
rows Image PixelRGBA8
img =
    ByteString -> ByteString -> ByteString
frameChunks (Int -> Int -> ByteString
baseCtrl (Image PixelRGBA8 -> Int
forall a. Image a -> Int
imageWidth Image PixelRGBA8
img) (Image PixelRGBA8 -> Int
forall a. Image a -> Int
imageHeight Image PixelRGBA8
img) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
virtualKeys) ByteString
rgba
  where
    rgba :: ByteString
rgba = Vector Word8 -> ByteString
vectorToBS (Image PixelRGBA8 -> Vector (PixelBaseComponent PixelRGBA8)
forall a. Image a -> Vector (PixelBaseComponent a)
imageData Image PixelRGBA8
img)
    virtualKeys :: ByteString
virtualKeys =
        String -> ByteString
BC.pack (String
",U=1,i=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
imgId String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
",c=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
cols String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
",r=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
rows String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
",q=2")

{- | Encode an image for direct placement at the cursor, scaled to @cols@×@rows@
cells, under image id @imgId@ / placement id 1. Re-emitting with the same ids
replaces the placement (no buildup). @C=1@ leaves the cursor where it is; @q=2@
suppresses the acknowledgement. Used to blit the notation over a reserved region
whose screen position the caller has moved the cursor to.
-}
encodeDirect :: Int -> Int -> Int -> Image PixelRGBA8 -> ByteString
encodeDirect :: Int -> Int -> Int -> Image PixelRGBA8 -> ByteString
encodeDirect Int
imgId Int
cols Int
rows Image PixelRGBA8
img =
    ByteString -> ByteString -> ByteString
frameChunks (Int -> Int -> ByteString
baseCtrl (Image PixelRGBA8 -> Int
forall a. Image a -> Int
imageWidth Image PixelRGBA8
img) (Image PixelRGBA8 -> Int
forall a. Image a -> Int
imageHeight Image PixelRGBA8
img) ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
keys) (Vector Word8 -> ByteString
vectorToBS (Image PixelRGBA8 -> Vector (PixelBaseComponent PixelRGBA8)
forall a. Image a -> Vector (PixelBaseComponent a)
imageData Image PixelRGBA8
img))
  where
    keys :: ByteString
keys =
        String -> ByteString
BC.pack
            (String
",i=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
imgId String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
",p=1,c=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
cols String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
",r=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
rows String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
",C=1,q=2")

{- | The control keys shared by every transmission: direct placement
(@a=T@), 32-bit RGBA (@f=32@), and the pixel dimensions.
-}
baseCtrl :: Int -> Int -> ByteString
baseCtrl :: Int -> Int -> ByteString
baseCtrl Int
w Int
h = String -> ByteString
BC.pack (String
"a=T,f=32,s=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
w String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
",v=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
h)

{- | Chunk the base64 payload into <=4096-byte APC frames with the given control
data on the first frame.
-}
frameChunks :: ByteString -> ByteString -> ByteString
frameChunks :: ByteString -> ByteString -> ByteString
frameChunks ByteString
ctrl ByteString
rgba =
    case Int -> ByteString -> [ByteString]
chunksOf Int
maxChunk (ByteString -> ByteString
B64.encode ByteString
rgba) of
        [] -> ByteString
BS.empty
        [ByteString
only] -> ByteString -> ByteString -> ByteString
frame (ByteString
ctrl ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
BC.pack String
",m=0") ByteString
only
        (ByteString
c0 : [ByteString]
rest) -> [ByteString] -> ByteString
BS.concat (ByteString -> ByteString -> ByteString
frame (ByteString
ctrl ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
BC.pack String
",m=1") ByteString
c0 ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString] -> [ByteString]
go [ByteString]
rest)
  where
    -- Kitty caps each transmission chunk at 4096 bytes of base64 payload.
    maxChunk :: Int
maxChunk = Int
4096
    go :: [ByteString] -> [ByteString]
go [] = []
    go [ByteString
lastC] = [ByteString -> ByteString -> ByteString
frame (String -> ByteString
BC.pack String
"m=0") ByteString
lastC]
    go (ByteString
c : [ByteString]
cs) = ByteString -> ByteString -> ByteString
frame (String -> ByteString
BC.pack String
"m=1") ByteString
c ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString] -> [ByteString]
go [ByteString]
cs

{- | Delete image @imgId@ (data and all its placements). Emitted before a
re-blit so Kitty actually refreshes the pixels instead of assuming an unchanged
placement is still current.
-}
deleteImage :: Int -> ByteString
deleteImage :: Int -> ByteString
deleteImage Int
imgId =
    [ByteString] -> ByteString
BS.concat [ByteString
apcStart, String -> ByteString
BC.pack (String
"a=d,d=I,i=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
imgId String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
",q=2"), ByteString
apcEnd]

-- | Wrap control data + payload in an @ESC _G <ctrl> ; <payload> ESC \\@ frame.
frame :: ByteString -> ByteString -> ByteString
frame :: ByteString -> ByteString -> ByteString
frame ByteString
ctrl ByteString
payload =
    [ByteString] -> ByteString
BS.concat [ByteString
apcStart, ByteString
ctrl, Char -> ByteString
BC.singleton Char
';', ByteString
payload, ByteString
apcEnd]

{- | The APC (Application Program Command) opening sequence Kitty graphics
commands are framed in: @ESC _G@.
-}
apcStart :: ByteString
apcStart :: ByteString
apcStart = [Word8] -> ByteString
BS.pack [Word8
0x1b] ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
BC.pack String
"_G"

-- | The APC closing (string terminator) sequence: @ESC \\@.
apcEnd :: ByteString
apcEnd :: ByteString
apcEnd = [Word8] -> ByteString
BS.pack [Word8
0x1b] ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> String -> ByteString
BC.pack String
"\\"

-- | Split a 'ByteString' into chunks of at most @n@ bytes.
chunksOf :: Int -> ByteString -> [ByteString]
chunksOf :: Int -> ByteString -> [ByteString]
chunksOf Int
n ByteString
bs
    | ByteString -> Bool
BS.null ByteString
bs = []
    | Bool
otherwise = let (ByteString
h, ByteString
t) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
n ByteString
bs in ByteString
h ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: Int -> ByteString -> [ByteString]
chunksOf Int
n ByteString
t

-- | Zero-copy view of a storable 'Word8' vector as a 'ByteString'.
vectorToBS :: VS.Vector Word8 -> ByteString
vectorToBS :: Vector Word8 -> ByteString
vectorToBS Vector Word8
v =
    let (ForeignPtr Word8
fptr, Int
len) = Vector Word8 -> (ForeignPtr Word8, Int)
forall a. Vector a -> (ForeignPtr a, Int)
VS.unsafeToForeignPtr0 Vector Word8
v
     in ForeignPtr Word8 -> Int -> Int -> ByteString
BSI.fromForeignPtr (ForeignPtr Word8 -> ForeignPtr Word8
forall a b. ForeignPtr a -> ForeignPtr b
castForeignPtr ForeignPtr Word8
fptr) Int
0 Int
len