{-# LANGUAGE ApplicativeDo #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Jet (run) where

import Control.Category ((>>>))
import Control.Comonad (extract)
import qualified Control.Comonad as Comonad
import Control.Comonad.Cofree
import qualified Control.Comonad.Trans.Cofree as CofreeF
import Control.Lens hiding ((:<))
import qualified Control.Lens.Cons as Cons
import Control.Monad.State
import Control.Monad.Trans.Maybe (MaybeT (MaybeT, runMaybeT))
import Data.Aeson (Value)
import qualified Data.Aeson as Aeson
import Data.Aeson.Encode.Pretty (encodePretty)
import Data.Aeson.Extra
import qualified Data.ByteString.Lazy.Char8 as BS
import qualified Data.ByteString.Lazy.Char8 as LBS
import Data.Functor.Classes (Eq1 (..), Ord1 (liftCompare))
import qualified Data.Functor.Foldable as FF
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import qualified Data.HashMap.Strict as HashMap
import Data.Hashable (Hashable)
import qualified Data.List as List
import Data.Maybe
import Data.Sequence (Seq)
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Text.Zipper as TZ
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import qualified Graphics.Vty as Vty
import Graphics.Vty.Input.Events
import qualified Jet.Render as Render
import Prettyprinter as P
import System.Environment (getArgs)
import System.Exit (exitFailure)
import System.Hclip
import System.IO (IOMode (ReadWriteMode), openFile)
import qualified System.IO as IO
import qualified System.Posix as Posix
import Text.Read (readMaybe)
import qualified Zipper.Recursive as Z

tabSize :: Int
tabSize = 2

maxUndoStates :: Int
maxUndoStates = 100

hoistMaybe :: Maybe a -> MaybeT Editor a
hoistMaybe = MaybeT . pure

data EditorState = EditorState
  { _undo :: UndoZipper (Z.Zipper JIndex ValueF FocusState),
    _mode :: Mode,
    _register :: ValueF (Cofree ValueF FocusState),
    _vty :: Vty.Vty,
    _flash :: Text,
    _save :: Z.Zipper JIndex ValueF FocusState -> Editor ()
  }

newtype Editor a = Editor {runEditor :: StateT EditorState IO a}
  deriving newtype (Functor, Applicative, Monad, MonadState EditorState, MonadIO)

mode_ :: Lens' EditorState Mode
mode_ = lens _mode (\s m -> s {_mode = m})

register_ :: Lens' EditorState (ValueF (Cofree ValueF FocusState))
register_ = lens _register (\s m -> s {_register = m})

undo_ :: Lens' EditorState (UndoZipper (Z.Zipper JIndex ValueF FocusState))
undo_ = lens _undo (\s m -> s {_undo = m})

vty_ :: Lens' EditorState Vty.Vty
vty_ = lens _vty (\s m -> s {_vty = m})

flash_ :: Lens' EditorState Text
flash_ = lens _flash (\s m -> s {_flash = m})

save_ :: Lens' EditorState (Z.Zipper JIndex ValueF FocusState -> Editor ())
save_ = lens _save (\s m -> s {_save = m})

recover :: a -> MaybeT Editor a -> Editor a
recover def m = do
  let e = runMaybeT m
  s <- get
  r <- liftIO $ flip runStateT s . runEditor $ e
  case r of
    (Just a, newS) -> put newS *> pure a
    (Nothing, _) -> pure def

data Focused = Focused | NotFocused
  deriving (Eq)

data Folded = Folded | NotFolded
  deriving (Eq)

type PrettyJSON = Doc (Either Render.Cursor Vty.Attr)

type Buffer = TZ.TextZipper Text

-- | Nodes are annotated with one of these.
-- This includes information about the node itself, but also
-- a cached render of the node, which allows us to re-render
-- the whole tree much faster.
data FocusState = FocusState
  { isFocused :: Focused,
    isFolded :: Folded,
    rendered :: PrettyJSON
  }

instance Eq FocusState where
  a == b =
    isFocused a == isFocused b
      && isFolded a == isFolded b

focused_ :: Lens' FocusState Focused
focused_ = lens isFocused (\fs new -> fs {isFocused = new})

folded_ :: Lens' FocusState Folded
folded_ = lens isFolded (\fs new -> fs {isFolded = new})

toggleFold :: Folded -> Folded
toggleFold Folded = NotFolded
toggleFold NotFolded = Folded

run :: IO ()
run = do
  (json, srcFile) <-
    getArgs >>= \case
      [] -> do
        json <-
          (Aeson.eitherDecode . BS.pack <$> getContents) >>= \case
            Left err -> do
              IO.hPutStrLn IO.stderr err
              exitFailure
            Right json -> pure json
        pure (json, Nothing)
      [f] -> do
        json <-
          Aeson.eitherDecodeFileStrict f >>= \case
            Left err -> do
              IO.hPutStrLn IO.stderr err
              exitFailure
            Right json -> pure json
        pure (json, Just f)
      _ -> IO.hPutStrLn IO.stderr "usage: structural-json FILE.json" *> exitFailure
  result <- edit srcFile $ json
  BS.putStrLn $ encodePretty result

edit :: Maybe FilePath -> Value -> IO Value
edit srcFile value = do
  -- Use tty so we don't interfere with stdin/stdout
  tty <- openFile "/dev/tty" ReadWriteMode >>= Posix.handleToFd
  config <- liftIO $ Vty.standardIOConfig
  vty <- (liftIO $ Vty.mkVty config {Vty.inputFd = Just tty, Vty.outputFd = Just tty})
  -- load the value into a zipper.
  let z = Z.zipper . toCofree $ value
  v <- flip evalStateT (editorState srcFile vty) . runEditor $ loop z
  Vty.shutdown vty
  pure (Z.flatten v)

loop ::
  Z.Zipper JIndex ValueF FocusState ->
  Editor (Z.Zipper JIndex ValueF FocusState)
loop z = do
  vty <- use vty_
  renderScreen z
  flash_ .= ""
  e <- liftIO $ Vty.nextEvent vty
  nextZ <- handleEvent e z
  if (shouldExit e)
    then pure nextZ
    else (loop nextZ)

renderScreen :: Z.Zipper JIndex ValueF FocusState -> Editor ()
renderScreen z = do
  vty <- use vty_
  (winWidth, winHeight) <- bounds
  rendered <- uses mode_ (\m -> fullRender m z)
  footer <- footerImg
  let screen = Vty.vertCat . Render.renderScreen (winHeight - Vty.imageHeight footer) . layoutPretty defaultLayoutOptions $ rendered
  let spacerHeight = winHeight - (Vty.imageHeight screen + Vty.imageHeight footer)
  let spacers = Vty.charFill Vty.defAttr ' ' winWidth spacerHeight
  liftIO $ Vty.update vty (Vty.picForImage (screen Vty.<-> spacers Vty.<-> footer))

-- | Get the current bounds of the current terminal screen.
bounds :: Editor (Int, Int)
bounds = use vty_ >>= liftIO . Vty.displayBounds . Vty.outputIface

-- | Render the footer bar to an image
footerImg :: Editor Vty.Image
footerImg = do
  (w, _) <- bounds
  flash <- gets _flash
  let attr = (Vty.defAttr `Vty.withForeColor` Vty.green `Vty.withStyle` Vty.reverseVideo)
      helpMsg = Vty.text' attr "| Press '?' for help"
      flashMsg = Vty.text' (attr `Vty.withStyle` Vty.bold) (" " <> flash)
  pure $
    Vty.horizCat
      [ flashMsg,
        Vty.charFill attr ' ' (w - (Vty.imageWidth helpMsg + Vty.imageWidth flashMsg)) 1,
        helpMsg
      ]

-- | Push the given zipper onto history iff it's distinct from the most recent undo state.
pushUndo :: Z.Zipper JIndex ValueF FocusState -> Editor ()
pushUndo z =
  undo_ %= \case
    (UndoZipper (ls Cons.:> _) _) | length ls >= maxUndoStates -> UndoZipper (z <| ls) Empty
    (UndoZipper ls _) -> UndoZipper (z <| ls) Empty

editorState :: Maybe FilePath -> Vty.Vty -> EditorState
editorState srcFile vty =
  EditorState
    { _undo = UndoZipper Empty Empty,
      _mode = Move,
      _register = NullF,
      _vty = vty,
      _flash = "Hello World",
      _save = saveFile
    }
  where
    saveFile = case srcFile of
      Nothing -> const (pure ())
      Just fp -> \z -> do
        liftIO $ BS.writeFile fp (z & Z.flatten & encodePretty @Value)
        flash_ .= "Saved to " <> Text.pack fp

shouldExit :: Vty.Event -> Bool
shouldExit = \case
  EvKey (KChar 'c') [Vty.MCtrl] -> True
  EvKey (KChar 'q') [] -> True
  _ -> False

bufferText :: Buffer -> Text
bufferText = Text.concat . TZ.getText

-- | Apply the state that's in the current mode's buffer to the selected node if possible.
applyBuf :: Z.Zipper JIndex ValueF FocusState -> Editor (Z.Zipper JIndex ValueF FocusState)
applyBuf z = do
  use mode_ >>= \case
    Edit buf -> do
      let txt = buf ^. to bufferText
      mode_ .= Move
      pure
        ( z & Z.unwrapped_ . _unwrap
            %~ ( \case
                   StringF _ -> StringF txt
                   (NumberF n) -> NumberF . fromMaybe n . readMaybe $ Text.unpack txt
                   x -> x
               )
        )
    KeyEdit key buf -> do
      let txt = buf ^. to bufferText
      mode_ .= (KeyMove txt)
      pure
        ( z & Z.unwrapped_ . _unwrap
            %~ ( \case
                   (ObjectF hm) -> ObjectF $ renameKey key txt hm
                   x -> x
               )
        )
    _ -> pure z

renameKey :: (Hashable k, Eq k) => k -> k -> HashMap k v -> HashMap k v
renameKey srcKey destKey hm =
  hm
    &~ do
      v <- use (at srcKey)
      at srcKey .= Nothing
      at destKey .= v

-- | Create a buffer using the text from the current value.
bufferForValueF :: ValueF x -> Maybe Buffer
bufferForValueF = \case
  (ObjectF _hm) -> Nothing
  (ArrayF _vec) -> Nothing
  StringF txt -> Just $ newBuffer txt
  (NumberF sci) ->
    Just $ newBuffer (Text.pack . show $ sci)
  (BoolF True) -> Just $ newBuffer "true"
  (BoolF False) -> Just $ newBuffer "true"
  NullF -> Just $ newBuffer "null"

boolText_ :: Prism' Text Bool
boolText_ = prism' toText toBool
  where
    toText True = "true"
    toText False = "false"
    toBool "true" = Just True
    toBool "false" = Just False
    toBool _ = Nothing

data Mode
  = Edit {_buf :: Buffer}
  | Move
  | KeyMove {_selectedKey :: Text}
  | KeyEdit {_selectedKey :: Text, _buf :: Buffer}
  deriving (Show)

buf_ :: Traversal' Mode Buffer
buf_ f = \case
  Edit b -> Edit <$> f b
  Move -> pure Move
  KeyMove txt -> pure (KeyMove txt)
  KeyEdit txt b -> KeyEdit txt <$> f b

-- | Main event handler
handleEvent :: Vty.Event -> Z.Zipper JIndex ValueF FocusState -> Editor (Z.Zipper JIndex ValueF FocusState)
handleEvent evt zipper = do
  use mode_ >>= \case
    KeyMove {} -> handleMove zipper
    Move {} -> handleMove zipper
    KeyEdit {} -> handleEdit zipper
    Edit {} -> handleEdit zipper
  where
    handleEdit ::
      ( Z.Zipper JIndex ValueF FocusState ->
        Editor (Z.Zipper JIndex ValueF FocusState)
      )
    handleEdit z =
      case evt of
        EvKey key [] ->
          -- Perform buffer updates:
          case key of
            KChar c -> do
              mode_ . buf_ %= TZ.insertChar c
              pure z
            KLeft -> do
              mode_ . buf_ %= TZ.moveLeft
              pure z
            KRight -> do
              mode_ . buf_ %= TZ.moveRight
              pure z
            KBS -> do
              mode_ . buf_ %= TZ.deletePrevChar
              pure z
            KEsc -> do
              newZ <- applyBuf z
              pure $ newZ
            _ -> pure z
        _ -> pure z
    handleMove ::
      ( Z.Zipper JIndex ValueF FocusState ->
        Editor (Z.Zipper JIndex ValueF FocusState)
      )
    handleMove z =
      case evt of
        EvKey key mods -> case key of
          -- move up
          KChar 'h' -> z & outOf
          -- move down
          KChar 'l' -> do
            z & Z.focus_ . folded_ .~ NotFolded
              & into
          -- next sibling
          KChar 'j' -> z & sibling Forward
          -- move down
          KChar 'J' -> do
            pushUndo z
            pure (z & moveElement Forward)
          -- prev sibling
          KChar 'k' -> z & sibling Backward
          -- move up
          KChar 'K' -> do
            pushUndo z
            pure (z & moveElement Backward)
          -- add new node
          KChar 'i' -> do
            pushUndo z
            insert z
          -- replace with boolean
          KChar 'b' -> do
            pushUndo z
            pure (z & setFocus (BoolF True))
          -- replace with object
          KChar 'o' -> do
            pushUndo z
            pure (z & setFocus (ObjectF mempty))
          -- replace with array
          KChar 'a' -> do
            pushUndo z
            pure (z & setFocus (ArrayF mempty))
          -- replace with number
          KChar 'n' -> do
            pushUndo z
            pure (z & setFocus (NumberF 0))
          -- replace with Null
          KChar 'N' -> do
            pushUndo z
            pure (z & setFocus NullF)
          -- Save file
          KChar 's'
            | [Vty.MCtrl] <- mods -> do
              saver <- use save_
              saver z
              pure z
          -- replace with string
          KChar 's' -> do
            pushUndo z
            pure (z & setFocus (StringF ""))
          -- undo
          KChar 'u' -> do
            flash_ .= "Undo"
            undo_ %%= \case
              (UndoZipper (l Cons.:< ls) rs) ->
                (l, UndoZipper ls (z Cons.:< rs))
              lz -> (z, lz)
          -- redo
          KChar 'r' | [Vty.MCtrl] <- mods -> do
            flash_ .= "Redo"
            undo_ %%= \case
              (UndoZipper ls (r Cons.:< rs)) -> (r, UndoZipper (z Cons.:< ls) rs)
              lz -> (z, lz)
          -- toggle bool
          KChar ' ' -> do
            pushUndo z
            pure (z & tryToggleBool)
          -- copy
          KChar 'y' -> do
            flash_ .= "Copied"
            copy z
          -- paste
          KChar 'p' -> do
            flash_ .= "Paste"
            pushUndo z
            paste z
          -- cut
          KChar 'x' -> do
            flash_ .= "Cut"
            pushUndo z
            copy z >>= delete
          -- help
          KChar '?' -> do
            vty <- use vty_
            liftIO $ Vty.update vty (Vty.picForImage helpImg)
            void $ liftIO $ Vty.nextEvent vty
            pure z
          -- add child
          KEnter -> do
            pushUndo z
            tryAddChild z
          -- toggle fold
          KChar '\t' -> do
            -- Exit KeyMove mode if we're in it.
            mode_ .= Move
            pure $ (z & Z.focus_ . folded_ %~ toggleFold)
          -- Fold all children
          KChar 'F' -> do
            -- Fold all child branches
            pure $ mapChildren (mapped . folded_ .~ Folded) z
          -- unfold all children
          KChar 'f' -> do
            -- Unfold all child branches
            pure $ mapChildren (mapped . folded_ .~ NotFolded) z
          -- delete node
          KBS -> do
            flash_ .= "Deleted"
            pushUndo z
            delete z
          _ -> pure z
        _ -> pure z
    paste z = do
      reg <- use register_
      pure (z & setFocus reg)
    copy z = do
      let curVal = Z.branches z
      register_ .= curVal
      liftIO $ setClipboard (encodeValueFCofree curVal)
      pure z
    insert z = do
      use mode_ >>= \case
        KeyMove k -> do
          mode_ .= KeyEdit k (newBuffer k)
          pure $ z & Z.focus_ . folded_ .~ NotFolded
        Move
          | Just editBuf <- bufferForValueF (z ^. Z.branches_) -> do
            mode_ .= Edit editBuf
            pure $ z & Z.focus_ . folded_ .~ NotFolded
        _ -> pure z

encodeValueFCofree :: ValueF (Cofree ValueF FocusState) -> String
encodeValueFCofree vf = LBS.unpack . encodePretty . FF.embed $ fmap (FF.cata alg) vf
  where
    alg :: CofreeF.CofreeF ValueF ann Value -> Value
    alg (_ CofreeF.:< vf') = FF.embed vf'

-- | Set the value of the focused node.
setFocus ::
  ValueF (Cofree ValueF FocusState) ->
  Z.Zipper JIndex ValueF FocusState ->
  Z.Zipper JIndex ValueF FocusState
setFocus f z = z & Z.branches_ .~ f & rerender

data Dir = Forward | Backward

-- | Move the current value within an array
moveElement :: Dir -> Z.Zipper JIndex ValueF FocusState -> Z.Zipper JIndex ValueF FocusState
moveElement dir z = fromMaybe z $ do
  i <- case Z.currentIndex z of
    Just (Index i) -> pure i
    _ -> Nothing
  parent <- z & rerender & Z.up
  pure $
    case parent ^. Z.branches_ of
      ArrayF arr ->
        let swapI = case dir of
              Forward -> i + 1
              Backward -> i - 1
            moves =
              [ (i, arr Vector.!? swapI),
                (swapI, arr Vector.!? i)
              ]
                & sequenceOf (traversed . _2)
                & fromMaybe []
         in parent
              & Z.branches_ .~ ArrayF (arr Vector.// moves)
              & fromMaybe z . Z.down (Index swapI)
      _ -> z

tryToggleBool :: Z.Zipper JIndex ValueF FocusState -> Z.Zipper JIndex ValueF FocusState
tryToggleBool z =
  z & Z.branches_ %~ \case
    BoolF b -> BoolF (not b)
    x -> x

tryAddChild :: Z.Zipper JIndex ValueF FocusState -> Editor (Z.Zipper JIndex ValueF FocusState)
tryAddChild z =
  z & Z.branches_ %%~ \case
    ObjectF hm -> do
      mode_ .= (KeyEdit "" $ newBuffer "")
      pure $ ObjectF $ HM.insert "" (toCofree Aeson.Null) hm
    ArrayF arr -> do
      mode_ .= Move
      pure $ ArrayF $ arr <> pure (toCofree Aeson.Null)
    x -> pure x

-- | Delete the current node
delete :: Z.Zipper JIndex ValueF FocusState -> Editor (Z.Zipper JIndex ValueF FocusState)
delete z = do
  curMode <- use mode_
  mode_ .= Move
  pure $ case z ^. Z.branches_ of
    -- If we're in a Key focus, delete that key
    ObjectF hm
      | KeyMove k <- curMode ->
        ( z & Z.branches_ .~ ObjectF (HM.delete k hm)
        )
    -- Otherwise move up a layer and delete the key we were in.
    _ -> case Z.currentIndex z of
      -- If we don't have a parent, set the current node to null
      Nothing ->
        z & Z.branches_ .~ NullF
      Just i -> fromMaybe z $ do
        parent <- z & rerender & Z.up
        pure $
          parent & Z.branches_ %~ \case
            ObjectF hm | Key k <- i -> ObjectF (HM.delete k hm)
            ArrayF arr | Index j <- i -> ArrayF (Vector.ifilter (\i' _ -> i' /= j) arr)
            x -> x

-- | Move to next/previous sibling.
sibling :: Dir -> Z.Zipper JIndex ValueF FocusState -> Editor (Z.Zipper JIndex ValueF FocusState)
sibling dir z = recover z $ do
  mode <- use mode_
  case (mode, Z.branches z) of
    (KeyMove k, ObjectF hm) -> do
      case findSiblingIndex (== k) $ HashMap.keys hm of
        Nothing -> pure z
        Just theKey -> do
          mode_ .= KeyMove theKey
          pure z
    _ -> do
      curI <- hoistMaybe $ Z.currentIndex z
      parent <- hoistMaybe $ (z & rerender & Z.up)
      let newI = case Z.branches parent of
            ObjectF hm -> do
              let keys = HM.keys hm
              newKey <- findSiblingIndex (\k -> Key k == curI) keys
              pure $ Key newKey
            ArrayF xs -> case curI of
              (Index i) -> alterIndex xs i
              _ -> Nothing
            StringF {} -> Nothing
            NumberF {} -> Nothing
            BoolF {} -> Nothing
            NullF -> Nothing
      case newI of
        Just i -> hoistMaybe $ Z.down i parent
        Nothing -> hoistMaybe Nothing
  where
    (findSiblingIndex, alterIndex) = case dir of
      Forward ->
        ( findAfter,
          \xs i -> if i < length xs - 1 then Just (Index (i + 1)) else Nothing
        )
      Backward ->
        ( findBefore,
          \_xs i -> if i > 0 then Just (Index (i -1)) else Nothing
        )

findAfter :: (a -> Bool) -> [a] -> Maybe a
findAfter p xs = fmap snd . List.find (p . fst) $ zip xs (drop 1 xs)

findBefore :: (a -> Bool) -> [a] -> Maybe a
findBefore p xs = fmap snd . List.find (p . fst) $ zip (drop 1 xs) xs

newBuffer :: Text -> Buffer
newBuffer txt = TZ.gotoEOF $ TZ.textZipper (Text.lines txt) Nothing

-- | Move into the current node
into :: Z.Zipper JIndex ValueF FocusState -> Editor (Z.Zipper JIndex ValueF FocusState)
into z = do
  mode <- use mode_
  case (Z.branches z, mode) of
    (ObjectF _, KeyMove key) -> do
      mode_ .= Move
      pure (Z.tug (Z.down (Key key)) z)
    (ObjectF hm, Move) -> do
      case (HM.keys hm) ^? _head of
        Just fstKey -> do
          mode_ .= KeyMove fstKey
          pure z
        _ -> pure z
    (ArrayF {}, _) -> do
      pure $ Z.tug (Z.down (Index 0)) z
    _ -> pure z

-- | Move out of the current node
outOf :: Z.Zipper JIndex ValueF FocusState -> Editor (Z.Zipper JIndex ValueF FocusState)
outOf z = do
  mode <- use mode_
  maybeParentKey <- case (Z.currentIndex z) of
    Just (Key k) -> pure $ Just k
    _ -> pure Nothing

  case (Z.branches z, mode) of
    (ObjectF _, KeyMove {}) -> do
      mode_ .= Move
      pure z
    _ -> do
      maybe (pure ()) (\k -> mode_ .= KeyMove k) maybeParentKey
      pure (Z.tug (rerender >>> Z.up) z)

-- | Render the full zipper using render caches stored in each node.
fullRender :: Mode -> Z.Zipper JIndex ValueF FocusState -> PrettyJSON
fullRender mode z = do
  let focusedRender =
        z & Z.focus_ . focused_ .~ Focused
          & Z.unwrapped_ %~ \(fs :< vf) ->
            let rerendered = renderSubtree fs mode (rendered . extract <$> vf)
             in (fs {rendered = rerendered} :< vf)
  rendered . foldSpine alg $ focusedRender
  where
    alg fs vf =
      fs {rendered = rerenderCached fs (rendered <$> vf)}
    rerenderCached fs = \case
      ObjectF o -> prettyObj (isFocused fs) mode o
      ArrayF a -> prettyArray (isFocused fs) a
      -- Nodes without children are never part of the spine, but just to have something
      -- we can render the cache.
      _ -> rendered fs

-- | Updates the cached render of the current focus, using cached renders for subtrees.
rerender :: Z.Zipper JIndex ValueF FocusState -> Z.Zipper JIndex ValueF FocusState
rerender = Z.unwrapped_ %~ rerenderCofree

-- Rerenders a layer of a cofree structure. Doesn't re-render the children.
rerenderCofree :: Cofree ValueF FocusState -> Cofree ValueF FocusState
rerenderCofree (fs :< vf) =
  let rerendered = (renderSubtree fs mode (rendered . extract <$> vf))
   in fs {rendered = rerendered} :< vf
  where
    -- Currently the mode is required by renderSubtree, but for the rerender cache it's
    -- irrelevant, because it only matters if we're 'focused', and if we're focused, we'll be
    -- manually rerendered later anyways.
    mode = Move

-- | Renders a subtree
renderSubtree :: FocusState -> Mode -> ValueF PrettyJSON -> PrettyJSON
renderSubtree (FocusState {isFolded = Folded, isFocused}) _ vf = case vf of
  ObjectF {} -> colored' Vty.white "{...}"
  ArrayF {} -> colored' Vty.white "[...]"
  StringF {} -> colored' Vty.green "\"...\""
  NumberF {} -> colored' Vty.blue "..."
  NullF {} -> colored' Vty.yellow "..."
  BoolF {} -> colored' Vty.magenta "..."
  where
    colored' :: Vty.Color -> String -> PrettyJSON
    colored' col txt =
      P.annotate (Right $ if isFocused == Focused then reverseCol col else Vty.defAttr `Vty.withForeColor` col) (pretty txt)
renderSubtree (FocusState {isFocused}) mode vf = case vf of
  (StringF txt) -> cursor isFocused $ case (isFocused, mode) of
    (Focused, Edit buf) ->
      colored' Vty.green "\"" <> renderBuffer Vty.green buf <> colored' Vty.green "\""
    _ -> colored' Vty.green "\"" <> colored' Vty.green (Text.unpack txt) <> colored' Vty.green "\""
  (NullF) -> cursor isFocused $ colored' Vty.yellow "null"
  (NumberF n) -> cursor isFocused $ case (isFocused, mode) of
    (Focused, Edit buf) -> renderBuffer Vty.blue buf
    _ -> colored' Vty.blue (show n)
  (BoolF b) -> cursor isFocused $ colored' Vty.magenta (Text.unpack $ boolText_ # b)
  (ArrayF xs) -> prettyArray isFocused xs
  (ObjectF xs) -> prettyObj isFocused mode xs
  where
    colored' :: Vty.Color -> String -> PrettyJSON
    colored' col txt =
      P.annotate (Right $ if isFocused == Focused then reverseCol col else Vty.defAttr `Vty.withForeColor` col) (pretty txt)

-- | Attr in reverse-video
reverseCol :: Vty.Color -> Vty.Attr
reverseCol col = Vty.defAttr `Vty.withForeColor` col `Vty.withStyle` Vty.reverseVideo

-- | Map over all children of the current node, re-rendering after changes.
mapChildren ::
  (Cofree ValueF FocusState -> Cofree ValueF FocusState) ->
  Z.Zipper JIndex ValueF FocusState ->
  Z.Zipper JIndex ValueF FocusState
mapChildren f = Z.branches_ . mapped %~ FF.cata alg
  where
    alg :: CofreeF.CofreeF ValueF FocusState (Cofree ValueF FocusState) -> Cofree ValueF FocusState
    alg (cf CofreeF.:< vf) = rerenderCofree $ f (cf :< vf)

prettyWith :: Pretty a => Vty.Attr -> a -> PrettyJSON
prettyWith ann a = annotate (Right ann) $ pretty a

colored :: Pretty a => Vty.Color -> a -> PrettyJSON
colored col a = annotate (Right $ Vty.defAttr `Vty.withForeColor` col) $ pretty a

renderBuffer :: Vty.Color -> Buffer -> PrettyJSON
renderBuffer col buf =
  let (prefix, suffix) = Text.splitAt (snd $ TZ.cursorPosition buf) (bufferText buf)
      suffixImg = case Text.uncons suffix of
        Nothing -> prettyWith (reverseCol col) ' '
        Just (c, rest) -> prettyWith (reverseCol col) c <> colored col rest
   in colored col prefix <> suffixImg

cursor :: Focused -> PrettyJSON -> PrettyJSON
cursor Focused = P.annotate (Left Render.Cursor)
cursor _ = id

prettyArray :: Focused -> Vector PrettyJSON -> PrettyJSON
prettyArray foc vs =
  let inner :: [PrettyJSON] =
        Vector.toList vs
          & imap (\i v -> v <> commaKey i)
   in cursor foc $ vsep $ [img "[", indent tabSize (vsep inner), img "]"]
  where
    img :: Text -> PrettyJSON
    img t = case foc of
      Focused -> prettyWith (reverseCol Vty.white) t
      NotFocused -> pretty t
    commaKey i
      | i == Vector.length vs - 1 = mempty
      | otherwise = ","

prettyObj :: Focused -> Mode -> HashMap Text PrettyJSON -> PrettyJSON
prettyObj focused mode vs =
  let inner :: PrettyJSON
      inner =
        vsep
          ( HM.toList vs
              & imap
                ( \i (k, v) ->
                    vsep [imgForKey k <> pretty @Text ": ", indent tabSize (v <> commaKey i)]
                )
          )
      rendered = vsep [img "{", indent tabSize inner, img "}"]
   in case mode of
        Move -> cursor focused rendered
        _ -> rendered
  where
    hmSize = HM.size vs
    commaKey i
      | i == hmSize - 1 = mempty
      | otherwise = ","
    imgForKey k = case focused of
      NotFocused -> colored Vty.cyan (show k)
      Focused -> case mode of
        KeyMove focKey | focKey == k -> cursor Focused $ prettyWith (reverseCol Vty.cyan) (show focKey)
        KeyEdit focKey buf | focKey == k -> cursor Focused $ colored Vty.cyan '"' <> renderBuffer Vty.cyan buf <> colored Vty.cyan '"'
        _ -> colored Vty.cyan (show k)
    img :: Text -> PrettyJSON
    img t = case (focused, mode) of
      (Focused, Move) -> prettyWith (reverseCol Vty.white) t
      _ -> pretty t

-- Orphan instances
instance Eq1 ValueF where
  liftEq f vf1 vf2 = case (vf1, vf2) of
    (ObjectF l, ObjectF r) -> liftEq f l r
    (ArrayF l, ArrayF r) -> liftEq f l r
    (NullF, NullF) -> True
    (StringF l, StringF r) -> l == r
    (NumberF l, NumberF r) -> l == r
    (BoolF l, BoolF r) -> l == r
    _ -> False

instance Ord1 ValueF where
  liftCompare f vf1 vf2 = case (vf1, vf2) of
    (ObjectF l, ObjectF r) -> liftCompare f l r
    (ArrayF l, ArrayF r) -> liftCompare f l r
    (NullF, NullF) -> EQ
    (StringF l, StringF r) -> compare l r
    (NumberF l, NumberF r) -> compare l r
    (BoolF l, BoolF r) -> compare l r
    (NullF, _) -> LT
    (_, NullF) -> GT
    (BoolF _, _) -> LT
    (_, BoolF _) -> GT
    (NumberF _, _) -> LT
    (_, NumberF _) -> GT
    (StringF _, _) -> LT
    (_, StringF _) -> GT
    (ArrayF _, _) -> LT
    (_, ArrayF _) -> GT

data JIndex
  = Index Int
  | Key Text
  deriving (Show, Eq, Ord)

instance FunctorWithIndex JIndex ValueF

instance FoldableWithIndex JIndex ValueF

instance TraversableWithIndex JIndex ValueF where
  itraverse f = \case
    NullF -> pure NullF
    StringF txt -> pure (StringF txt)
    NumberF sci -> pure (NumberF sci)
    BoolF b -> pure (BoolF b)
    ObjectF hm -> ObjectF <$> itraverse (\k a -> f (Key k) a) hm
    ArrayF arr -> ArrayF <$> itraverse (\k a -> f (Index k) a) arr

type instance Index (ValueF a) = JIndex

type instance IxValue (ValueF a) = a

instance Ixed (ValueF a) where
  ix (Index i) f (ArrayF xs) = ArrayF <$> ix i f xs
  ix (Key k) f (ObjectF xs) = ObjectF <$> ix k f xs
  ix _ _ x = pure x

toCofree :: (Value -> Cofree ValueF FocusState)
toCofree t = FF.hylo alg FF.project $ t
  where
    defaultFs = FocusState NotFocused NotFolded mempty
    mode = Move
    alg :: ValueF (Cofree ValueF FocusState) -> Cofree ValueF FocusState
    alg vf = defaultFs {rendered = renderSubtree defaultFs mode (rendered . extract <$> vf)} :< vf

helpImg :: Vty.Image
helpImg =
  let helps =
        [ ("h", "ascend"),
          ("l", "descend"),
          ("j", "next sibling"),
          ("k", "previous sibling"),
          ("J", "move down (in array)"),
          ("K", "move up (in array)"),
          ("i", "enter edit mode (string/number)"),
          ("<C-s>", "save file"),
          ("<SPACE>", "toggle boolean"),
          ("<ESC>", "exit edit mode"),
          ("<BS>", "delete key/element"),
          ("<ENTER>", "add new key/element (object/array)"),
          ("<TAB>", "toggle fold"),
          ("f", "unfold all children"),
          ("F", "fold all children"),
          ("s", "replace element with string"),
          ("b", "replace element with bool"),
          ("n", "replace element with number"),
          ("N", "replace element with null"),
          ("a", "replace element with array"),
          ("o", "replace element with object"),
          ("u", "undo last change (undo buffer keeps 100 states)"),
          ("<C-r>", "redo from undo states"),
          ("y", "copy current value into buffer (and clipboard)"),
          ("p", "paste value from buffer over current value"),
          ("x", "cut a value, equivalent to a copy -> delete"),
          ("q | ctrl-c", "quit without saving. Due to a bug, tap twice")
        ]

      (keys, descs) =
        unzip
          ( helps <&> \(key, desc) ->
              ( Vty.text' (Vty.defAttr `Vty.withForeColor` Vty.green) (key <> ": "),
                Vty.text' Vty.defAttr desc
              )
          )
   in (Vty.vertCat keys Vty.<|> Vty.vertCat descs)

-- | Recomputes the spine at the current position, then at every position from that point
-- upwards until the zipper is closed, returning the result.
foldSpine :: (Functor f, Z.Idx i f a) => (a -> f a -> a) -> Z.Zipper i f a -> a
foldSpine f z =
  case Z.up z of
    Nothing -> z ^. Z.focus_
    Just parent ->
      let next = f (parent ^. Z.focus_) (fmap Comonad.extract . Z.branches $ parent)
       in foldSpine f (parent & Z.focus_ .~ next)

data UndoZipper a
  = UndoZipper
      (Seq a)
      -- ^ undo states
      (Seq a)
      -- ^ redo states