content: {-# LANGUAGE OverloadedStrings #-}
  
  module Jet.Render where
  
  import Control.Lens
  import Control.Monad.State
  import qualified Data.List as List
  import Data.Maybe
  import qualified Graphics.Vty as Vty
  import Prettyprinter
  import Prettyprinter.Render.Util.StackMachine (renderSimplyDecoratedA)
  
  data Cursor = Cursor
    deriving (Show, Eq)
  
  data Output = LineBreak | ImgChunk Vty.Image
  
  renderScreen :: Int -> SimpleDocStream (Either Cursor Vty.Attr) -> [Vty.Image]
  renderScreen winHeight doc = do
    let (outputs, (_, _, mcursor)) =
          doc
            & toLineStream
            & flip runState ([], 0, Nothing)
  
    let allLines =
          outputs & List.foldr collapse ([], [])
            -- clean up anything remaining in the buffer
            & ( \(buf, rest) ->
                  if null buf
                    then rest
                    else Vty.horizCat buf : rest
              )
  
    let cropped = case mcursor of
          Nothing -> allLines
          Just cursorPos -> takeSurroundingCursor winHeight cursorPos allLines
    cropped
    where
      collapse out (buf, rest) =
        case out of
          LineBreak -> ([], Vty.horizCat buf : rest)
          ImgChunk img -> (img : buf, rest)
  
  -- Take only enough lines to fill the screen, with the cursor centered don't print any more.
  takeSurroundingCursor :: Int -> Int -> [a] -> [a]
  takeSurroundingCursor height cursorPos xs
    | cursorPos - half > 0 =
      xs
        & drop (cursorPos - half)
        & take height
    | otherwise = take height xs
    where
      half = height `div` 2
  
  toLineStream ::
    SimpleDocStream (Either Cursor Vty.Attr) ->
    State ([Vty.Attr], Int, Maybe Int) [Output]
  toLineStream doc =
    renderSimplyDecoratedA
      renderText
      pushAnn
      popAnn
      doc
    where
      popAnn _ = do
        _1 %= drop 1
        -- Need to clear existing colors or they bleed to the right.
        pure [ImgChunk (Vty.text' Vty.defAttr "")]
      pushAnn = \case
        Left Cursor -> do
          cursorLine <- use _2
          _3 ?= cursorLine
          pure []
        Right ann -> do
          _1 %= (ann :)
          pure mempty
      -- prettyprinter always renders lines as a single text fragment
      renderText "\n" = do
        _2 += 1
        pure [LineBreak]
      -- prettyprinter never passes text with newlines here
      renderText txt = do
        attr <- uses _1 (fromMaybe Vty.defAttr . listToMaybe)
        pure [ImgChunk (Vty.text' attr txt)]
  