<commit_before>{-# LANGUAGE OverloadedStrings #-}

import           Control.Monad
import           Control.Monad.ST

import qualified Data.ByteString              as B
import           Data.Char                    (isPrint, isSpace)
import           Data.Function                (on)
import qualified Data.Map.Strict              as M
import           Data.Maybe                   (fromMaybe)
import           Data.Monoid                  ((<>))
import           Data.Text                    (Text)
import qualified Data.Text                    as T
import qualified Data.Text.Encoding           as TE
import qualified Data.Text.Encoding.Error     as Err
import           Data.Text.IO                 as T
import           Data.Vector                  (Vector)
import qualified Data.Vector                  as V
import qualified Data.Vector.Algorithms.Intro as VI

import           System.Exit                  (exitFailure, exitSuccess)
import           System.IO                    (Handle, IOMode (..), hClose,
                                               openFile, hGetChar, hReady)

import           System.Console.ANSI

import           Score
import           Tty

{- TODO: implement something like withTty that handles restoring tty state -}
{- TODO: getting unweildy passing around currMatchCount and choicesToShow-}
{- TODO: look into resource monad for ensuring tty gets closed? -}
{- TODO: state monad? passing around matches and memo tables is cumbersome -}

prompt :: Text
prompt = "> "

data KeyPress = CtrlC
              | CtrlN
              | CtrlP
              | CtrlW
              | CtrlH
              | CtrlU
              | Enter
              | Invisible
              | Backspace
              | PlainChar Char

type Memo = M.Map Text (Vector Text)

data Search = Search
    { query     :: !Text
    , choices   :: !(Vector Text)
    , selection :: !Int
    , matches   :: !(Vector Text)
    }

instance Show Search where
  show _ = "[search]"

data SelecthState = SelecthState
    { s_search        :: !Search
    , s_choicesToShow :: !Int
    , s_scoreMemo     :: !Memo
    } deriving Show

data RenderedSearch = RenderedSearch
    { queryString   :: !Text
    , renderedLines :: !(Vector (Text, SGR))
    } deriving Show


data Action = SearchAction !SearchAction | ExitAction !ExitAction
  deriving Show


data SearchAction = Extend !Text
                  | DropWord
                  | DropChars !Int
                  | Clear
                  | SelectPrev
                  | SelectNext
                  | Ignore
  deriving Show

data ExitAction   = Abort
                  | MakeChoice !Search
  deriving Show

specialChars :: M.Map Char KeyPress
specialChars = M.fromList [ ('\ETX', CtrlC)
                          , ('\r',   Enter)
                          , ('\DEL', Backspace)
                          , ('\BS',  CtrlH)
                          , ('\NAK', CtrlU)
                          , ('\SO',  CtrlN)
                          , ('\DLE', CtrlP)
                          , ('\ETB', CtrlW)
                          ]

charToKeypress :: Char -> KeyPress
charToKeypress c = fromMaybe (if isPrint c then PlainChar c else Invisible)
                             (M.lookup c specialChars)

findMatches :: Memo -> Text -> Vector Text -> (Vector Text, Memo)
findMatches memo qry chs =
    case M.lookup qry memo of
       Just mtchs -> (mtchs, memo)
       Nothing    -> (newMatches, M.insert qry newMatches memo)
  where
    newMatches = getMatches qry chs
    getMatches q cs = V.map fst
                        . sortImmutableVec (flip compare `on` snd)
                        . V.filter (\(_,cScore) -> cScore > 0)
                        $ scoreAll q cs
    sortImmutableVec f v = runST $ do
                              vec <- V.thaw v
                              VI.sortBy f vec
                              V.freeze vec

render :: Search -> Int -> RenderedSearch
render (Search {query=q, matches=matched, selection=selIndex}) csToShow =
    RenderedSearch { queryString   = queryLine
                   , renderedLines = V.cons (queryLine, Reset)
                                            renderedMatchLines
                   }
  where
    renderedMatchLines = swapBackground selIndex
                           . V.map (\m -> (m, Reset))
                           $ matchLines
    matched'   = V.take csToShow matched
    matchLines = pad csToShow matched'
    queryLine  = prompt <> q
    swapBackground i vec =
      V.update
          vec
          (V.singleton
              (selIndex, (fst (vec V.! i), SetSwapForegroundBackground True)))
    pad n xs = xs <> V.replicate (n - V.length xs) " "

draw :: Handle -> RenderedSearch -> IO ()
draw tty rendered = do
    withCursorHidden tty $
        writeLines tty (renderedLines rendered)
    hSetCursorColumn tty . T.length $ queryString rendered

-- TODO: Handle exceptions, ensure tty gets restored to default settings
-- TODO: Ensure handle to tty is closed?  See what GB ensures.
{-restoreTTY :: IO ()-}

dropLastWord :: Text -> Text
dropLastWord = T.stripEnd . T.dropWhileEnd (not . isSpace)

-- TODO: ugh, this should indeed have used choices to show,
-- there can easily be 100s of matches.
writeSelection :: Handle -> Search -> IO ()
writeSelection tty (Search {matches=matches', selection=sel}) = do
    hCursorDown tty (V.length matches')
    T.hPutStr tty "\n"
    T.putStrLn $ fromMaybe (error "Failed to write selection.")
                           (matches' V.!? sel)

handleInput :: Text -> Search -> [Action]
handleInput inputText search = T.foldr buildActions [] inputText
  where
    buildActions c xs = actionFor c : xs
    actionFor c =
      case charToKeypress c of
          Enter     -> ExitAction (MakeChoice search)
          CtrlC     -> ExitAction Abort
          CtrlN     -> SearchAction SelectNext
          CtrlP     -> SearchAction SelectPrev
          CtrlU     -> SearchAction Clear
          CtrlW     -> SearchAction DropWord
          CtrlH     -> SearchAction (DropChars 1)
          Backspace -> SearchAction (DropChars 1)
          Invisible -> SearchAction Ignore
          PlainChar char -> SearchAction $ Extend (T.singleton char)

collapseActions :: [Action] -> [Action]
collapseActions = foldr collapse []
  where
    collapse (SearchAction (Extend t)) (SearchAction (Extend t'):as) =
        SearchAction (Extend (t <> t')) : as
    collapse (SearchAction (DropChars n)) (SearchAction (DropChars m):as) =
        SearchAction (DropChars $ n + m) : as
    collapse nextA as = nextA : as

buildSearch ::SearchAction -> Search -> Int -> Memo -> (Search, Memo)
buildSearch action search choicesToShow memo =
    case action of
        Extend qAddition -> (search', memo')
          where
            search' = search { query   = q'
                             , matches = matches'
                             , selection = 0 }
            q' = query search <> qAddition
            (matches', memo') = findMatches memo q' (matches search)
        DropWord    -> memoSearch $ dropLastWord (query search)
        DropChars n -> memoSearch $ T.dropEnd n (query search)
        Clear       -> if T.null (query search)
                       then (search, memo) -- why isn't this being hit?
                       else memoSearch ""
        Ignore     -> (search, memo)
        SelectNext -> moveSel 1
        SelectPrev -> moveSel (-1)
  where
    memoSearch q =  let (matches', memo') = findMatches memo q (choices search)
                    in  (search { query = q
                                , matches = matches'
                                , selection = 0}
                        , memo')
    moveSel n = ( search { selection = if possibleSelections > 0
                                       then mod (selection search + n)<commit_msg>formatting<commit_after>{-# LANGUAGE OverloadedStrings #-}

import           Control.Monad
import           Control.Monad.ST

import qualified Data.ByteString              as B
import           Data.Char                    (isPrint, isSpace)
import           Data.Function                (on)
import qualified Data.Map.Strict              as M
import           Data.Maybe                   (fromMaybe)
import           Data.Monoid                  ((<>))
import           Data.Text                    (Text)
import qualified Data.Text                    as T
import qualified Data.Text.Encoding           as TE
import qualified Data.Text.Encoding.Error     as Err
import           Data.Text.IO                 as T
import           Data.Vector                  (Vector)
import qualified Data.Vector                  as V
import qualified Data.Vector.Algorithms.Intro as VI

import           System.Exit                  (exitFailure, exitSuccess)
import           System.IO                    (Handle, IOMode (..), hClose,
                                               openFile, hGetChar, hReady)

import           System.Console.ANSI

import           Score
import           Tty

{- TODO: implement something like withTty that handles restoring tty state -}
{- TODO: getting unweildy passing around currMatchCount and choicesToShow-}
{- TODO: look into resource monad for ensuring tty gets closed? -}
{- TODO: state monad? passing around matches and memo tables is cumbersome -}

prompt :: Text
prompt = "> "

data KeyPress = CtrlC
              | CtrlN
              | CtrlP
              | CtrlW
              | CtrlH
              | CtrlU
              | Enter
              | Invisible
              | Backspace
              | PlainChar Char

type Memo = M.Map Text (Vector Text)

data Search = Search
    { query     :: !Text
    , choices   :: !(Vector Text)
    , selection :: !Int
    , matches   :: !(Vector Text)
    }

instance Show Search where
  show _ = "[search]"

data SelecthState = SelecthState
    { s_search        :: !Search
    , s_choicesToShow :: !Int
    , s_scoreMemo     :: !Memo
    } deriving Show

data RenderedSearch = RenderedSearch
    { queryString   :: !Text
    , renderedLines :: !(Vector (Text, SGR))
    } deriving Show


data Action = SearchAction !SearchAction | ExitAction !ExitAction
  deriving Show


data SearchAction = Extend !Text
                  | DropWord
                  | DropChars !Int
                  | Clear
                  | SelectPrev
                  | SelectNext
                  | Ignore
  deriving Show

data ExitAction   = Abort
                  | MakeChoice !Search
  deriving Show

specialChars :: M.Map Char KeyPress
specialChars = M.fromList [ ('\ETX', CtrlC)
                          , ('\r',   Enter)
                          , ('\DEL', Backspace)
                          , ('\BS',  CtrlH)
                          , ('\NAK', CtrlU)
                          , ('\SO',  CtrlN)
                          , ('\DLE', CtrlP)
                          , ('\ETB', CtrlW)
                          ]

charToKeypress :: Char -> KeyPress
charToKeypress c = fromMaybe (if isPrint c then PlainChar c else Invisible)
                             (M.lookup c specialChars)

findMatches :: Memo -> Text -> Vector Text -> (Vector Text, Memo)
findMatches memo qry chs =
    case M.lookup qry memo of
       Just mtchs -> (mtchs, memo)
       Nothing    -> (newMatches, M.insert qry newMatches memo)
  where
    newMatches = getMatches qry chs
    getMatches q cs = V.map fst
                        . sortImmutableVec (flip compare `on` snd)
                        . V.filter (\(_,cScore) -> cScore > 0)
                        $ scoreAll q cs
    sortImmutableVec f v = runST $ do
                              vec <- V.thaw v
                              VI.sortBy f vec
                              V.freeze vec

render :: Search -> Int -> RenderedSearch
render (Search {query=q, matches=matched, selection=selIndex}) csToShow =
    RenderedSearch { queryString   = queryLine
                   , renderedLines = V.cons (queryLine, Reset)
                                            renderedMatchLines
                   }
  where
    renderedMatchLines = swapBackground selIndex
                           . V.map (\m -> (m, Reset))
                           $ matchLines
    matched'   = V.take csToShow matched
    matchLines = pad csToShow matched'
    queryLine  = prompt <> q
    swapBackground i vec =
      V.update
          vec
          (V.singleton
              (selIndex, (fst (vec V.! i), SetSwapForegroundBackground True)))
    pad n xs = xs <> V.replicate (n - V.length xs) " "

draw :: Handle -> RenderedSearch -> IO ()
draw tty rendered = do
    withCursorHidden tty $
        writeLines tty (renderedLines rendered)
    hSetCursorColumn tty . T.length $ queryString rendered

-- TODO: Handle exceptions, ensure tty gets restored to default settings
-- TODO: Ensure handle to tty is closed?  See what GB ensures.
{-restoreTTY :: IO ()-}

dropLastWord :: Text -> Text
dropLastWord = T.stripEnd . T.dropWhileEnd (not . isSpace)

-- TODO: ugh, this should indeed have used choices to show,
-- there can easily be 100s of matches.
writeSelection :: Handle -> Search -> IO ()
writeSelection tty (Search {matches=matches', selection=sel}) = do
    hCursorDown tty (V.length matches')
    T.hPutStr tty "\n"
    T.putStrLn $ fromMaybe (error "Failed to write selection.")
                           (matches' V.!? sel)

handleInput :: Text -> Search -> [Action]
handleInput inputText search = T.foldr buildActions [] inputText
  where
    buildActions c xs = actionFor c : xs
    actionFor c =
      case charToKeypress c of
          Enter     -> ExitAction (MakeChoice search)
          CtrlC     -> ExitAction Abort
          CtrlN     -> SearchAction SelectNext
          CtrlP     -> SearchAction SelectPrev
          CtrlU     -> SearchAction Clear
          CtrlW     -> SearchAction DropWord
          CtrlH     -> SearchAction (DropChars 1)
          Backspace -> SearchAction (DropChars 1)
          Invisible -> SearchAction Ignore
          PlainChar char -> SearchAction $ Extend (T.singleton char)

collapseActions :: [Action] -> [Action]
collapseActions = foldr collapse []
  where collapse (SearchAction (Extend t)) (SearchAction (Extend t'):as) =
          SearchAction (Extend (t <> t')) : as
        collapse (SearchAction (DropChars n)) (SearchAction (DropChars m):as) =
          SearchAction (DropChars $ n + m) : as
        collapse nextA as = nextA : as

buildSearch :: SearchAction -> Search -> Int -> Memo -> (Search, Memo)
buildSearch action search choicesToShow memo =
    case action of
        Extend qAddition -> (search', memo')
          where
            search' = search { query   = q'
                             , matches = matches'
                             , selection = 0 }
            q' = query search <> qAddition
            (matches', memo') = findMatches memo q' (matches search)
        DropWord    -> memoSearch $ dropLastWord (query search)
        DropChars n -> memoSearch $ T.dropEnd n (query search)
        Clear       -> if T.null (query search)
                       then (search, memo) -- why isn't this being hit?
                       else memoSearch ""
        Ignore     -> (search, memo)
        SelectNext -> moveSel 1
        SelectPrev -> moveSel (-1)
  where
    memoSearch q =  let (matches', memo') = findMatches memo q (choices search)
                    in  (search { query = q
                                , matches = matches'
                                , selection = 0}
                        , memo')
    moveSel n = ( search { selection = if possibleSelections > 0
                                       then mod (selection search + n)
{'fraction_non_alphanumeric': 0.08539239533990527, 'fraction_numerical': 0.0014082703879144795, 'mean_word_length': 1.5411515940143137, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 97, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '45759', 'n_tokens_mistral': 4206, 'n_tokens_neox': 3998, 'n_words': 1335}