{-# 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)


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


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"


stripSGR :: Blessings a -> Blessings a
stripSGR = \case
    Append t1 t2 -> Append (stripSGR t1) (stripSGR t2)
    SGR _ t -> stripSGR t
    Plain x -> Plain x
    Empty -> Empty


normalize' :: (Eq a, Monoid a) => Blessings a -> Blessings a -> Blessings a
normalize' t t' =
    if t' == t
      then t
      else normalize t'


normalize :: (Eq a, Monoid a) => Blessings a -> Blessings a
normalize = \case
    Append Empty t -> normalize t
    Append t Empty -> normalize t
    Append (Append t1 t2) t3 -> normalize $ Append t1 (Append t2 t3)
    Append (Plain s1) (Plain s2) -> normalize $ Plain (s1 <> s2)

    Append (Plain s1) (Append (Plain s2) t1) ->
      normalize (Append (Plain (s1 <> s2)) t1)

    t@(Append t1@(SGR pm1 t11) t2@(Append (SGR pm2 t21) t22)) ->
      let
        pm1' = normalizePm pm1
        pm2' = normalizePm pm2
      in
        if pm1' == pm2'
          then normalize (Append (SGR pm1 (Append t11 t21)) t22)
          else normalize' t $ Append (normalize t1) (normalize t2)

    t@(Append t1@(SGR pm1 t11) t2@(SGR pm2 t21)) ->
      let
        pm1' = normalizePm pm1
        pm2' = normalizePm pm2
      in
        if pm1' == pm2'
          then normalize (SGR pm1' (Append t11 t21))
          else normalize' t $ Append (normalize t1) (normalize t2)

    t@(Append t1 t2) -> normalize' t $ Append (normalize t1) (normalize t2)
    SGR _ Empty -> Empty
    SGR [] t -> normalize t
    t@(SGR pm t1) -> normalize' t $ SGR (normalizePm pm) (normalize t1)
    Plain x | x == mempty -> Empty
    t@(Plain _) -> t
    Empty -> Empty


normalizeHead :: (Eq a, Monoid a) => Blessings a -> Blessings a
normalizeHead = \case
    Append Empty t -> normalizeHead t
    Append t1 t2 ->
      let
        t1' = normalizeHead t1
      in
      if t1' == Empty
        then normalizeHead t2
        else Append t1' t2
    SGR _ Empty -> Empty
    SGR [] t -> normalizeHead t
    SGR pm t ->
      let
        pm' = normalizePm pm
        t' = normalizeHead t
      in
      if pm' == []
        then t'
        else SGR pm' t'
    Plain x | x == mempty -> Empty
    t@(Plain _) -> t
    Empty -> Empty


data NormalizedPm
    = NormalizedPm
        { foregroundColor :: [Word8]
        , backgroundColor :: [Word8]
        }

emptyNormalizedPm :: NormalizedPm
emptyNormalizedPm
    = NormalizedPm
        { foregroundColor = []
        , backgroundColor = []
        }

normalizePm :: [Word8] -> [Word8]
normalizePm pm0 =
    collectEffective emptyNormalizedPm $ skipCanceled pm0 pm0
  where
    collectEffective p = \case
        -- direct-color
        (38 : 2 : r : g : b : pm) -> collectEffective (p { foregroundColor = [38, 2, r, g, b] }) pm
        (48 : 2 : r : g : b : pm) -> collectEffective (p { backgroundColor = [48, 2, r, g, b] }) pm
        -- indexed-color
        (38 : 5 : i : pm) -> collectEffective (p { foregroundColor = [38, 5, i] }) pm
        (48 : 5 : i : pm) -> collectEffective (p { backgroundColor = [48, 5, i] }) pm
        (ps : pm)
          -- 8-color (must be analyzed after direct- and indexed-colors)
          | 30 <= ps && ps <= 39 -> collectEffective (p { foregroundColor = [ps] }) pm
          | 40 <= ps && ps <= 49 -> collectEffective (p { backgroundColor = [ps] }) pm
          -- ignore everything else
          | otherwise -> ps : collectEffective p pm
        [] -> foregroundColor p <> backgroundColor p

    skipCanceled xs = \case
        (38 : 2 : _ : _ : _ : pm) -> skipCanceled xs pm
        (38 : 5 : _ : pm) -> skipCanceled xs pm
        xs'@(0 : pm) -> skipCanceled xs' pm
        (_ : pm) -> skipCanceled xs pm
        [] -> xs


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
            normalizeHead $ Append t1' t2'
      Plain s ->
          normalizeHead $ Plain (S.drop n s)
      SGR pm t ->
          normalizeHead $ 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
            normalizeHead $
              if n' > 0
                then t1' <> S.take n' t2
                else t1'
      Plain s ->
          normalizeHead $ Plain (S.take n s)
      SGR pm t ->
          normalizeHead $ SGR pm (S.take n t)
      Empty ->
          Empty

  splitAt n = \case
      Append t1 t2 ->
        both normalizeHead $
          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 (normalizeHead . Plain) $ S.splitAt n s
      SGR pm t ->
        both (normalizeHead . SGR pm) $ S.splitAt n t
      Empty ->
        (Empty, Empty)

  break p = \case
      Append t1 t2 ->
        both normalizeHead $
          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 (normalizeHead . 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