module UI.Widgets.Common
  ( module UI.Widgets.Common
  , module UI.Terminal.IO
  , module Control.Monad
  , module Data.Text
  , module Control.Monad.IO.Class
  , module Control.Monad.State.Strict
  , module Data.Constraint
  , module GHC.Stack
  ) where

import Common
import Control.Concurrent.STM (atomically)
import Control.Concurrent.STM.TChan
import Control.Exception
import qualified System.IO as SIO
import Control.Monad.IO.Class
import Control.Monad.Loops (iterateWhile)
import Data.Bits
import Data.Word
import Data.Constraint
import Data.Kind (Type)
import Data.Map.Strict as M hiding (keys)
import Data.Maybe
import Data.Text as T
import Data.Text hiding (lines)
import qualified Data.Text as C
import Data.Text.IO as T
import Data.Typeable (Proxy(..), Typeable, cast, typeRep)
import Data.Vector.Mutable (IOVector)
import qualified Data.Vector.Mutable as MV
import GHC.Stack
import System.Random
import qualified System.Terminal as TERM
import UI.Chars
import UI.Terminal.IO

import Highlighter.Highlighter
import Control.Monad
import Control.Monad.State.Strict
import qualified System.Console.ANSI as A
import qualified System.IO as S

data WidgetState = WidgetState
  { wsWidgets         :: Map Int SomeWidget
  , wsCursorWidget    :: Maybe SomeKeyInputWidget -- Widget that authoritativly decide the status/location of cursor. Does not decide what widgets receive keyboard input
  , wsScreenState     :: ScreenState
  , wsCursorVisible   :: Bool
  , wsScreenStateBack :: ScreenState
  }

data ScreenState = ScreenState
  { ssLines          :: IOVector [StyledText]
  , ssCursorPos      :: ScreenPos
  , ssColumns        :: Int
  , ssCursorOverflow :: Bool
  }

setCursorVisibility :: WidgetC m => Bool ->  m ()
setCursorVisibility b = modify (\ws -> ws { wsCursorVisible = b })

emptyScreenState :: Int -> Int -> IO ScreenState
emptyScreenState rows cols = do
  stLines <- MV.generate rows (\_ -> [Plain (T.replicate cols " ")])
  pure (ScreenState stLines (ScreenPos 0 0) cols True)

emptyWidgetState :: Int -> Int -> IO WidgetState
emptyWidgetState lineCount columns = do
  ss <- emptyScreenState lineCount columns
  ssBack <- emptyScreenState lineCount columns
  pure $ WidgetState mempty Nothing ss True ssBack

type WidgetM m a = MonadIO m => StateT WidgetState m a

runWidgetM' :: MonadIO m => WidgetM m a -> m (a, WidgetState)
runWidgetM' act = do
  ws <- liftIO $ emptyWidgetState 0 0
  flip runStateT ws act

runWidgetM :: MonadIO m => WidgetM m a -> m a
runWidgetM act =  fst <$> runWidgetM' act

type WidgetC m =
  ( HasCallStack
  , HasCharScreen m
  , HasRandom m
  , HasLog m
  , HasTerminal m
  , MonadState WidgetState m
  , MonadIO m
  )

getScreenBounds :: WidgetC m => m Dimensions
getScreenBounds = do
  screenState <- wsScreenState <$> get
  let
    screenLines = ssLines screenState
    screenColumns = ssColumns screenState
  pure $ Dimensions screenColumns (MV.length screenLines)

instance MonadIO m => HasRandom (StateT WidgetState m) where
  getRandom = liftIO randomIO

