{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
module Sixel
    ( PaletteColor
    , RGBColor
    , render
    ) where

import Data.Bits (shiftL)
import Data.Bool (bool)
import Data.ByteString (ByteString)
import Data.Word (Word8)
import Pager.Types (Geometry(..))
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Char8 as ByteString.Char8
import qualified Data.Char as Char
import qualified Data.List as List
import qualified Data.List.Split as List (chunksOf)
import qualified Foreign.C.Types
import qualified Numeric.Probability.Trace as List (zipListWith)
import qualified Test.Speculate.Utils as Tuple (uncurry6)


showByteString :: Show a => a -> ByteString
showByteString =
    ByteString.Char8.pack . show


type Bit = Foreign.C.Types.CBool

type Sixel = (Bit, Bit, Bit, Bit, Bit, Bit)

type PaletteColor = Word8

type RGBColor = (Word8, Word8, Word8)


render :: Geometry -> [RGBColor] -> [PaletteColor] -> ByteString
render (Geometry x y width _) rgbColors canvas =
    "\ESCP0;0;q" <> colors <> sixeldata <> "\ESC\\"
  where
    bandsToSkip = floor (fromIntegral y / (6 :: Double))

    paddedCanvas =
      if y < 0 then
        drop (-(fromIntegral y) * fromIntegral width) canvas
      else
        let
          paddingTop = replicate (fromIntegral offsetY * fromIntegral width) 0
          offsetY = y - bandsToSkip * 6
        in
        paddingTop <> canvas

    palette = map fromIntegral [0..length rgbColors - 1]
    colors = mconcat $ map (uncurry setColorMapRegister) (zip palette rgbColors)
    channels = splitChannels palette paddedCanvas
    bitmaps = map (toScanlines width) channels
    images = map (map (runLengthEncode . toByteString) . toSixels) bitmaps
    sixeldata = skipBands (fromIntegral bandsToSkip) <> overstrikeBands (map (map (shiftX x)) images)


skipBands :: Int -> ByteString
skipBands n =
    ByteString.replicate n newline
  where
    newline = fromIntegral (Char.ord '-')


shiftX :: Int -> ByteString -> ByteString
shiftX x s =
    "!" <> showByteString x <> "?" <> s


setColorMapRegister :: PaletteColor -> RGBColor -> ByteString
setColorMapRegister i _rgbColor256@(r256,g256,b256) =
    "#" <> mconcat (List.intersperse ";" (map showByteString [i, 2, r100, g100, b100]))
  where
    (r100,g100,b100) = (to100 r256, to100 g256, to100 b256)
    to100 = round . (*(100/256 :: Double)) . fromIntegral


useColorMapRegister :: PaletteColor -> ByteString -> ByteString
useColorMapRegister color s =
    "#" <> showByteString color <> s


splitChannels :: [PaletteColor] -> [PaletteColor] -> [[Bit]]
splitChannels channels canvas =
    map (flip getChannel canvas) channels
  where
    getChannel :: PaletteColor -> [PaletteColor] -> [Bit]
    getChannel color = map (bool 0 1 . (==color))


toScanlines :: Int -> [Bit] -> [[Bit]]
toScanlines width =
    List.chunksOf (fromIntegral width)


-- Turn scanlines into sixelbands.
-- Empty lines will be added as necessary to construct the final band.
toSixels :: [[Bit]] -> [[Sixel]]
toSixels =
    map (Tuple.uncurry6 List.zip6) . rec
  where
    rec :: [[Bit]] -> [([Bit],[Bit],[Bit],[Bit],[Bit],[Bit])]
    rec (a:b:c:d:e:f:rest) = (a,b,c,d,e,f) : rec rest
    rec (a:b:c:d:e:[])     = (a,b,c,d,e,z) : []
    rec (a:b:c:d:[])       = (a,b,c,d,z,z) : []
    rec (a:b:c:[])         = (a,b,c,z,z,z) : []
    rec (a:b:[])           = (a,b,z,z,z,z) : []
    rec (a:[])             = (a,z,z,z,z,z) : []
    rec ([])               = []
    z = repeat 0


toByteString :: [Sixel] -> ByteString
toByteString =
    ByteString.pack . map ((+63) . toWord8)


toWord8 :: Sixel -> Word8
toWord8 (a, b, c, d, e, f) =
  shiftL (fromIntegral a) 0 +
  shiftL (fromIntegral b) 1 +
  shiftL (fromIntegral c) 2 +
  shiftL (fromIntegral d) 3 +
  shiftL (fromIntegral e) 4 +
  shiftL (fromIntegral f) 5


runLengthEncode :: ByteString -> ByteString
runLengthEncode input =
    case ByteString.uncons input of
      Just (c, input') ->
        let
          (c_last, n_last, out) = ByteString.foldl f (c, 1, "") input'
        in
        encode c_last n_last out

      Nothing ->
        input
  where
    f :: (Word8, Int, ByteString) -> Word8 -> (Word8, Int, ByteString)
    f (c, n, out) c_in =
      if c_in == c then
        (c, n + 1, out)
      else
        (c_in, 1, encode c n out)

    encode :: Word8 -> Int -> ByteString -> ByteString
    encode c n output =
        if n > 3 then
          output <> "!" <> showByteString n <> ByteString.singleton c
        else
          output <> ByteString.replicate n c


overstrikeBands :: [[ByteString]] -> ByteString
overstrikeBands =
    mconcat . List.intersperse "-" . List.zipListWith overstrikeBand


overstrikeBand :: [ByteString] -> ByteString
overstrikeBand =
    mconcat . List.intersperse "$" . map (uncurry useColorMapRegister) . zip [0..]