{-# 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)

-- TODO rename to RegisteredColor, IndexedColor?
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


-- TODO what's the correct name?
-- TODO reword: channels :: [BitArray Int]
-- TODO reword:  XXX channels must be sorted by color index (to compute sixeldata)
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)


-- TODO maybe use BitArray
-- 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


-- #gitginore TODO use Vector.imap instead of map zip [0..]?
overstrikeBand :: [ByteString] -> ByteString
overstrikeBand =
    mconcat . List.intersperse "$" . map (uncurry useColorMapRegister) . zip [0..]