{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}

import Control.Concurrent (threadDelay)
import Foreign.C.String
import Data.String.AnsiEscapeCodes.Strip.Text
import qualified Data.ByteString.UTF8 as B
import System.Posix.Pty
import GHC.IO.Encoding
import Control.Exception
import Control.Monad
import Control.Monad.Trans
import Control.Monad.Trans.State
import Data.Bits
import Data.Function
import Data.Maybe
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.Text.Rope (Position (..), Rope)
import qualified Data.Text.Rope as R
import Foreign.C.Types
import Graphics.X11.Types
import Graphics.X11.Xft
import Graphics.X11.Xlib hiding (Buffer, Position)
import Graphics.X11.Xlib.Extras
import Graphics.X11.Xrender
import qualified Streamly.Data.Fold as F
import qualified Streamly.Data.Stream.Prelude as S
import System.Environment
import System.IO.Unsafe
import System.Process

data Buffer = Buffer
  { pos :: Word,
    content :: Rope,
    cursor :: (Word, Word),
    selection :: Maybe (CInt, CInt),
    pty :: Pty
  }

drawWin display win bgcolour fgcolour linecolour selcolour font = do
  buf <- get
  lift $ do
    (_, _, _, n, m, _, _) <- getGeometry display win
    resizePty (pty buf) (fromIntegral n`div`fromIntegral width, fromIntegral m`div`fromIntegral height)
    bracket (createGC display win) (freeGC display) $ \gc -> do
      bracket (createPixmap display win n m (defaultDepthOfScreen (defaultScreenOfDisplay display))) (freePixmap display) $ \p -> do
        withXftDraw display p (defaultVisualOfScreen (defaultScreenOfDisplay display)) (defaultColormap display 0) $ \draw -> do
          xftDrawRect draw bgcolour 0 0 n m

          -- Get rope covering window area
          let winlines = fromIntegral m `div` fromIntegral height
          let (pre, R.splitAtLine winlines -> (vis, _)) = R.splitAtLine (pos buf) (content buf)

          -- relative cursor position
          let (s, e) = cursor buf
              s' :: Int = fromIntegral s - fromIntegral (R.length pre)
              e' :: Int = fromIntegral (max e $ s + 1) - fromIntegral (R.length pre)

          let (R.splitAt (fromIntegral $ max 0 s') -> (presel, sel), postsel) = R.splitAt (fromIntegral $ max 0 e') vis
              Position l1 c1 = R.lengthAsPosition presel
              Position l2 c2 = R.lengthAsPosition (presel <> sel)

          xftDrawRect draw linecolour 0 (l1 * fromIntegral height + 4) n height
          drawRope False draw fgcolour bgcolour 0 height (R.lines presel)
          drawRope True draw fgcolour selcolour (fromIntegral c1 * width) (fromIntegral (l1 + 1) * height) (R.lines sel)
          drawRope False draw fgcolour bgcolour (fromIntegral c2 * width) (fromIntegral (l2 + 1) * height) (R.lines postsel)
        copyArea display p win gc 0 0 n m 0 0
  where
    drawRope _ _ _ _ _ _ [] = pure ()
    drawRope drawbg draw fgcolour bgcolour x y (l : ls) = do
      when drawbg $ xftDrawRect draw bgcolour x (y - height + 4) (width * T.length l) height
      xftDrawString draw fgcolour font x y (T.unpack l)
      drawRope drawbg draw fgcolour bgcolour 0 (y + height) ls

    (height, width) = fontSize display font

-- assume fixed width fonts
fontSize :: Display -> XftFont -> (Int, Int)
fontSize display font = unsafePerformIO $ do
  extents <- xftTextExtents display font "a"
  height <- xftfont_height font
  pure (height, xglyphinfo_xOff extents)

key mod k = case (mod, k) of
  (_, 32) -> " "
  (0, 39) -> "'"
  (0, 44) -> ","
  (0, 45) -> "-"
  (0, 46) -> "."
  (0, 47) -> "/"
  (0, 48) -> "0"
  (0, 49) -> "1"
  (0, 50) -> "2"
  (0, 51) -> "3"
  (0, 52) -> "4"
  (0, 53) -> "5"
  (0, 54) -> "6"
  (0, 55) -> "7"
  (0, 56) -> "8"
  (0, 57) -> "9"
  (0, 59) -> ";"
  (0, 61) -> "="
  (0, 91) -> "["
  (0, 92) -> "\\"
  (0, 93) -> "]"
  (0, 96) -> "`"
  (0, 97) -> "a"
  (0, 98) -> "b"
  (0, 99) -> "c"
  (0, 100) -> "d"
  (0, 101) -> "e"
  (0, 102) -> "f"
  (0, 103) -> "g"
  (0, 104) -> "h"
  (0, 105) -> "i"
  (0, 106) -> "j"
  (0, 107) -> "k"
  (0, 108) -> "l"
  (0, 109) -> "m"
  (0, 110) -> "n"
  (0, 111) -> "o"
  (0, 112) -> "p"
  (0, 113) -> "q"
  (0, 114) -> "r"
  (0, 115) -> "s"
  (0, 116) -> "t"
  (0, 117) -> "u"
  (0, 118) -> "v"
  (0, 119) -> "w"
  (0, 120) -> "x"
  (0, 121) -> "y"
  (0, 122) -> "z"
  (1, 39) -> "\""
  (1, 44) -> "<"
  (1, 45) -> "_"
  (1, 46) -> ">"
  (1, 47) -> "?"
  (1, 48) -> ")"
  (1, 49) -> "!"
  (1, 50) -> "@"
  (1, 51) -> "#"
  (1, 52) -> "$"
  (1, 53) -> "%"
  (1, 54) -> "^"
  (1, 55) -> "&"
  (1, 56) -> "*"
  (1, 57) -> "("
  (1, 59) -> ":"
  (1, 61) -> "+"
  (1, 91) -> "{"
  (1, 92) -> "|"
  (1, 93) -> "}"
  (1, 96) -> "~"
  (1, 97) -> "A"
  (1, 98) -> "B"
  (1, 99) -> "C"
  (1, 100) -> "D"
  (1, 101) -> "E"
  (1, 102) -> "F"
  (1, 103) -> "G"
  (1, 104) -> "H"
  (1, 105) -> "I"
  (1, 106) -> "J"
  (1, 107) -> "K"
  (1, 108) -> "L"
  (1, 109) -> "M"
  (1, 110) -> "N"
  (1, 111) -> "O"
  (1, 112) -> "P"
  (1, 113) -> "Q"
  (1, 114) -> "R"
  (1, 115) -> "S"
  (1, 116) -> "T"
  (1, 117) -> "U"
  (1, 118) -> "V"
  (1, 119) -> "W"
  (1, 120) -> "X"
  (1, 121) -> "Y"
  (1, 122) -> "Z"
  (_, 65289) -> "\t"
  (_, 65293) -> "\n"
  _ -> ""

data HandledEvent = GraphicsExpose | NoExpose | Expose | Key XKeyEvent | T Text | Selection Event | Notify

waitEvent :: Display -> XEventPtr -> IO HandledEvent
waitEvent display event = do
  nextEvent display event
  ty <- get_EventType event
  if ty == expose
    then
      pure Expose
    else
      if ty == graphicsExpose
        then
          pure GraphicsExpose
        else
          if ty == noExpose
            then
              pure NoExpose
            else
              if ty == selectionRequest
                then
                  Selection <$> getEvent event
                else
                  if ty == selectionNotify
                    then
                      pure Notify
                    else do
                      ev <- get_KeyEvent event
                      pure $ Key ev

handleEvent :: Display -> Drawable -> XftColor -> XftColor -> XftColor -> XftColor -> XftFont -> HandledEvent -> StateT Buffer IO ()
handleEvent display win bgcolour fgcolour linecolour selcolour font event = do
  b <- get
  let (start, end) = cursor b
  case event of
    Key (_, _, time, x, y, _, _, mod, keycode, _) -> do
      keysym <- lift $ keycodeToKeysym display keycode 0
      let toAppend = key mod keysym
      case (mod, keycode) of -- handle buttons
        (0, 1) -> do
          -- left click down
          p <- lift $ getPos x y b
          put $ b {cursor = (p, p), selection = Just (x, y)}
        (_, 1) ->
          -- left click up
          put $ b {selection = Nothing}
        (0, 4) -> scrollup 3 -- scroll wheel
        (0, 5) -> scrolldown 3
        _ ->
          case selection b of
            Just (max 0 -> x0, max 0 -> y0) -> do
              p <- lift $ getPos x0 y0 b
              q <- lift $ getPos x y b
              put $ b {cursor = (min p q, max p q)}
            _ -> do
              case (mod, keysym) of
                (4, 99) -> lift $ do
                  -- ctrl-c (copy)
                  let (R.splitAt start -> (_, sel), _) = R.splitAt end (content b)
                  storeBuffer display (R.toText sel & T.unpack) 0
                  xSetSelectionOwner display pRIMARY win time
                  xSetSelectionOwner display sECONDARY win time
                (4, 120) -> do
                  -- ctrl-x (cut)
                  let (R.splitAt start -> (_, sel), _) = R.splitAt end (content b)
                  lift $ storeBuffer display (R.toText sel & T.unpack) 0
                  put $ b {content = delins start end (content b) "", cursor = (start + 1, start + 1)}
                (4, 118) -> lift $ do
                  -- ctrl-v (paste)
                  -- win <- xGetSelectionOwner display pRIMARY
                  prop <- internAtom display "JTERM_CLIPBOARD" True
                  xConvertSelection display pRIMARY xUtf8String prop win time
                (4, 113) -> lift $ do
                  -- ctrl-q (write log)
                  home <- getEnv "HOME"
                  T.writeFile (home <> "/jterm.log") (R.toText (content b))
                (4, 115) -> do
                  -- ctrl-s (send)
                  let (R.splitAt start -> (_, sel), _) = R.splitAt end (content b)
                  lift $ send b (sel <> "\n")
                (_, 65361) -> do
                  -- left
                  let s = start - min start 1
                  put $ b {cursor = (s, s)}
                (_, 65363) -> do
                  -- right
                  let e = min (R.length (content b)) $ end + 1
                  put $ b {cursor = (e, e)}
                (_, 65535) -> lift $ sendCommand b Interrupt
                (_, 65360) -> put $ b {pos = 0} -- home
                (_, 65367) -> do
                  -- end
                  let e = R.length (content b)
                  put $ b {cursor = (e, e)}
                  scrolldown (maxBound - pos b)
                (_, 65362) -> scrollup 1
                (_, 65364) -> scrolldown 1
                (_, 65365) -> do
                  (_, _, _, _, m, _, _) <- lift $ getGeometry display win
                  let (height, width) = fontSize display font
                      pageheight = fromIntegral m `div` fromIntegral height
                  scrollup ((681 * pageheight) `div` 1000)
                (_, 65366) -> do
                  (_, _, _, _, m, _, _) <- lift $ getGeometry display win
                  let (height, width) = fontSize display font
                      pageheight = fromIntegral m `div` fromIntegral height
                  scrolldown ((681 * pageheight) `div` 1000)
                (_, 65288) -> do
                  -- backspace
                  let s = if start > 0 then start - 1 else 0
                  put $ b {content = delins s end (content b) "", cursor = (s, s)}
                  lift $ sendCommand b Erase
                _ -> when (toAppend /= "") $ do
                  if start == R.length (content b) then
                    lift $ send b toAppend
                  else put $ b {content = delins start end (content b) toAppend, cursor = (start + 1, start + 1)}
      redraw
    T toAppend -> do
      if toAppend == "\b \b" then
        -- Erase character
        let (c, _) = R.splitAt l (content b)
            l = R.length (content b)
            p = min start l
        in put $ b {content = c, cursor = (p, p)}
      else do
        let c = content b <> R.fromText expandedTabs
            pos = R.length c
            noesc = stripAnsiEscapeCodes toAppend
            expandedTabs = T.intercalate "\n" $ map expandTabs $ T.splitOn "\n" noesc
        put $ b {content = c}
        when (start == R.length (content b)) $ modify $ \b -> b {cursor = (pos, pos)}
      redraw
    Notify -> do
      prop <- lift $ internAtom display "JTERM_CLIPBOARD" False
      Just (T.pack . map castCCharToChar -> str) <- lift $ getWindowProperty8 display prop win
      let s = start + n
          n = fromIntegral (T.length str)
      put $ b {content = delins start end (content b) (R.fromText str), cursor = (s, s)}
      redraw
    Expose -> redraw
    Selection e -> lift $ allocaXEvent $ \ev -> do
      setEventType ev selectionNotify
      if ev_target e == xUtf8String
        then do
          clip <- fetchBuffer display 0
          changeProperty8 display (ev_requestor e) (ev_property e) (ev_target e) propModeReplace (map (fromIntegral . fromEnum) clip)
          setSelectionNotify ev (ev_requestor e) (ev_selection e) (ev_target e) (ev_property e) (ev_time e)
        else setSelectionNotify ev (ev_requestor e) (ev_selection e) (ev_target e) 0 (ev_time e)
      sendEvent display (ev_requestor e) False noEventMask ev
      flush display
    _ -> pure ()
  where
    xUtf8String = 315 :: Atom

    sendCommand b c = do
                  attrs <- getTerminalAttributes (pty b)
                  writePty (pty b) (B.fromString [fromJust $ controlChar attrs c])

    redraw = do
      drawWin display win bgcolour fgcolour linecolour selcolour font
      lift $ sync display False

    delins :: Word -> Word -> Rope -> Rope -> Rope
    delins s e rp ins =
      let (before, _) = R.splitAt s rp
          (_, after) = R.splitAt e rp
       in before <> ins <> after

    xy2lc :: Word -> Word -> IO (Word, Word)
    xy2lc x y = do
      let (height, width) = fontSize display font
      let c = x `div` fromIntegral width
          l = y `div` fromIntegral height
      pure (l, c)

    lc2pos :: Word -> Word -> Rope -> Word
    lc2pos l c rp =
      let p = R.splitAtPosition (Position l c) rp & fst & R.length
          pmax = R.splitAtLine (1 + l) rp & fst & R.length
       in min p pmax

    getPos :: CInt -> CInt -> Buffer -> IO Word
    getPos x y b = do
      (l, c) <- xy2lc (max 0 $ fromIntegral x) (max 0 $ fromIntegral y)
      let p = lc2pos (pos b + l) c (content b)
      pure p

    scrollup k = modify $ \b -> b {pos = pos b - min (pos b) k}
    scrolldown k = modify $ \b -> b {pos = min (R.lengthInLines (content b) - 1) (pos b + k)}

    send b cmd = writePty (pty b) (R.toText cmd & T.unpack & B.fromString)

    expandTabs = go . T.splitOn "\t"
      where
        go (a:b:xs) = a <> T.pack (replicate (8-T.length a `mod` 8) ' ') <> go (b : xs)
        go x = mconcat x

watch :: Pty -> IO HandledEvent
watch h = T . T.filter (/= '\r') . T.pack . B.toString <$> readPty h

main :: IO ()
main = do
  setLocaleEncoding utf8
  env <- getEnvironment
  let shell = case filter (\(var,_) -> var == "SHELL") env of [(_,x)] -> x; _ -> "sh"
  (pty, _) <- spawnWithPty (Just $ env ++ [("TERM", "dumb")]) True shell [] (80, 25)

  void $ bracket (openDisplay "") closeDisplay $ \display -> do
    let screenno = defaultScreen display
        border = blackPixel display screenno
        background = whitePixel display screenno
        screen = defaultScreenOfDisplay display
        visual = defaultVisualOfScreen screen
        colourmap = defaultColormap display screenno
    font <- xftFontOpen display screen "Iosevka Extended:size=12"
    rootw <- rootWindow display screenno
    win <- createSimpleWindow display rootw 0 0 100 100 1 border background
    mapWindow display win


    -- ensure clipboard is not empty
    storeBuffer display "" 0

    withXftColorName display visual colourmap "black" $ \fgcolour ->
      withXftColorValue display visual colourmap (XRenderColor 0xffff 0xffff 0xeaea 0xffff) $ \bgcolour -> do
        withXftColorValue display visual colourmap (XRenderColor 0x3030 0xd4d4 0xd4d4 0xffff) $ \selcolour -> do
         withXftColorValue display visual colourmap (XRenderColor 0xaeae 0xeeee 0xeeee 0xffff) $ \linecolour -> do
          selectInput display win (exposureMask .|. buttonPressMask .|. buttonReleaseMask .|. keyPressMask .|. pointerMotionMask)
          allocaXEvent $ \e ->
            flip runStateT (Buffer 0 "" (0, 0) Nothing pty) $
              S.parList (S.eager True) [S.repeatM (lift $ waitEvent display e), S.repeatM (lift $ watch pty)] & S.fold (F.drainMapM (handleEvent display win bgcolour fgcolour linecolour selcolour font))