{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

module Blessings where

import Control.Applicative
import Data.Ix (inRange)
import Data.List (genericDrop)
import Data.List qualified as L
import Data.MonoTraversable (Element, GrowingAppend, MonoFoldable(..), MonoFunctor(..), MonoPointed(..), MonoTraversable(..))
import Data.Sequences (Index, IsSequence, SemiSequence(..))
import Data.Sequences qualified as S
import Data.String
import Data.Tuple.Extra (both, first, second)
import Data.Word (Word8)


type Blessable t =
  ( Eq t
  , IsString t
  , Monoid t
  , IsSequence t
  , Index t ~ Int
  )

type Ps = Word8
type Pm = [Ps]

data Blessings a
    = Plain a
    | SGR Pm (Blessings a)
    | Append (Blessings a) (Blessings a)
    | Empty
  deriving (Eq, Show)


cataBlessings :: Monoid a => (a -> r) -> (Pm -> r -> r) -> (r -> r -> r) -> Blessings a -> r
cataBlessings plain sgr append = go
  where
    go (Plain s)      = plain s
    go (SGR pm t)     = sgr pm (go t)
    go (Append t1 t2) = append (go t1) (go t2)
    go Empty          = plain mempty


instance Foldable Blessings where
    foldMap f = \case
        Append t1 t2 -> foldMap f t1 <> foldMap f t2
        Plain s -> f s
        SGR _ t -> foldMap f t
        Empty -> mempty


instance Functor Blessings where
    fmap f = \case
        Append t1 t2 -> Append (fmap f t1) (fmap f t2)
        Plain s -> Plain (f s)
        SGR pm t -> SGR pm (fmap f t)
        Empty -> Empty


instance Semigroup (Blessings a) where
    t <> Empty = t
    Empty <> t = t
    Append t1 t2 <> t3 = Append t1 (t2 <> t3)
    t1 <> t2 = Append t1 t2


instance Monoid (Blessings a) where
    mempty = Empty
    mconcat = \case
      x:[] -> x
      x:xs -> Append x $ mconcat xs
      [] -> Empty


instance IsString a => IsString (Blessings a) where
    fromString = Plain . fromString


type instance Element (Blessings a) = a

instance MonoFoldable (Blessings a)

instance MonoPointed (Blessings a) where
    opoint = Plain

instance MonoFunctor (Blessings a) where
    omap = fmap

instance MonoTraversable (Blessings a) where
    otraverse f = \case
        Plain a       -> Plain <$> f a
        SGR pm t      -> SGR pm <$> otraverse f t
        Append t1 t2  -> Append <$> otraverse f t1 <*> otraverse f t2
        Empty         -> pure Empty

instance GrowingAppend (Blessings a)

instance Blessable a => SemiSequence (Blessings a) where
  type Index (Blessings a) = Int

  cons a b
    | a == mempty   = b
  cons a (Plain b)  = Plain (a <> b)
  cons a Empty      = Plain a
  cons a b          = Append (Plain a) b

  snoc b a
    | a == mempty   = b
  snoc (Plain b) a  = Plain (b <> a)
  snoc Empty a      = Plain a
  snoc b a          = Append b (Plain a)

  intersperse sep xs =
    case otoList xs of
      []     -> Empty
      (y:ys) ->
        foldl'
          (\acc z -> snoc (snoc acc sep) z)
          (Plain y)
          ys

  find p = ofoldr (\x acc -> if p x then Just x else acc) Nothing

  sortBy cmp xs =
    foldr cons Empty (L.sortBy cmp (otoList xs))

  reverse xs =
      foldr cons Empty (L.reverse (otoList xs))


class IsPm a where
    toPm :: a -> Pm
    fromPm :: Pm -> Maybe a


data FColor = ECMA48FColor Ps   -- ECMA-48 / ISO 6429 / ANSI X3.64
            | Xterm256FColor Ps
            | ISO8613_3FColor Ps Ps Ps
  deriving (Eq, Show)

instance IsPm FColor where
    toPm (ECMA48FColor i) = [i]
    toPm (Xterm256FColor i) = [38,5,i]
    toPm (ISO8613_3FColor r g b) = [38,2,r,g,b]
    fromPm = fromSGRPm SGRPm
               { def8Ps = 39
               , extPs = 38
               , lo8Ps = 30
               , hi8Ps = 37
               , makeECMA48Color = ECMA48FColor
               , makeXterm256Color = Xterm256FColor
               , makeISO8613_3Color = ISO8613_3FColor
               }
           . filterPm sgrBColor


data BColor = ECMA48BColor Ps
            | Xterm256BColor Ps
            | ISO8613_3BColor Ps Ps Ps
  deriving (Eq, Show)


instance IsPm BColor where
    toPm (ECMA48BColor i) = [i]
    toPm (Xterm256BColor i) = [48,5,i]
    toPm (ISO8613_3BColor r g b) = [48,2,r,g,b]
    fromPm = fromSGRPm SGRPm
                 { def8Ps = 49
                 , extPs = 48
                 , lo8Ps = 40
                 , hi8Ps = 47
                 , makeECMA48Color = ECMA48BColor
                 , makeXterm256Color = Xterm256BColor
                 , makeISO8613_3Color = ISO8613_3BColor
               }
           . filterPm sgrFColor


data Blink = Blink | NoBlink
  deriving (Eq, Show)

instance IsPm Blink where
    toPm Blink = [5]
    toPm NoBlink = [25]
    fromPm = rec . filterPm sgrColor
      where
        rec xs = case filter (`elem` ([5,25] :: [Word8])) xs of
            [] -> Nothing
            xs' -> case last xs' of
                5 -> Just Blink
                25 -> Just NoBlink
                _ -> error "filter broken in fromPm :: Pm -> Maybe Blink"


data Bold = Bold | NoBold
  deriving (Eq, Show)

instance IsPm Bold where
    toPm Bold = [1]
    toPm NoBold = [22]
    fromPm = rec . filterPm sgrColor
      where
        rec xs = case filter (`elem` ([1,22] :: [Word8])) xs of
            [] -> Nothing
            xs' -> case last xs' of
                1 -> Just Bold
                22 -> Just NoBold
                _ -> error "filter broken in fromPm :: Pm -> Maybe Bold"


data Underline = Underline | NoUnderline
  deriving (Eq, Show)

instance IsPm Underline where
    toPm Underline = [4]
    toPm NoUnderline = [24]
    fromPm = rec . filterPm sgrColor
      where
        rec xs = case filter (`elem` ([4,24] :: [Word8])) xs of
            [] -> Nothing
            xs' -> case last xs' of
                4 -> Just Underline
                24 -> Just NoUnderline
                _ -> error "filter broken in fromPm :: Pm -> Maybe Underline"


data SGRPm c = SGRPm
    { def8Ps :: Ps
    , extPs :: Ps
    , lo8Ps :: Ps
    , hi8Ps :: Ps
    , makeECMA48Color :: Ps -> c
    , makeXterm256Color :: Ps -> c
    , makeISO8613_3Color :: Ps -> Ps -> Ps -> c
    }


fromSGRPm :: IsPm c => SGRPm c -> Pm -> Maybe c
fromSGRPm SGRPm{..} = rec Nothing
  where
    rec mb_c (x:xs)
        | x == extPs = case xs of
            (2:r:g:b:xs') -> rec (Just $ makeISO8613_3Color r g b) xs'
            (5:i:xs')     -> rec (Just $ makeXterm256Color i) xs'
            _             -> rec mb_c xs
        | x == def8Ps = rec (Just $ makeECMA48Color def8Ps) xs
        | inRange (lo8Ps, hi8Ps) x = rec (Just $ makeECMA48Color x) xs
        | otherwise = rec mb_c xs
    rec mb_c _ = mb_c


-- filterPm is used to preprocess Pm before searching with fromPm in
-- order to remove (longer) sequences that could contain subsequences
-- that look like the (shorter) sequences we're searching.
-- E.g. we could find [1] (bold) in any extended color sequence.
-- TODO Can we combine this whole from*Pm with Scanner?
filterPm :: (Pm -> Maybe Word8) -> Pm -> Pm
filterPm f = rec []
  where
    rec ys xs@(xhead:xtail) = maybe (rec (ys ++ [xhead]) xtail)
                                    (rec ys . flip genericDrop xs)
                                    (f xs)
    rec ys _ = ys

sgrColor, sgrFColor, sgrBColor :: Pm -> Maybe Word8

sgrColor xs = sgrFColor xs <|> sgrBColor xs

sgrFColor (38:5:_) = Just 3
sgrFColor (38:2:_) = Just 5
sgrFColor _ = Nothing

sgrBColor (48:5:_) = Just 3
sgrBColor (48:2:_) = Just 5
sgrBColor _ = Nothing


data Style = Style
  { fg        :: FColor
  , bg        :: BColor
  , blink     :: Blink
  , bold      :: Bold
  , underline :: Underline
  }
  deriving (Eq, Show)

defaultStyle :: Style
defaultStyle =
  Style
    { fg        = ECMA48FColor 39
    , bg        = ECMA48BColor 49
    , blink     = NoBlink
    , bold      = NoBold
    , underline = NoUnderline
    }


--------------------------------------------------------------------------------
-- Semantics
--------------------------------------------------------------------------------

type ColoredChar a = (Style, Element a)

sem :: Blessable a => Blessings a -> [ColoredChar a]
sem = semWith defaultStyle

semWith :: Blessable a  => Style -> Blessings a -> [ColoredChar a]
semWith st = \case
  Empty         -> []
  Plain t       -> [ (st, c) | c <- otoList t ]
  Append a b    -> semWith st a ++ semWith st b
  SGR pm a ->
    let st' = applyPm st pm
    in semWith st' a


--------------------------------------------------------------------------------
-- SGR interpretation
--------------------------------------------------------------------------------

-- apply a full SGR list to a Style
applyPm :: Style -> [Word8] -> Style
applyPm st pm = go st pm
  where
    go s [] = s

    -- reset
    go _ (0:rest) =
      go defaultStyle rest

    -- bold on/off
    go s (1:rest)  = go s{ bold = Bold } rest
    go s (22:rest) = go s{ bold = NoBold } rest

    -- underline on/off
    go s (4:rest)  = go s{ underline = Underline } rest
    go s (24:rest) = go s{ underline = NoUnderline } rest

    -- blink on/off
    go s (5:rest)  = go s{ blink = Blink } rest
    go s (25:rest) = go s{ blink = NoBlink } rest

    -- 8-color fg
    go s (c:rest)
      | 30 <= c && c <= 37 || c == 39 =
          go s{ fg = ECMA48FColor c } rest

    -- 8-color bg
    go s (c:rest)
      | 40 <= c && c <= 47 || c == 39 =
          go s{ bg = ECMA48BColor c } rest

    -- xterm-256 fg: 38;5;i
    go s (38:5:i:rest) =
      go s{ fg = Xterm256FColor i } rest

    -- xterm-256 bg: 48;5;i
    go s (48:5:i:rest) =
      go s{ bg = Xterm256BColor i } rest

    -- truecolor fg: 38;2;r;g;b
    go s (38:2:r:g:b:rest) =
      go s{ fg = ISO8613_3FColor r g b } rest

    -- truecolor bg: 48;2;r;g;b
    go s (48:2:r:g:b:rest) =
      go s{ bg = ISO8613_3BColor r g b } rest

    -- anything else / incomplete sequences: skip
    go s (_:rest) = go s rest


pmHasVisibleEffect :: Style -> [Word8] -> Bool
pmHasVisibleEffect st pm =
  applyPm st pm /= st


--------------------------------------------------------------------------------
-- Normalizer
--------------------------------------------------------------------------------

normalize :: Blessable a => Blessings a -> Blessings a
normalize = fromSem . sem

fromSem :: Blessable a => [ColoredChar a] -> Blessings a
fromSem [] = Empty
fromSem cs =
  foldr1 Append
    [ chunkToBlessings st (S.fromList s)
    | (st, s) <- chunks cs
    , not (null s)
    ]

chunks :: Eq t => [(t, a)] -> [(t, [a])]
chunks [] = []
chunks ((st0,c0):rest) = go st0 [c0] rest
  where
    go curSt acc [] =
      [(curSt, L.reverse acc)]
    go curSt acc ((st,c):xs)
      | st == curSt = go curSt (c:acc) xs
      | otherwise   = (curSt, L.reverse acc) : go st [c] xs

chunkToBlessings :: Blessable a => Style -> a -> Blessings a
chunkToBlessings st s
  | s == mempty        = Empty
  | st == defaultStyle = Plain s
  | otherwise          = SGR (styleToPm st) (Plain s)

styleToPm :: Style -> [Word8]
styleToPm st =
  fgCodes ++ bgCodes ++ blinkCodes ++ boldCodes ++ underlineCodes
  where
    fgCodes =
      case fg st of
        ECMA48FColor c          -> [c]
        Xterm256FColor i        -> [38,5,i]
        ISO8613_3FColor r g b   -> [38,2,r,g,b]

    bgCodes =
      case bg st of
        ECMA48BColor c          -> [c]
        Xterm256BColor i        -> [48,5,i]
        ISO8613_3BColor r g b   -> [48,2,r,g,b]

    blinkCodes =
      case blink st of
        Blink   -> [5]
        NoBlink -> []

    boldCodes =
      case bold st of
        Bold   -> [1]
        NoBold -> []

    underlineCodes =
      case underline st of
        Underline   -> [4]
        NoUnderline -> []


--------------------------------------------------------------------------------
-- Renderer
--------------------------------------------------------------------------------

type RenderState = [(FColor, BColor, Blink, Bold, Underline)]


emptyRenderState :: RenderState
emptyRenderState = [(ECMA48FColor 39, ECMA48BColor 49, NoBlink, NoBold, NoUnderline)]


render :: Blessable a => RenderState -> Blessings a -> a -> a

render _ (Plain s) y = s <> y

-- TODO merge successive sequences: \ESC[32m\ESC[1m -> \ESC[31;1m
render rs@((fc, bc, bl, b, u):_) (SGR c t) y =
    renderSGR bra <> render rs' t (renderSGR ket <> y)
  where
    fc' = maybe fc id $ fromPm c
    bc' = maybe bc id $ fromPm c
    bl' = maybe bl id $ fromPm c
    b'  = maybe  b id $ fromPm c
    u'  = maybe  u id $ fromPm c
    rs' = (fc', bc', bl', b', u') : rs
    bra = braket >>= fst
    ket = braket >>= snd
    braket =
        (if fc' /= fc then (toPm fc', toPm fc) else ([],[])) :
        (if bc' /= bc then (toPm bc', toPm bc) else ([],[])) :
        (if bl' /= bl then (toPm bl', toPm bl) else ([],[])) :
        (if b'  /=  b then (toPm  b', toPm  b) else ([],[])) :
        (if u'  /=  u then (toPm  u', toPm  u) else ([],[])) : []

