{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module Blessings
    ( module Export
    , module Blessings
    ) where

import Blessings.Internal as Export (Blessable)
import qualified Blessings.Internal as Bless
import Control.Applicative
import Data.Ix (inRange)
import Data.List (genericDrop)
import Data.String
import Data.Tuple.Extra (both, first, second)
import Data.Word (Word8)
import qualified Prelude
import Prelude hiding (drop, length, take)

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


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["<>) . (<>"m") . Bless.intercalate ";" . map Bless.fromWord8 $ pm


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


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 ""


instance (Eq a, Blessable a) => Blessable (Blessings a) where
  length (Plain x) = Bless.length x
  length (SGR _ t) = Bless.length t
  length (Append t1 t2) = Bless.length t1 + Bless.length t2
  length Empty = 0

  drop n = \case
      Append t1 t2 ->
          let
            n1 = Bless.length (Bless.take n t1)
            n2 = n - n1
            t1' = Bless.drop n1 t1
            t2' = Bless.drop n2 t2
            isEmpty = (==0) . Bless.length . Bless.take 1
          in
            if n1 /= n || isEmpty t1'
              then t2'
              else Append t1' t2'
      Plain s ->
          normalize $ Plain (Bless.drop n s)
      SGR pm t ->
          normalize $ SGR pm (Bless.drop n t)
      Empty ->
          Empty

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

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

  intercalate i = \case
      [] -> mempty
      [t] -> t
      (t:ts) -> normalize $ t <> i <> Bless.intercalate i ts

  fromWord8 = Plain . Bless.fromWord8


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