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

import Control.Concurrent
import Control.Exception
import Control.Monad
import Control.Monad.Trans
import Control.Monad.Trans.State
import Data.Bits
import qualified Data.ByteString.UTF8 as B
import Data.Function
import qualified Data.Map as M
import Data.Maybe
import Data.String.AnsiEscapeCodes.Strip.Text
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 Debug.Trace
import Foreign.C.String
import Foreign.C.Types
import GHC.IO.Encoding
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.Posix.Pty

doubleClickDelay = 300

a -. b
  | b > a = 0
  | otherwise = a - b
infixl 1 -.

memo f = unsafePerformIO $ do
  mvar <- newMVar M.empty
  pure $ \x -> unsafePerformIO $ do
    lut <- takeMVar mvar
    case M.lookup x lut of
      Just v -> pure v
      Nothing -> do
        let v = f x
        putMVar mvar $ M.insert x v lut
        pure v

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

search :: Rope -> Rope -> Word -> Maybe Word
search needle haystack p = go (n - 1) p
  where
    index rope i = R.splitAt (i + 1) rope & fst & R.splitAt i & snd
    go i j
      | j >= R.length haystack = Nothing
      | index needle i == index haystack j = if i == 0 then Just j else go (i - 1) (j - 1)
      | otherwise = go (n - 1) $ j + n - memo (match needle) (i, index haystack j)

    match needle (i, c)
      | i == 0 = 0
      | index needle (i - 1) == c = i
      | otherwise = match needle (i - 1, c)

    n = R.length needle

bsearch :: Rope -> Rope -> Word -> Maybe Word
bsearch needle haystack p = go 0 p
  where
    index rope i = R.splitAt (i + 1) rope & fst & R.splitAt i & snd
    go i j
      | j == 0 = Nothing
      | i == n = Just j
      | index needle i == index haystack j = go (i + 1) (j + 1)
      | otherwise = go 0 $ j -. memo (match needle) (i, index haystack j) + 1

    match needle (i, c)
      | i + 1 == n = n -. 1
      | index needle (i + 1) == c = i
      | otherwise = match needle (i + 1, c)

    n = R.length needle

expandAround :: Rope -> Word -> (Word, Word)
expandAround rp p =
  case (getc p, getc (p + 2)) of
    ("(", _) -> (p, searchFor True 0 (== "(") (== ")") (p + 1))
    (_, ")") -> (searchFor False 0 (== ")") (== "(") p, p + 1)
    ("[", _) -> (p, searchFor True 0 (== "[") (== "]") (p + 1))
    (_, "]") -> (searchFor False 0 (== "]") (== "[") p, p + 1)
    ("{", _) -> (p, searchFor True 0 (== "{") (== "}") (p + 1))
    (_, "}") -> (searchFor False 0 (== "}") (== "{") p, p + 1)
    ("'", _) -> (p, searchFor True 0 (== "'") (== "'") (p + 1))
    (_, "'") -> (searchFor False 0 (== "'") (== "'") p, p + 1)
    ("\"", _) -> (p, searchFor True 0 (== "\"") (== "\"") (p + 1))
    (_, "\"") -> (searchFor False 0 (== "\"") (== "\"") p, p + 1)
    ("`", _) -> (p, searchFor True 0 (== "`") (== "`") (p + 1))
    (_, "`") -> (searchFor False 0 (== "`") (== "`") p, p + 1)
    _ -> (searchFor False 0 ws ws p, searchFor True 0 ws ws p)
  where
    l = R.length rp

    ws x = x `elem` [" ", "\t", "(", ")", "[", "]", "{", "}", "\n", "'", "\"", ",", ";", ":", "|", "!"]

    getc i
      | i == 0 = let (sel, _) = R.splitAt 1 rp in R.toText sel
      | i < l = let (R.splitAt (i - 1) -> (_, sel), _) = R.splitAt i rp in R.toText sel
    getc _ = ""

    searchFor fwd n a b p =
      case getc p of
        "" -> if fwd then p - 1 else p
        c ->
          if b c && n == 0
            then if fwd then p - 1 else p
            else searchFor fwd (n + if a c then 1 else if b c then -1 else 0) a b (if fwd then p + 1 else p - 1)

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
      selstr = R.splitAt end (content b) & fst & R.splitAt start & snd
  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
          if (time - lastLeftClick b) < doubleClickDelay
            then put $ b {cursor = expandAround (content b) p, selection = Nothing, lastLeftClick = time}
            else put $ b {cursor = (p, p), selection = Just (x, y), lastLeftClick = time}
        (_, 1) ->
          -- left click up
          put $ b {selection = Nothing}
        (0, 3) ->
          -- right click down
          case search selstr (content b) (1 + end) of
            Just p -> do
              (_, _, _, _, m, _, _) <- lift $ getGeometry display win
              let (height, width) = fontSize display font
                  pageheight = fromIntegral m `div` fromIntegral height
              put $ b {cursor = (p, p + end - start), pos = max (pos b) $ R.lengthInLines (R.splitAt p (content b) & fst) + 1 -. pageheight}
            Nothing -> pure ()
        (0, 2) ->
          -- middle click down
          case bsearch selstr (content b) (start -. 1) of
            Just p -> do
              put $ b {cursor = (p + start - end, p), pos = min (pos b) $ R.lengthInLines (R.splitAt p (content b) & fst) -. 1}
            Nothing -> pure ()
        (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, 117) ->
                  -- ctrl-u (delete content before cursor)
                  put $ b {content = R.splitAt start (content b) & snd, cursor = (0, end-start), pos = 0}
                (4, 105) ->
                  -- ctrl-i (delete content after cursor)
                  put $ b {content = R.splitAt end (content b) & fst}
                (4, 108) ->
                  -- ctrl-l (restrip buffer of ansi codes)
                  put $ b {content = content b & R.toText & stripAnsiEscapeCodes & R.fromText}
                (4, 99) -> lift $ do
                  -- ctrl-c (copy)
                  storeBuffer display (R.toText selstr & T.unpack) 0
                  xSetSelectionOwner display pRIMARY win time
                  xSetSelectionOwner display sECONDARY win time
                (4, 120) -> do
                  -- ctrl-x (cut)
                  lift $ storeBuffer display (R.toText selstr & T.unpack) 0
                  put $ b {content = delins start end (content b) "", cursor = (start + 1, start + 1)}
                (4, 118) -> lift $ do
                  -- ctrl-v (paste)
                  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)
                  lift $ send b (selstr <> "\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 = start -. 1
                  put $ b {content = delins s end (content b) "", cursor = (s, s)}
                  when (start == R.length (content b)) $ 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 (R.fromText . T.pack . map castCCharToChar -> str) <- lift $ getWindowProperty8 display prop win
      if start == R.length (content b)
        then lift $ send b str
        else do
          let s = start + n
              n = fromIntegral (R.length str)
          put $ b {content = delins start end (content b) 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 0xf7f7 0xeaea 0xffff 0xffff) $ \selcolour -> do
          withXftColorValue display visual colourmap (XRenderColor 0xaeae 0xf7f7 0xffff 0xffff) $ \linecolour -> do
            selectInput display win (exposureMask .|. buttonPressMask .|. buttonReleaseMask .|. keyPressMask .|. pointerMotionMask)
            allocaXEvent $ \e ->
              flip runStateT (Buffer 0 "" (0, 0) Nothing pty 0) $
                S.parList id [S.repeatM (lift $ watch pty), S.repeatM (lift $ waitEvent display e)] & S.fold (F.drainMapM (handleEvent display win bgcolour fgcolour linecolour selcolour font))