render _ (SGR _ _) _ =
    error "render called w/o proper initial state"
    -- where a proper initial state is s.th. like emptyRenderState

render r (Append t1 t2) y =
    render r t1 $ render r t2 y

render _ Empty y = y


renderSGR :: Blessable a => Pm -> a
renderSGR [] = mempty
renderSGR pm =
  "\ESC[" <> mconcat (L.intersperse ";" (map (fromString . show) pm)) <> "m"


pp :: Blessable a => Blessings a -> a
pp t = render emptyRenderState t mempty


instance Blessable a => IsSequence (Blessings a) where
  lengthIndex = ofoldl' (\acc w -> acc + S.lengthIndex w) 0

  drop n = \case
      Append t1 t2 ->
          let
            n1 = S.lengthIndex (S.take n t1)
            n2 = n - n1
            t1' = S.drop n1 t1
            t2' = S.drop n2 t2
          in
            Append t1' t2'
      Plain s ->
          Plain (S.drop n s)
      SGR pm t ->
          SGR pm (S.drop n t)
      Empty ->
          Empty

  take n = \case
      Append t1 t2 ->
          let
            t1' = S.take n t1
            n' = n - S.lengthIndex t1'
          in
              if n' > 0
                then t1' <> S.take n' t2
                else t1'
      Plain s ->
          Plain (S.take n s)
      SGR pm t ->
          SGR pm (S.take n t)
      Empty ->
          Empty

  splitAt n = \case
      Append t1 t2 ->
          let
            nt1 = S.lengthIndex t1
          in
            if n <= nt1
              then second (<>t2) $ S.splitAt n t1
              else first (t1<>) $ S.splitAt (n - nt1) t2
      Plain s ->
        both Plain $ S.splitAt n s
      SGR pm t ->
        both (SGR pm) $ S.splitAt n t
      Empty ->
        (Empty, Empty)

  break p = \case
      Append t1 t2 ->
          case S.break p t1 of
            (t1l, t1r)
              | t1r == mempty -> first (t1l<>) $ S.break p t2
              | otherwise -> (t1l, t1r <> t2)
      Plain s
        | p s       -> (Empty, Plain s)
        | otherwise -> (Plain s, Empty)
      SGR pm t ->
        both (SGR pm) $ S.break p t
      Empty ->
        (Empty, Empty)


chunksOf :: Blessable a => Int -> a -> [a]
chunksOf k = rec
  where
    rec t =
        case S.splitAt k t of
          (tl, tr)
            | tl == mempty -> []
            | otherwise -> tl : rec tr