instance MonadIO m => HasCharScreen (StateT WidgetState m) where
  csInitialize (Dimensions cols rows) = do
    -- Initialize the screen memory for the dimensions
    -- and initialize to whitespaces.
    (ss, ssBack) <- liftIO $ do
      ss <- emptyScreenState rows cols
      ssBack <- emptyScreenState rows cols
      pure (ss, ssBack)
    modify (\ws -> ws { wsScreenState = ss, wsScreenStateBack = ssBack })

  csClear = do
    -- Clears the back buffer before starting to write
    -- stuff.
    bb <- wsScreenStateBack <$> get
    liftIO $ MV.set (ssLines bb) [Plain (T.replicate (ssColumns bb) " ")]

  csDraw = do
    -- ^ Compares the stuff that has been written to backbuffer
    -- with the stuff already on frontbuffer, and send the instructions
    -- to draw the changes. Then switch frontbuffer and backbuffers to
    -- prepare for the next draw cycle.
    WidgetState { wsScreenState = (ssLines -> ss), wsScreenStateBack = (ssLines -> ssb) } <- get
    liftIO $ MV.imapM_ (\idx neLine -> do
      oldLine <- MV.read ss idx
      if (oldLine /= neLine)
        then do
          A.setCursorPosition idx 0
          -- mapM_ (\x -> do T.putStr x; S.hFlush stdout; threadDelay 10000;) (stRender <$> neLine)
          mapM_ T.putStr (stRender <$> neLine)
          S.hFlush stdout
        else pure ()
      ) ssb
    wsCursorVisible <$> get >>= \case
      False -> pure ()
      True ->
        (wsCursorWidget <$> get) >>= \case
          Just (SomeKeyInputWidget fref) -> getCursorInfo fref >>= \case
            Just (cl, csst) -> do
              liftIO $ A.setCursorPosition (sY cl) (sX cl)
              putTextFlush $ cursorStyleCode csst
            Nothing -> pure ()
          Nothing -> pure ()
    modify (\ws -> ws { wsScreenStateBack = wsScreenState ws, wsScreenState = wsScreenStateBack ws })

  csPutText t = do
    -- Write stuff to the backbuffer. If the cursor is in an overflow position, then do nothing.
    (ScreenState {ssLines = ssLns, ssCursorOverflow = cursorOverflow, ssCursorPos = ScreenPos cx cy}) <- wsScreenStateBack <$> get
    if cursorOverflow
      then pure ()
      else liftIO $ flip (MV.modify ssLns) cy $ \l -> stInsert l cx t

  csSetCursorPosition x y = do
    -- Sets the cursor position in the backbuffer.
    modify (\ws ->
      let
        screenState = wsScreenStateBack ws
        screenLines = ssLines screenState
        screenColumns = ssColumns screenState
      in if (x >= 0 && x < screenColumns) && (y >= 0 && y < (MV.length screenLines))
        then ws { wsScreenStateBack = screenState { ssCursorOverflow = False, ssCursorPos = ScreenPos x y }}
        else ws { wsScreenStateBack = screenState { ssCursorOverflow = True }})

getTerminalSizeIO :: IO (Maybe (Int, Int))
getTerminalSizeIO = do
  A.getTerminalSize >>= \case
    Just (y, x) -> pure $ Just (x, y)
    Nothing     -> pure Nothing

instance MonadIO m => HasTerminal (StateT WidgetState m) where
  setCursorPosition x y = do
    liftIO $ A.setCursorPosition y x
    hFlush
  hideCursor = liftIO A.hideCursor
  showCursor = do
    liftIO A.showCursor
    hFlush
  putText t = liftIO $ do
    T.putStr t
  putTextFlush t = do
    putText t
    hFlush
  hFlush = liftIO $ S.hFlush S.stdout
  hSetEcho h b = liftIO $ S.hSetEcho h b
  hGetChar = liftIO $ S.hGetChar S.stdin
  hSetBuffering h b = liftIO $ S.hSetBuffering h b
  hWaitForInput = liftIO $ S.hWaitForInput stdin 0
  clearscreen = do
    liftIO A.clearScreen
    hFlush
  clearline = liftIO $ A.hClearFromCursorToLineEnd stdout

instance MonadIO m => HasLog (StateT WidgetState m) where
  appendLog a = liftIO (appendLog a)

-- Below, the type parameter `a` is left in case we need to use tagged
-- references, like an IORef.
newtype WRef (a :: Type) = WRef Int
  deriving newtype (Eq, Ord)

strToKeyEvent :: String -> [KeyEvent]
-- gnome-terminal
strToKeyEvent ('\DEL': rst)  =  (KeyCtrl False False False Backspace):strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'A' : rst)  =  (KeyCtrl False False False ArrowUp) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'B' : rst)  =  (KeyCtrl False False False ArrowDown) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'H' : rst)  =  (KeyCtrl False False False Home) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'F' : rst)  =  (KeyCtrl False False False End) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'C' : rst)  =  (KeyCtrl False False False ArrowRight) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : 'D' : rst)  =  (KeyCtrl False False False ArrowLeft) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '2' : '~': rst)  =  (KeyCtrl False False False Insert) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '3' : '~': rst)  =  (KeyCtrl False False False Del) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1' : '5': '~': rst)  =  (KeyCtrl False False False (Fun 5)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1' : '7': '~': rst)  =  (KeyCtrl False False False (Fun 6)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '2' : '1': '~': rst)  =  (KeyCtrl False False False (Fun 10)) : strToKeyEvent rst
strToKeyEvent ('\ESC': 'O' : 'P': rst)  =  (KeyCtrl False False False (Fun 1)) : strToKeyEvent rst
strToKeyEvent ('\ESC': 'O' : 'Q': rst)  =  (KeyCtrl False False False (Fun 2)) : strToKeyEvent rst
strToKeyEvent ('\ESC': 'O' : 'R': rst)  =  (KeyCtrl False False False (Fun 3)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': '9': ';' : '5' : '~' : rst)  = (KeyCtrl True False False (Fun 8)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': ';' : '2' : 'C' : rst)  = (KeyCtrl True False False (ArrowRight)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': ';' : '2' : 'D' : rst)  = (KeyCtrl True False False (ArrowLeft)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': ';' : '2' : 'B' : rst)  = (KeyCtrl True False False (ArrowDown)) : strToKeyEvent rst
strToKeyEvent ('\ESC': '[' : '1': ';' : '2' : 'A' : rst)  = (KeyCtrl True False False (ArrowUp)) : strToKeyEvent rst

