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 Control.Monad.IO.Class
import Control.Monad.Loops (iterateWhile)
import Data.Constraint
import Data.Kind (Type)
import qualified Data.List as DL
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 qualified Graphics.Vty as VTY
import System.Random
import UI.Chars
import UI.Terminal.IO

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

initializeVty :: IO VTY.Vty
initializeVty =
  VTY.standardIOConfig >>= VTY.mkVty

shutdownVty :: VTY.Vty -> IO ()
shutdownVty vty = VTY.shutdown vty

readVtyEvent :: VTY.Vty -> IO [TerminalEvent]
readVtyEvent vty = VTY.nextEvent vty >>= \x -> do
  case x of
    VTY.EvResize w h -> pure [TerminalResize w h]
    VTY.EvKey k mods -> pure $ TerminalKey <$> case k of
      VTY.KChar c -> [setModifiers mods $ KeyChar False False False c]
      VTY.KUp -> [setModifiers mods $ KeyCtrl False False False ArrowUp]
      VTY.KDown -> [setModifiers mods $ KeyCtrl False False False ArrowDown]
      VTY.KRight -> [setModifiers mods $ KeyCtrl False False False ArrowRight]
      VTY.KLeft -> [setModifiers mods $ KeyCtrl False False False ArrowLeft]
      VTY.KEsc -> [setModifiers mods $ KeyCtrl False False False Esc]
      VTY.KEnter -> [setModifiers mods $ KeyCtrl False False False Return]
      VTY.KFun 32 -> [setModifiers (VTY.MCtrl : mods) $ KeyCtrl False False False (Fun 8)]
      VTY.KFun n -> [setModifiers mods $ KeyCtrl False False False (Fun n)]
      VTY.KBS -> [setModifiers mods $ KeyCtrl False False False Backspace]
      VTY.KHome -> [setModifiers mods $ KeyCtrl False False False Home]
      VTY.KEnd -> [setModifiers mods $ KeyCtrl False False False End]
      VTY.KDel -> [setModifiers mods $ KeyCtrl False False False Del]
      _ -> []
    _ -> pure []
  where
    setModifiers :: [VTY.Modifier] -> KeyEvent -> KeyEvent
    setModifiers mods key = DL.foldl' foldFn key mods

    foldFn :: KeyEvent -> VTY.Modifier -> KeyEvent
    foldFn (KeyCtrl c _ a v) VTY.MShift = KeyCtrl c True a v
    foldFn (KeyCtrl _ s a v) VTY.MCtrl = KeyCtrl True s a v
    foldFn (KeyCtrl c s _ v) VTY.MMeta = KeyCtrl c s True v
    foldFn (KeyCtrl c s _ v) VTY.MAlt = KeyCtrl c s True v
    foldFn (KeyChar c _ a v) VTY.MShift = KeyChar c True a v
    foldFn (KeyChar _ s a v) VTY.MCtrl = KeyChar True s a v
    foldFn (KeyChar c s _ v) VTY.MMeta = KeyChar c s True v
    foldFn (KeyChar c s _ v) VTY.MAlt = KeyChar c s True v

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
  | KeyCtrl Bool Bool Bool CtrlKey
  deriving (Show, Eq, Ord)

data TerminalEvent
  = TerminalKey KeyEvent
  | TerminalResize Int Int
  deriving Show

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

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]
    )