-- This file is part of Diohsc
-- Copyright (C) 2020 Martin Bays <mbays@sdf.org>
--
-- This program is free software: you can redistribute it and/or modify
-- it under the terms of version 3 of the GNU General Public License as
-- published by the Free Software Foundation, or any later version.
--
-- You should have received a copy of the GNU General Public License
-- along with this program.  If not, see http://www.gnu.org/licenses/.

{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Safe              #-}

-- Basic ansi attributes, using only most widely supported ansi terminal codes
module ANSIColour
    ( applyIf
    , resetCode
    , withColour
    , withBold
    , withReverse
    , withUnderline
    , withColourStr
    , withBoldStr
    , withReverseStr
    , withUnderlineStr
    , stripCSI
    , visibleLength
    , splitAtVisible
    , escapePromptCSI
    , sanitiseNonCSI
    , Colour(..)
    ) where

import           Control.Exception.Base (bracket_)

import qualified Data.Text.Lazy         as T
import qualified Data.Text.Lazy.IO      as T

import           MetaString
import           WCWidth

data Colour = Black | Red | Green | Yellow
    | Blue | Magenta | Cyan | White
    | BoldBlack | BoldRed | BoldGreen | BoldYellow
    | BoldBlue | BoldMagenta | BoldCyan | BoldWhite
    deriving (Eq,Ord,Show,Read)

resetCode, boldCode, unboldCode, reverseCode,
    unreverseCode, underlineCode , ununderlineCode, resetColourCode
    :: MetaString a => a
resetCode = "\ESC[0m"
boldCode = "\ESC[1m"
underlineCode = "\ESC[4m"
reverseCode = "\ESC[7m"
unboldCode = "\ESC[22m"
ununderlineCode = "\ESC[24m"
unreverseCode = "\ESC[27m"
resetColourCode = "\ESC[39m\ESC[22m"

colourCode :: MetaString a => Colour -> a
colourCode c = (if isBold c then boldCode else "") <> "\ESC[3" <> fromString (colNum c) <> "m"
    where
    isBold = flip elem [BoldBlack, BoldRed, BoldGreen, BoldYellow,
        BoldBlue, BoldMagenta, BoldCyan, BoldWhite]
    colNum Black       = "0"
    colNum Red         = "1"
    colNum Green       = "2"
    colNum Yellow      = "3"
    colNum Blue        = "4"
    colNum Magenta     = "5"
    colNum Cyan        = "6"
    colNum White       = "7"
    colNum BoldBlack   = "0"
    colNum BoldRed     = "1"
    colNum BoldGreen   = "2"
    colNum BoldYellow  = "3"
    colNum BoldBlue    = "4"
    colNum BoldMagenta = "5"
    colNum BoldCyan    = "6"
    colNum BoldWhite   = "7"

withStyle :: T.Text -> T.Text -> IO a -> IO a
withStyle c r = T.putStr c `bracket_` T.putStr r
withColour :: Colour -> IO a -> IO a
withColour c = withStyle (colourCode c) resetColourCode
withBold, withReverse, withUnderline :: IO a -> IO a
withBold = withStyle boldCode unboldCode
withReverse = withStyle reverseCode unreverseCode
withUnderline = withStyle underlineCode ununderlineCode

withStyleStr :: MetaString a => a -> a -> a -> a
withStyleStr c r s = c <> s <> r
withColourStr :: MetaString a => Colour -> a -> a
withColourStr c = withStyleStr (colourCode c) resetColourCode
withBoldStr, withReverseStr, withUnderlineStr :: MetaString a => a -> a
withBoldStr = withStyleStr boldCode unboldCode
withReverseStr = withStyleStr reverseCode unreverseCode
withUnderlineStr = withStyleStr underlineCode ununderlineCode

-- |"applyIf cond f" is shorthand for "if cond then f else id"
applyIf :: Bool -> (a -> a) -> (a -> a)
applyIf True  = id
applyIf False = const id


endCSI :: Char -> Bool
endCSI c = '@' <= c && c <= '~'

-- |strip all CSI escape sequences
stripCSI :: T.Text -> T.Text
stripCSI s =
    let (pre,post) = T.breakOn "\ESC[" s
    in if T.null post then pre
        else (pre <>) . stripCSI . T.drop 1 .
            T.dropWhile (not . endCSI) $ T.drop 2 post

visibleLength :: (Integral i) => T.Text -> i
visibleLength = fromIntegral . wcLength . stripCSI

wcLength :: T.Text -> Int
wcLength = sum . (max 0 . wcwidth <$>) . T.unpack

splitAtWC :: Int -> T.Text -> (T.Text,T.Text)
splitAtWC m = go m T.empty where
    go !n !acc t
        | Just (c,r) <- T.uncons t = let w = max 0 $ wcwidth c in
            if w > max 0 n then (T.reverse acc,t)
            else go (n - w) (T.cons c acc) r
        | otherwise = (T.reverse acc, T.empty)


splitAtVisible :: (Integral i) => i -> T.Text -> (T.Text,T.Text)
splitAtVisible n t =
    let (pre,post) = T.breakOn "\ESC[" t
        n' = fromIntegral n
        (a,b) = splitAtWC n' pre
        catFst s (s',s'') = (s<>s',s'')
    in a `catFst`
        if not (T.null b) || T.null post then ("",b<>post)
        else
            let (s,r) = T.splitAt 2 post
                (s',r') = T.break endCSI r
                (s'',rest) = T.splitAt 1 r'
                csi = s <> s' <> s''
            in csi `catFst` splitAtVisible (n' - wcLength a) rest

-- |sanitise non-CSI escape sequences by turning \ESC into \\ESC
-- (buggy terminals make these sequences a potential security hole;
-- see e.g. https://nvd.nist.gov/vuln/detail/CVE-2020-9366 )
sanitiseNonCSI :: T.Text -> T.Text
sanitiseNonCSI s =
    let (pre,post) = T.breakOn "\ESC" s
    in if T.null post then pre else
        let post' = T.drop 1 post
            isCSI = T.take 1 post' == "["
        in pre <> (if isCSI then "\ESC" else "\\ESC") <> sanitiseNonCSI post'

-- |append \STX to each CSI sequence, as required in Haskeline prompts.
-- See https://github.com/judah/haskeline/wiki/ControlSequencesInPrompt
escapePromptCSI :: String -> String
escapePromptCSI s = case break (== '\ESC') s of
    (pre,'\ESC':'[':post) -> ((pre <> "\ESC[") <>) $
        case break endCSI post of
            (pre',e:post') -> (pre' <>) $ e : '\STX' : escapePromptCSI post'
            (pre',[])      -> pre'
    (pre,[]) -> pre
    (pre,e:post) -> (pre <>) $ e : escapePromptCSI post