{-# LANGUAGE BangPatterns               #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE InstanceSigs               #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE OverloadedStrings          #-}
{-# LANGUAGE ScopedTypeVariables        #-}

module HaskellWorks.Data.Json.LightJson where

import           Control.Arrow
import qualified Data.Attoparsec.ByteString.Char8           as ABC
import qualified Data.ByteString.Char8                      as BSC
import qualified Data.ByteString                            as BS
import qualified Data.DList                                 as DL
import qualified Data.List                                  as L
import           Data.String
import           Data.Word
import           Data.Word8
import           HaskellWorks.Data.AtLeastSize
import           HaskellWorks.Data.Bits.BitWise
import qualified HaskellWorks.Data.BalancedParens           as BP
import           HaskellWorks.Data.Drop
import           HaskellWorks.Data.Entry
import           HaskellWorks.Data.Json.CharLike
import           HaskellWorks.Data.Json.Conduit.Words
import           HaskellWorks.Data.Json.Succinct
import           HaskellWorks.Data.Micro
import           HaskellWorks.Data.Mini
import           HaskellWorks.Data.MQuery
import           HaskellWorks.Data.Positioning
import           HaskellWorks.Data.RankSelect.Base.Rank0
import           HaskellWorks.Data.RankSelect.Base.Rank1
import           HaskellWorks.Data.RankSelect.Base.Select1
import           HaskellWorks.Data.Row
import           HaskellWorks.Data.TreeCursor
import           HaskellWorks.Data.Uncons
import           Prelude hiding (drop)
import           Text.PrettyPrint.ANSI.Leijen

data LightJson c
  = LightJsonString String
  | LightJsonNumber BS.ByteString
  | LightJsonObject [(String, c)]
  | LightJsonArray [c]
  | LightJsonBool Bool
  | LightJsonNull
  | LightJsonError String
  deriving Show

instance Eq (LightJson c) where
  (==) (LightJsonString a) (LightJsonString b)  = a == b
  (==) (LightJsonNumber a) (LightJsonNumber b)  = a == b
  (==) (LightJsonBool   a) (LightJsonBool   b)  = a == b
  (==)  LightJsonNull       LightJsonNull       = True
  (==)  _                   _                   = False

-- instance Ord (LightJson c) where
--   compare (LightJsonString a) (LightJsonString b)  = a `compare` b
--   compare (LightJsonNumber a) (LightJsonNumber b)  = a `compare` b
--   compare (LightJsonBool   a) (LightJsonBool   b)  = a `compare` b
--   compare  LightJsonNull       LightJsonNull       = True

data (LightJsonField c) = LightJsonField String (LightJson c)

class LightJsonAt a where
  lightJsonAt :: a -> LightJson a

wSpace :: Word8
wSpace = 0x20

data JsonState
  = Escaped
  | InJson
  | InString
  | InNumber
  | InIdent

slurpByteString :: BS.ByteString -> BS.ByteString
slurpByteString bs = let (!cs, _) = BS.unfoldrN (BS.length bs) genString (InJson, bs) in cs
  where genString :: (JsonState, BS.ByteString) -> Maybe (Word8, (JsonState, BS.ByteString))
        genString (InJson, cs) = case BS.uncons cs of
          Just (!e, !es) | e == _quotedbl     -> genString            (InString , es)
          -- TODO: Only match whitespace
          Just (!_, !es)                      -> genString            (InJson   , es)
          Nothing                             -> Nothing
        genString (InString, ds) = case BS.uncons ds of
          Just (!e, !es) | e == _backslash    -> genString            (Escaped  , es)
          Just (!e, !_ ) | e == _quotedbl     -> Nothing
          Just (e , !es)                      -> Just (e            , (InString , es))
          Nothing                             -> Nothing
        genString (Escaped, ds) = case BS.uncons ds of
          Just (_ , !es)                      -> Just (_period      , (InString , es))
          Nothing                             -> Nothing
        genString (_, _) = Nothing

slurpString :: BS.ByteString -> String
slurpString bs = L.unfoldr genString (InJson, BSC.unpack bs)
  where genString :: (JsonState, String) -> Maybe (Char, (JsonState, String))
        genString (InJson, ds) = case ds of
          (e:es)  | e == '"'  -> genString  (InString , es)
          (_:es)              -> genString  (InJson   , es)
          _                   -> Nothing
        genString (InString, ds) = case ds of
          (e:es) | e == '\\'  -> genString  (Escaped  , es)
          (e:_ ) | e == '"'   -> Nothing
          (e:es)              -> Just (e,   (InString , es))
          _                   -> Nothing
        genString (Escaped, ds) = case ds of
          (_:es)              -> Just ('.', (InString , es))
          _                   -> Nothing
        genString (_, _) = Nothing

slurpNumber :: BS.ByteString -> BS.ByteString
slurpNumber bs = let (!cs, _) = BS.unfoldrN (BS.length bs) genNumber (InJson, bs) in cs
    where genNumber :: (JsonState, BS.ByteString) -> Maybe (Word8, (JsonState, BS.ByteString))
          genNumber (InJson, cs) = case BS.uncons cs of
            Just (!d, !ds) | isLeadingDigit d   -> Just (d           , (InNumber , ds))
            Just (!d, !ds)                      -> Just (d           , (InJson   , ds))
            Nothing -> Nothing
          genNumber (InNumber, cs) = case BS.uncons cs of
            Just (!d, !ds) | isTrailingDigit d  -> Just (d           , (InNumber , ds))
            Just (!d, !ds) | d == _quotedbl     -> Just (_parenleft  , (InString , ds))
            _                                   -> Nothing
          genNumber (_, _) = Nothing

instance (BP.BalancedParens w, Rank0 w, Rank1 w, Select1 v, TestBit w) => LightJsonAt (JsonCursor BS.ByteString v w) where
  lightJsonAt k = case uncons remainder of
    Just (!c, _) | isLeadingDigit2 c  -> LightJsonNumber  (slurpNumber remainder)
    Just (!c, _) | isQuotDbl c        -> LightJsonString  (slurpString remainder)
    Just (!c, _) | isChar_t c         -> LightJsonBool    True
    Just (!c, _) | isChar_f c         -> LightJsonBool    False
    Just (!c, _) | isChar_n c         -> LightJsonNull
    Just (!c, _) | isBraceLeft c      -> LightJsonObject (mapValuesFrom   (firstChild k))
    Just (!c, _) | isBracketLeft c    -> LightJsonArray  (arrayValuesFrom (firstChild k))
    Just _                            -> LightJsonError "Invalid Json Type"
    Nothing                           -> LightJsonError "End of data"
    where ik                = interests k
          bpk               = balancedParens k
          p                 = lastPositionOf (select1 ik (rank1 bpk (cursorRank k)))
          remainder         = drop (toCount p) (cursorText k)
          arrayValuesFrom   = L.unfoldr (fmap (id &&& nextSibling))
          mapValuesFrom j   = pairwise (arrayValuesFrom j) >>= asField
          pairwise (a:b:rs) = (a, b) : pairwise rs
          pairwise _        = []
          asField (a, b)    = case lightJsonAt a of
                                LightJsonString s -> [(s, b)]
                                _                 -> []

toLightJsonField :: (String, LightJson c) -> LightJsonField c
toLightJsonField (k, v) = LightJsonField k v

instance LightJsonAt c => Pretty (LightJsonField c) where
  pretty (LightJsonField k v) = text (show k) <> text ": " <> pretty v

hEncloseSep :: Doc -> Doc -> Doc -> [Doc] -> Doc
hEncloseSep l r s ds
    = case ds of
        []  -> l <> r
        [d] -> l <> d <> r
        _   -> hcat (zipWith (<>) (l : repeat s) ds) <> r

instance LightJsonAt c => Pretty (LightJson c) where
  pretty c = case c of
    LightJsonString s   -> dullgreen  (text (show s))
    LightJsonNumber n   -> cyan       (text (show n))
    LightJsonObject []  -> text "{}"
    LightJsonObject kvs -> hEncloseSep (text "{") (text "}") (text ",") ((pretty . toLightJsonField . second lightJsonAt) `map` kvs)
    LightJsonArray vs   -> hEncloseSep (text "[") (text "]") (text ",") ((pretty . lightJsonAt) `map` vs)
    LightJsonBool w     -> red (text (show w))
    LightJsonNull       -> text "null"
    LightJsonError s    -> text "<error " <> text s <> text ">"

instance Pretty (Micro (LightJson c)) where
  pretty (Micro (LightJsonString s )) = dullgreen (text (show s))
  pretty (Micro (LightJsonNumber n )) = cyan      (text (show n))
  pretty (Micro (LightJsonObject [])) = text "{}"
  pretty (Micro (LightJsonObject _ )) = text "{..}"
  pretty (Micro (LightJsonArray [] )) = text "[]"
  pretty (Micro (LightJsonArray _  )) = text "[..]"
  pretty (Micro (LightJsonBool w   )) = red (text (show w))
  pretty (Micro  LightJsonNull      ) = text "null"
  pretty (Micro (LightJsonError s  )) = text "<error " <> text s <> text ">"

instance Pretty (Micro (String, LightJson c)) where
  pretty (Micro (fieldName, jpv)) = red (text (show fieldName)) <> text ": " <> pretty (Micro jpv)

instance LightJsonAt c => Pretty (Mini (LightJson c)) where
  pretty mjpv = case mjpv of
    Mini (LightJsonString s   ) -> dullgreen  (text (show s))
    Mini (LightJsonNumber n   ) -> cyan       (text (show n))
    Mini (LightJsonObject []  ) -> text "{}"
    Mini (LightJsonObject kvs ) -> case kvs of
      (_:_:_:_:_:_:_:_:_:_:_:_:_) -> text "{" <> prettyKvs (map (second lightJsonAt) kvs) <> text ", ..}"
      []                          -> text "{}"
      _                           -> text "{" <> prettyKvs (map (second lightJsonAt) kvs) <> text "}"
    Mini (LightJsonArray []   ) -> text "[]"
    Mini (LightJsonArray vs   ) | vs `atLeastSize` 11 -> text "[" <> nest 2 (prettyVs ((Micro . lightJsonAt) `map` take 10 vs)) <> text ", ..]"
    Mini (LightJsonArray vs   ) | vs `atLeastSize` 1  -> text "[" <> nest 2 (prettyVs ((Micro . lightJsonAt) `map` take 10 vs)) <> text "]"
    Mini (LightJsonArray _    )                       -> text "[]"
    Mini (LightJsonBool w     ) -> red (text (show w))
    Mini  LightJsonNull         -> text "null"
    Mini (LightJsonError s    ) -> text "<error " <> text s <> text ">"

instance LightJsonAt c => Pretty (Mini (String, LightJson c)) where
  pretty (Mini (fieldName, jpv)) = text (show fieldName) <> text ": " <> pretty (Mini jpv)

instance LightJsonAt c => Pretty (MQuery (LightJson c)) where
  pretty = pretty . Row 120 . mQuery

instance LightJsonAt c => Pretty (MQuery (Entry String (LightJson c))) where
  pretty (MQuery das) = pretty (Row 120 das)

-- hasKV :: LightJsonAt c => BS.ByteString -> LightJson c -> LightJson c -> MQuery (LightJson c)
-- hasKV k v (LightJsonObject xs)  = let ys = second lightJsonAt `map` xs in
--                                   if (k, v) `elem` ys then MQuery (DL.singleton (LightJsonObject xs)) else MQuery DL.empty
-- hasKV _ _  _                    = MQuery DL.empty

item :: LightJsonAt c => LightJson c -> MQuery (LightJson c)
item jpv = case jpv of
  LightJsonArray es -> MQuery $ DL.fromList (lightJsonAt `map` es)
  _                 -> MQuery   DL.empty

entry :: LightJsonAt c => LightJson c -> MQuery (Entry String (LightJson c))
entry jpv = case jpv of
  LightJsonObject fs  -> MQuery $ DL.fromList ((uncurry Entry . second lightJsonAt) `map` fs)
  _                   -> MQuery   DL.empty

asString :: LightJson c -> MQuery String
asString jpv = case jpv of
  LightJsonString s -> MQuery $ DL.singleton s
  _                 -> MQuery   DL.empty

asDouble :: LightJson c -> MQuery Double
asDouble jpv = case jpv of
  LightJsonNumber sn  -> case ABC.parse ABC.rational sn of
    ABC.Fail    {}    -> MQuery DL.empty
    ABC.Partial f     -> case f " " of
      ABC.Fail    {}    -> MQuery DL.empty
      ABC.Partial _     -> MQuery DL.empty
      ABC.Done    _ r   -> MQuery (DL.singleton r)
    ABC.Done    _ r   -> MQuery (DL.singleton r)
  _                   -> MQuery   DL.empty

asInteger :: LightJson c -> MQuery Integer
asInteger jpv = do
  d <- asDouble jpv
  return (floor d)

castAsInteger :: LightJson c -> MQuery Integer
castAsInteger jpv = case jpv of
  LightJsonString n -> MQuery $ DL.singleton (read n)
  LightJsonNumber _ -> asInteger jpv
  _                 -> MQuery   DL.empty

named :: String -> Entry String (LightJson c) -> MQuery (LightJson c)
named fieldName (Entry fieldName' jpv) | fieldName == fieldName'  = MQuery $ DL.singleton jpv
named _         _                                                 = MQuery   DL.empty

jsonKeys :: LightJson c -> [String]
jsonKeys jpv = case jpv of
  LightJsonObject fs  -> fst `map` fs
  _                   -> []

hasKey :: String -> LightJson c -> Bool
hasKey fieldName jpv = fieldName `elem` jsonKeys jpv

jsonSize :: LightJson c -> MQuery Integer
jsonSize jpv = case jpv of
  LightJsonArray  es  -> MQuery (DL.singleton (fromIntegral (length es)))
  LightJsonObject es  -> MQuery (DL.singleton (fromIntegral (length es)))
  _                   -> MQuery (DL.singleton 0)