-- Linux term
strToKeyEvent ('\ESC': '[' : '[' : 'E': rst)  =  (KeyCtrl False False False (Fun 5)) : strToKeyEvent rst
--
strToKeyEvent ('\ESC': c : rst )  =  (KeyChar False False True c) : strToKeyEvent rst
strToKeyEvent ('\ESC' : rst)  =  (KeyCtrl False False False Esc) : strToKeyEvent rst
strToKeyEvent ('\n': rst)  =  (KeyCtrl False False False Return) : strToKeyEvent rst
strToKeyEvent str       =  KeyChar False False False <$> str

-- Some terminals sets the eighth bit, instead of sending
-- Escape code to indicate alt key sequences. So this function
-- checks for the 8th bit, and if it is set, sets the alt modifier
-- and changes the original char to one with 8th bit unset.
convertFrom8Bit :: KeyEvent -> KeyEvent
convertFrom8Bit kc@(KeyChar c s _ char) = let
  cw = (toEnum $ fromEnum char) :: Word8
  mask' = complement (bit 7)
  in if ((bit 7 .&. cw) > 0) then (KeyChar c s True (toEnum $ fromEnum (cw .&. mask'))) else kc
convertFrom8Bit c = c

readTerminalEvent :: (TERM.MonadScreen m, TERM.MonadInput m, MonadIO m) => m [TerminalEvent]
readTerminalEvent = TERM.awaitEvent >>= \case
    Left _ -> do
      pure [TerminalInterrupt]
    Right x ->
      case x of
        TERM.WindowEvent TERM.WindowSizeChanged -> do
          TERM.Size h w <- TERM.getWindowSize
          pure [TerminalResize w h]
        TERM.KeyEvent k mods -> pure $ TerminalKey <$> case k of
          TERM.CharKey c -> [convertFrom8Bit $ setModifiers mods $ KeyChar False False False c]
          TERM.ArrowKey TERM.Upwards -> [setModifiers mods $ KeyCtrl False False False ArrowUp]
          TERM.ArrowKey TERM.Downwards -> [setModifiers mods $ KeyCtrl False False False ArrowDown]
          TERM.ArrowKey TERM.Rightwards -> [setModifiers mods $ KeyCtrl False False False ArrowRight]
          TERM.ArrowKey TERM.Leftwards -> [setModifiers mods $ KeyCtrl False False False ArrowLeft]
          TERM.EscapeKey -> [setModifiers mods $ KeyCtrl False False False Esc]
          TERM.EnterKey -> [setModifiers mods $ KeyCtrl False False False Return]
          TERM.FunctionKey n -> [setModifiers mods $ KeyCtrl False False False (Fun n)]
          TERM.BackspaceKey -> [setModifiers mods $ KeyCtrl False False False Backspace]
          TERM.HomeKey -> [setModifiers mods $ KeyCtrl False False False Home]
          TERM.EndKey -> [setModifiers mods $ KeyCtrl False False False End]
          TERM.DeleteKey -> [setModifiers mods $ KeyCtrl False False False Del]
          _ -> []
        _ -> pure []
  where
    setModifiers :: TERM.Modifiers -> KeyEvent -> KeyEvent
    setModifiers modifiers (KeyCtrl _ _ _ c) = KeyCtrl ((modifiers .&. TERM.ctrlKey) /= mempty) ((modifiers .&. TERM.shiftKey) /= mempty) ((modifiers .&. TERM.altKey) /= mempty) c
    setModifiers modifiers (KeyChar _ _ _ c) = KeyChar ((modifiers .&. TERM.ctrlKey) /= mempty) ((modifiers .&. TERM.shiftKey) /= mempty) ((modifiers .&. TERM.altKey) /= mempty) c

readKey :: IO [KeyEvent]
readKey = do
  k <- readKey_
  pure $ strToKeyEvent k

readKey_ :: IO String
readKey_ = do
  char <- S.hGetChar S.stdin
  readRest [char]
  where
    readRest :: [Char] -> IO [Char]
    readRest t = S.hWaitForInput stdin 0 >>= \case
      True -> do
        c <- S.hGetChar S.stdin
        readRest (c:t)
      False -> pure $ Prelude.reverse t

uiLoop :: forall m event. (Show event, WidgetC m) => TChan event -> (event -> m Bool) -> m ()
uiLoop es cback = do
  void $ iterateWhile id $ do
    event <- liftIO $ atomically $ readTChan es
    r <- cback event
    pure r

cursorStyleCode :: CursorStyle -> Text
cursorStyleCode Bar       = T.pack $ "\ESC[5 q" <> A.showCursorCode
cursorStyleCode Underline = T.pack $ "\ESC[4 q" <> A.showCursorCode
cursorStyleCode Hidden    = T.pack A.hideCursorCode

readWRef :: forall a m. (WidgetC m, Widget a) => WRef a -> m a
readWRef (WRef ref) = do
  (fromMaybe (error "not found") . M.lookup ref . wsWidgets) <$> get >>= \case
    SomeWidget w -> case cast w of
      Just a  -> pure a
      Nothing -> error "Unexpected type"

modifyWRef :: (WidgetC m, Widget a) => WRef a -> (a -> a) -> m ()
modifyWRef (WRef ref) fn =
  modify $ \s -> s { wsWidgets = M.update (Just . (modifySomeWidget fn)) ref $ wsWidgets s }

modifyWRefM :: (WidgetC m, Widget a) => WRef a -> (a -> m a) -> m ()
modifyWRefM (WRef ref) fn = do
  m <- wsWidgets <$> get
  case M.lookup ref m of
    Just sw -> do
      nSw <- modifySomeWidgetM fn sw
      modify $ \s -> s { wsWidgets = M.update (\_ -> Just nSw) ref $ wsWidgets s }
    Nothing -> pure ()

modifySomeWidget :: Widget a => (a -> a) -> SomeWidget -> SomeWidget
modifySomeWidget fn (SomeWidget w) = case cast w of
  Just a  -> (SomeWidget (fn a))
  Nothing -> error "unexpected type"

modifySomeWidgetM :: (Monad m, Widget a) => (a -> m a) -> SomeWidget -> m SomeWidget
modifySomeWidgetM fn (SomeWidget w) = case cast w of
  Just a  -> do
    n <- fn a
    pure (SomeWidget n)
  Nothing -> error "unexpected type"

-- Insert the new widget at a random key in the Widget state map
-- and return the key.
newWRef :: (WidgetC m, Widget a) => a -> m (WRef a)
newWRef a = do
  ref <- getRandom
  modify $ \s -> s { wsWidgets = M.insert ref (SomeWidget a) $ wsWidgets s }
  pure (WRef ref)

data CtrlKey
  = Del
  | Esc
  | Insert
  | End
  | Home
  | ArrowLeft
  | ArrowRight
  | ArrowUp
  | ArrowDown
  | Backspace
  | Fun Int
  | Return
  deriving (Show, Ord, Eq)

data KeyEvent
  = KeyChar Bool Bool Bool Char -- Bool fields for modifiers for ctrl, shift, alt
  | KeyCtrl Bool Bool Bool CtrlKey
  deriving (Show, Eq, Ord)

data TerminalEvent
  = TerminalKey KeyEvent
  | TerminalResize Int Int
  | TerminalInterrupt
  deriving Show

pushTerminalEventToHandle :: SIO.Handle -> TerminalEvent -> IO ()
pushTerminalEventToHandle handle' kv = case kv of
  TerminalKey (KeyChar False False False c) -> do
    SIO.hPutChar handle' c
    SIO.hFlush handle'
  TerminalKey (KeyCtrl False False False Backspace) -> do
    SIO.hPutChar handle' '\BS'
    SIO.hFlush handle'
  TerminalKey (KeyCtrl False False False Esc) -> do
    SIO.hPutChar handle' '\ESC'
    SIO.hFlush handle'
  TerminalKey (KeyCtrl False False False Return) -> do
    SIO.hPutChar handle' '\n'
    SIO.hFlush handle'
  _ -> pass

data TerminalException
  = TerminalException Text
  deriving (Show)

instance Exception TerminalException

class HasRandom m where
  getRandom :: Random a => m a

class HasCursor m where
  getCursor :: m CursorInfo

class HasCharScreen m where
  csInitialize :: Dimensions -> m ()
  csClear :: m ()
  csDraw :: m ()
  csPutText :: StyledText -> m ()
  csSetCursorPosition :: Int -> Int -> m ()

class Layout a where
  addWidget :: (WidgetC m, Widget child) => WRef a -> Text -> WRef child -> m ()
  setTextFocus :: WidgetC m => WRef a -> Text -> m ()

class Drawable a where
  draw :: (WidgetC m) => WRef a -> m ()
  setVisibility :: WidgetC m => WRef a -> Bool -> m ()
  getVisibility :: WidgetC m => WRef a -> m Bool

class Moveable a where
  move ::  WidgetC m => WRef a -> ScreenPos -> m ()
  getPos :: WidgetC m => WRef a -> m ScreenPos
  getDim :: WidgetC m => WRef a -> m Dimensions
  resize :: WidgetC m => WRef a -> (Dimensions -> Dimensions) -> m ()

class Container a c | a -> c where
  setContent :: WidgetC m => WRef a -> c -> m ()
  getContent :: WidgetC m => WRef a -> m c

class Selectable a where
  getSelection :: WidgetC m => WRef a -> m Text

class (Typeable a, Drawable a) => KeyInput a where
  getCursorInfo :: WidgetC m => WRef a -> m (Maybe CursorInfo)
  handleInput :: WidgetC m => WRef a -> KeyEvent -> m ()

data WidgetCapability a (c :: Constraint) where
  KeyInputCap :: WRef a -> WidgetCapability a (KeyInput a)
  MoveableCap :: WRef a -> WidgetCapability a (Moveable a)
  SelectableCap :: WRef a -> WidgetCapability a (Selectable a)
  DrawableCap :: WRef a -> WidgetCapability a (Drawable a)
  ContainerCap :: Typeable cnt => WRef a -> Proxy cnt -> WidgetCapability a (Container a cnt)

class Typeable a => Widget a where
  hasCapability :: WidgetCapability a c -> Maybe (Dict c)

withCapability :: forall a c m b . (WidgetC m, Widget a, Typeable c) => WidgetCapability a c -> (c => m b)  -> m b
withCapability cap fn = case hasCapability cap of
  Just Dict -> fn
  Nothing   -> error ("No capability:" <> (show $ typeRep (Proxy @c)))

data SomeWidget where
  SomeWidget :: forall a. Widget a => a -> SomeWidget

data SomeWidgetRef where
  SomeWidgetRef :: forall a. (Typeable a, Widget a) => WRef a -> SomeWidgetRef

data SomeKeyInputWidget where
  SomeKeyInputWidget :: KeyInput a => WRef a -> SomeKeyInputWidget

wSetCursor :: (WidgetC m, HasTerminal m) => ScreenPos -> m ()
wSetCursor ScreenPos {..} = csSetCursorPosition sX sY

wSetCursorRel :: (HasTerminal m, WidgetC m) => ScreenPos -> ScreenPos -> m ScreenPos
wSetCursorRel o rel = do
  let n = moveRight (sX rel) $ moveDown (sY rel) o
  wSetCursor n
  pure n

drawTitleLine :: WidgetC m => ScreenPos -> Int -> Int -> Maybe Text -> m ()
drawTitleLine sp width titleOffset (fromMaybe "" -> title) = do
  wSetCursor sp
  let titleLength = T.length title
  csPutText $ StyledText NoStyle [Plain $ C.concat [C.replicate (titleOffset - 1) (C.singleton horizontalLine), C.singleton verticalRight], colorText A.White A.Blue title, Plain (C.singleton verticalLeft <> C.replicate (width - titleLength - titleOffset - 1) (C.singleton horizontalLine)) ]

drawBorderBox :: WidgetC m => ScreenPos -> Dimensions -> m ()
drawBorderBox sp Dimensions {..} = do
  wSetCursor sp
  csPutText $ Plain $ C.concat [C.singleton cornerLT, C.replicate (diW - 2) (C.singleton horizontalLine), C.singleton cornerRT]
  wSetCursor $ moveDown (diH - 1) sp
  csPutText $ Plain $ C.concat [C.singleton cornerLB, C.replicate (diW - 2) (C.singleton horizontalLine), C.singleton cornerRB]
  forM_ [1..(diH - 2)] (\r -> do
    wSetCursor $ moveDown r sp
    csPutText $ Plain $ C.concat [C.singleton verticalLine]
    wSetCursor $ moveRight (diW - 1) $ moveDown r sp
    csPutText $ Plain $ C.concat [C.singleton verticalLine]
    )