{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PackageImports        #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE UndecidableInstances  #-}

module Data.String.Interpolate.Conversion
  ( IsCustomSink, InterpSink(..), Interpolatable(..)
  , SpaceChompable(..)
  , bsToTextBuilder, lbsToTextBuilder, encodeCharUTF8
  )
where

import Data.Char             ( isSpace )
import Data.Int              ( Int64 )
import Data.Kind             ( Type )
import Data.Proxy
import Data.String           ( IsString, fromString )
import Data.Text.Conversions

import qualified Data.ByteString         as B
import qualified Data.ByteString.Builder as LB
import qualified Data.ByteString.Lazy    as LB
import qualified Data.Text               as T
import qualified Data.Text.Lazy          as LT hiding ( singleton )
import qualified Data.Text.Lazy.Builder  as LT

import qualified "utf8-string" Data.ByteString.Lazy.UTF8 as LUTF8
import qualified "utf8-string" Data.ByteString.UTF8      as UTF8

import "base" Control.Category ( (>>>) )

-- Remove some imports above GHC 8.8.X
#if MIN_VERSION_base(4,13,0)
#else
import "base" Data.Monoid ( (<>) )
import "base" Text.Show   ( ShowS, showChar, showString )
#endif

-- |
-- We wrap the builders in B so that we can add a phantom type parameter.
-- This gives the inner `interpolate's enough information to know where
-- they're going and pick an instance, forcing all the types into lockstep.
newtype B dst a = B { unB :: a }
  deriving (Eq, Show)

-- | Does this type require special behavior when something is interpolated /into/ it?
type family IsCustomSink dst where
  IsCustomSink T.Text = 'True
  IsCustomSink LT.Text = 'True
  IsCustomSink LT.Builder = 'True
  IsCustomSink B.ByteString = 'True
  IsCustomSink LB.ByteString = 'True
  IsCustomSink LB.Builder = 'True
  IsCustomSink _ = 'False

-- | Something that can be interpolated into.
class IsCustomSink dst ~ flag => InterpSink (flag :: Bool) dst where
  type Builder flag dst :: Type

  -- | Meant to be used only for verbatim parts of the interpolation.
  ofString :: Proxy flag -> String -> B dst (Builder flag dst)
  -- |
  -- `build' should be 'in-order'; that is, the left builder comes from
  -- a string on the left, and the right builder comes from a string on the right.
  build :: Proxy flag -> B dst (Builder flag dst) -> B dst (Builder flag dst) -> B dst (Builder flag dst)
  finalize :: Proxy flag -> B dst (Builder flag dst) -> dst

-- |
-- Represents that we can interpolate objects of type src into a an
-- interpolation string that returns type dst.
class InterpSink flag dst => Interpolatable (flag :: Bool) src dst where
  interpolate :: Proxy flag -> src -> B dst (Builder flag dst)

-- |
-- We can collapse whitespace in the given type.
class SpaceChompable a where
  chompSpaces :: a -> a

instance (IsCustomSink str ~ 'False, IsString str) => InterpSink 'False str where
  type Builder 'False str = ShowS

  ofString _ = B . showString
  build _ (B f) (B g) = B $ f . g
  finalize _ = fromString . ($ "") . unB

instance InterpSink 'True T.Text where
  type Builder 'True T.Text = LT.Builder

  ofString _ = B . LT.fromString
  build _ (B l) (B r) = B $ l <> r
  finalize _ = LT.toStrict . LT.toLazyText . unB

instance InterpSink 'True LT.Text where
  type Builder 'True LT.Text = LT.Builder

  ofString _ = B . LT.fromString
  build _ (B l) (B r) = B $ l <> r
  finalize _ = LT.toLazyText . unB

instance InterpSink 'True LT.Builder where
  type Builder 'True LT.Builder = LT.Builder

  ofString _ = B . LT.fromString
  build _ (B l) (B r) = B $ l <> r
  finalize _ = unB

instance InterpSink 'True B.ByteString where
  type Builder 'True B.ByteString = LB.Builder

  ofString _ = B . LB.byteString . unUTF8 . convertText
  build _ (B l) (B r) = B $ l <> r
  finalize _ = LB.toStrict . LB.toLazyByteString . unB

instance InterpSink 'True LB.ByteString where
  type Builder 'True LB.ByteString = LB.Builder

  ofString _ = B . LB.lazyByteString . unUTF8 . convertText
  build _ (B l) (B r) = B $ l <> r
  finalize _ = LB.toLazyByteString . unB

instance InterpSink 'True LB.Builder where
  type Builder 'True LB.Builder = LB.Builder

  ofString _ = B . LB.lazyByteString . unUTF8 . convertText
  build _ (B l) (B r) = B $ l <> r
  finalize _ = unB

instance {-# OVERLAPPABLE #-} (Show src, IsString dst, IsCustomSink dst ~ 'False) => Interpolatable 'False src dst where
  interpolate _ = B . shows
instance {-# OVERLAPS #-} (IsString dst, IsCustomSink dst ~ 'False) => Interpolatable 'False Char dst where
  interpolate _ = B . showChar
instance {-# OVERLAPS #-} (IsString dst, IsCustomSink dst ~ 'False) => Interpolatable 'False String dst where
  interpolate _ = B . showString
instance {-# OVERLAPS #-} (IsString dst, IsCustomSink dst ~ 'False) => Interpolatable 'False T.Text dst where
  interpolate _ = B . showString . T.unpack
instance {-# OVERLAPS #-} (IsString dst, IsCustomSink dst ~ 'False) => Interpolatable 'False LT.Text dst where
  interpolate _ = B . showString . LT.unpack
instance {-# OVERLAPS #-} (IsString dst, IsCustomSink dst ~ 'False) => Interpolatable 'False LT.Builder dst where
  interpolate _ = B . showString . LT.unpack . LT.toLazyText
instance {-# OVERLAPS #-} (IsString dst, IsCustomSink dst ~ 'False) => Interpolatable 'False B.ByteString dst where
  interpolate _ = B . showString . UTF8.toString
instance {-# OVERLAPS #-} (IsString dst, IsCustomSink dst ~ 'False) => Interpolatable 'False LB.ByteString dst where
  interpolate _ = B . showString . LUTF8.toString
instance {-# OVERLAPS #-} (IsString dst, IsCustomSink dst ~ 'False) => Interpolatable 'False LB.Builder dst where
  interpolate _ = B . showString . LUTF8.toString . LB.toLazyByteString

instance {-# OVERLAPPABLE #-} Show src => Interpolatable 'True src T.Text where
  interpolate _ = B . LT.fromString . show
instance {-# OVERLAPS #-} Interpolatable 'True Char T.Text where
  interpolate _ = B . LT.singleton
instance {-# OVERLAPS #-} Interpolatable 'True String T.Text where
  interpolate _ = B . LT.fromString
instance {-# OVERLAPS #-} Interpolatable 'True T.Text T.Text where
  interpolate _ = B . LT.fromText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Text T.Text where
  interpolate _ = B . LT.fromLazyText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Builder T.Text where
  interpolate _ = B
instance {-# OVERLAPS #-} Interpolatable 'True B.ByteString T.Text where
  interpolate _ = B . bsToTextBuilder
instance {-# OVERLAPS #-} Interpolatable 'True LB.ByteString T.Text where
  interpolate _ = B . lbsToTextBuilder
instance {-# OVERLAPS #-} Interpolatable 'True LB.Builder T.Text where
  interpolate _ = B . lbsToTextBuilder . LB.toLazyByteString

instance {-# OVERLAPPABLE #-} Show src => Interpolatable 'True src LT.Text where
  interpolate _ = B . LT.fromString . show
instance {-# OVERLAPS #-} Interpolatable 'True Char LT.Text where
  interpolate _ = B . LT.singleton
instance {-# OVERLAPS #-} Interpolatable 'True String LT.Text where
  interpolate _ = B . LT.fromString
instance {-# OVERLAPS #-} Interpolatable 'True T.Text LT.Text where
  interpolate _ = B . LT.fromText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Text LT.Text where
  interpolate _ = B . LT.fromLazyText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Builder LT.Text where
  interpolate _ = B
instance {-# OVERLAPS #-} Interpolatable 'True B.ByteString LT.Text where
  interpolate _ = B . bsToTextBuilder
instance {-# OVERLAPS #-} Interpolatable 'True LB.ByteString LT.Text where
  interpolate _ = B . lbsToTextBuilder
instance {-# OVERLAPS #-} Interpolatable 'True LB.Builder LT.Text where
  interpolate _ = B . lbsToTextBuilder . LB.toLazyByteString

instance {-# OVERLAPPABLE #-} Show src => Interpolatable 'True src LT.Builder where
  interpolate _ = B . LT.fromString . show
instance {-# OVERLAPS #-} Interpolatable 'True Char LT.Builder where
  interpolate _ = B . LT.singleton
instance {-# OVERLAPS #-} Interpolatable 'True String LT.Builder where
  interpolate _ = B . LT.fromString
instance {-# OVERLAPS #-} Interpolatable 'True T.Text LT.Builder where
  interpolate _ = B . LT.fromText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Text LT.Builder where
  interpolate _ = B . LT.fromLazyText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Builder LT.Builder where
  interpolate _ = B
instance {-# OVERLAPS #-} Interpolatable 'True B.ByteString LT.Builder where
  interpolate _ = B . bsToTextBuilder
instance {-# OVERLAPS #-} Interpolatable 'True LB.ByteString LT.Builder where
  interpolate _ = B . lbsToTextBuilder
instance {-# OVERLAPS #-} Interpolatable 'True LB.Builder LT.Builder where
  interpolate _ = B . lbsToTextBuilder . LB.toLazyByteString

instance {-# OVERLAPPABLE #-} Show src => Interpolatable 'True src B.ByteString where
  interpolate _ = B . LB.byteString . unUTF8 . convertText . show
instance {-# OVERLAPS #-} Interpolatable 'True Char B.ByteString where
  interpolate _ = B . encodeCharUTF8
instance {-# OVERLAPS #-} Interpolatable 'True String B.ByteString where
  interpolate _ = B . LB.byteString . unUTF8 . convertText
instance {-# OVERLAPS #-} Interpolatable 'True T.Text B.ByteString where
  interpolate _ = B . LB.byteString . unUTF8 . convertText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Text B.ByteString where
  interpolate _ = B . LB.byteString . unUTF8 . convertText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Builder B.ByteString where
  interpolate _ = B . LB.byteString . unUTF8 . convertText . LT.toLazyText
instance {-# OVERLAPS #-} Interpolatable 'True B.ByteString B.ByteString where
  interpolate _ = B . LB.byteString
instance {-# OVERLAPS #-} Interpolatable 'True LB.ByteString B.ByteString where
  interpolate _ = B . LB.lazyByteString
instance {-# OVERLAPS #-} Interpolatable 'True LB.Builder B.ByteString where
  interpolate _ = B

instance {-# OVERLAPPABLE #-} Show src => Interpolatable 'True src LB.ByteString where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText . show
instance {-# OVERLAPS #-} Interpolatable 'True Char LB.ByteString where
  interpolate _ = B . encodeCharUTF8
instance {-# OVERLAPS #-} Interpolatable 'True String LB.ByteString where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText
instance {-# OVERLAPS #-} Interpolatable 'True T.Text LB.ByteString where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Text LB.ByteString where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Builder LB.ByteString where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText . LT.toLazyText
instance {-# OVERLAPS #-} Interpolatable 'True B.ByteString LB.ByteString where
  interpolate _ = B . LB.byteString
instance {-# OVERLAPS #-} Interpolatable 'True LB.ByteString LB.ByteString where
  interpolate _ = B . LB.lazyByteString
instance {-# OVERLAPS #-} Interpolatable 'True LB.Builder LB.ByteString where
  interpolate _ = B

instance {-# OVERLAPPABLE #-} Show src => Interpolatable 'True src LB.Builder where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText . show
instance {-# OVERLAPS #-} Interpolatable 'True Char LB.Builder where
  interpolate _ = B . encodeCharUTF8
instance {-# OVERLAPS #-} Interpolatable 'True String LB.Builder where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText
instance {-# OVERLAPS #-} Interpolatable 'True T.Text LB.Builder where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Text LB.Builder where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText
instance {-# OVERLAPS #-} Interpolatable 'True LT.Builder LB.Builder where
  interpolate _ = B . LB.lazyByteString . unUTF8 . convertText . LT.toLazyText
instance {-# OVERLAPS #-} Interpolatable 'True B.ByteString LB.Builder where
  interpolate _ = B . LB.byteString
instance {-# OVERLAPS #-} Interpolatable 'True LB.ByteString LB.Builder where
  interpolate _ = B . LB.lazyByteString
instance {-# OVERLAPS #-} Interpolatable 'True LB.Builder LB.Builder where
  interpolate _ = B

instance {-# OVERLAPPABLE #-} (Show a, IsString a) => SpaceChompable a where
  chompSpaces = fromString . chompSpaces . show
instance {-# OVERLAPS #-} SpaceChompable String where
  chompSpaces = unwords . words
instance {-# OVERLAPS #-} SpaceChompable T.Text where
  chompSpaces = T.unwords . T.words
instance {-# OVERLAPS #-} SpaceChompable LT.Text where
  chompSpaces = LT.unwords . LT.words

-- | For storing state while we fold over the ByteString.
data BSChomper = BSChomper
  { bscNumWS   :: !Int64
  , bscBuilder :: !(Maybe LB.Builder)  -- ^ We use Maybe here to know if we've processed
                                       --   non-whitespace characters yet.
  }

chompBS :: BSChomper -> Char -> BSChomper
chompBS bsc c = case (isSpace c, bscNumWS bsc, bscBuilder bsc) of
  (True, _, Nothing) -> bsc
  (True, n, Just _) -> bsc { bscNumWS = n + 1 }
  (False, _, Nothing) -> bsc { bscBuilder = Just (encodeCharUTF8 c) }
  (False, 0, Just builder) -> bsc { bscBuilder = Just (builder <> encodeCharUTF8 c) }
  (False, _, Just builder) -> bsc { bscBuilder = Just (builder <> encodeCharUTF8 ' ' <> encodeCharUTF8 c)
                                  , bscNumWS = 0
                                  }

finalizeBSC :: BSChomper -> LB.ByteString
finalizeBSC bsc = case bscBuilder bsc of
  Nothing      -> mempty
  Just builder -> LB.toLazyByteString builder

instance {-# OVERLAPS #-} SpaceChompable B.ByteString where
  chompSpaces = UTF8.foldl chompBS (BSChomper 0 Nothing)
    >>> finalizeBSC
    >>> LB.toStrict
instance {-# OVERLAPS #-} SpaceChompable LB.ByteString where
  chompSpaces = LUTF8.foldl chompBS (BSChomper 0 Nothing)
    >>> finalizeBSC

-- |
-- Convert a strict ByteString into a Text `LT.Builder', converting any invalid
-- characters into the Unicode replacement character � (U+FFFD).
bsToTextBuilder :: B.ByteString -> LT.Builder
bsToTextBuilder = UTF8.foldr (\char bldr -> LT.singleton char <> bldr) mempty

-- |
-- Convert a lazy ByteString into a Text `LT.Builder', converting any invalid
-- characters into the Unicode replacement character � (U+FFFD).
lbsToTextBuilder :: LB.ByteString -> LT.Builder
lbsToTextBuilder = LUTF8.foldr (\char bldr -> LT.singleton char <> bldr) mempty

-- |
-- "Data.ByteString.Builder" provides `charUtf8' to do this, but it doesn't
-- correctly handle invalid characters.
encodeCharUTF8 :: Char -> LB.Builder
encodeCharUTF8 c =
  let normalized = case c of
        '\xFFFE' -> '\xFFFD'
        '\xFFFF' -> '\xFFFD'
        _        -> c
  in LB.charUtf8 normalized
