-- start snippet imports
{-# LANGUAGE DeriveGeneric, TupleSections, LambdaCase #-}
module JSONParser where

import Control.Applicative (Alternative(..), optional)
import Control.Monad (replicateM)
import Data.Bits (shiftL)
import Data.Char (isDigit, isHexDigit, isSpace, chr, ord, digitToInt)
import Data.Functor (($>))
import Data.List (intercalate)
import GHC.Generics (Generic)
import Numeric (showHex)
import Text.Printf (printf)
import Test.QuickCheck hiding (Positive, Negative)
import Data.Maybe (fromMaybe)
-- end snippet imports

-- start snippet type
data JValue = JNull
            | JBool Bool
            | JString String
            | JNumber { int :: Integer, frac :: [Int], exponent :: Integer}
            | JArray [JValue]
            | JObject [(String, JValue)]
            deriving (Eq, Generic)
-- end snippet type

-- start snippet show-instance
instance Show JValue where
  show value = case value of
    JNull          -> "null"
    JBool True     -> "true"
    JBool False    -> "false"
    JString s      -> showJSONString s
    JNumber s [] 0 -> show s
    JNumber s f 0  -> show s ++ "." ++ concatMap show f
    JNumber s [] e -> show s ++ "e" ++ show e
    JNumber s f e  -> show s ++ "." ++ concatMap show f ++ "e" ++ show e
    JArray a       -> "[" ++ intercalate ", " (map show a) ++ "]"
    JObject o      -> "{" ++ intercalate ", " (map showKV o) ++ "}"
    where
      showKV (k, v) = showJSONString k ++ ": " ++ show v

showJSONString :: String -> String
showJSONString s = "\"" ++ concatMap showJSONChar s ++ "\""

isControl :: Char -> Bool
isControl c = c `elem` ['\0' .. '\31']

showJSONChar :: Char -> String
showJSONChar c = case c of
  '\'' -> "'"
  '\"' -> "\\\""
  '\\' -> "\\\\"
  '/'  -> "\\/"
  '\b' -> "\\b"
  '\f' -> "\\f"
  '\n' -> "\\n"
  '\r' -> "\\r"
  '\t' -> "\\t"
  _ | isControl c -> "\\u" ++ showJSONNonASCIIChar c
  _ -> [c]
  where
    showJSONNonASCIIChar c =
      let a = "0000" ++ showHex (ord c) "" in drop (length a - 4) a
-- end snippet show-instance

-- start snippet scalargens
jNullGen :: Gen JValue
jNullGen = pure JNull

jBoolGen :: Gen JValue
jBoolGen = JBool <$> arbitrary

jNumberGen :: Gen JValue
jNumberGen = JNumber <$> arbitrary <*> listOf (choose (0, 9)) <*> arbitrary
-- end snippet scalargens

-- start snippet stringgen
jsonStringGen :: Gen String
jsonStringGen =
  concat <$> listOf (oneof [ vectorOf 1 arbitraryUnicodeChar
                           , escapedUnicodeChar ])
  where
    escapedUnicodeChar = ("\\u" ++) <$> vectorOf 4 (elements hexDigitLetters)
    hexDigitLetters    = ['0'..'9'] ++ ['a'..'f'] ++ ['A'..'F']

jStringGen :: Gen JValue
jStringGen = JString <$> jsonStringGen
-- end snippet stringgen

-- start snippet compositegens
jArrayGen :: Int -> Gen JValue
jArrayGen = fmap JArray . scale (`div` 2) . listOf . jValueGen . (`div` 2)

jObjectGen :: Int -> Gen JValue
jObjectGen = fmap JObject . scale (`div` 2) . listOf . objKV . (`div` 2)
  where
    objKV n = (,) <$> jsonStringGen <*> jValueGen n
-- end snippet compositegens

-- start snippet valuegen
jValueGen :: Int -> Gen JValue
jValueGen n = if n < 5
  then frequency [(4, oneof scalarGens), (1, oneof (compositeGens n))]
  else frequency [(1, oneof scalarGens), (4, oneof (compositeGens n))]
  where
    scalarGens      = [jNullGen , jBoolGen , jNumberGen , jStringGen]
    compositeGens n = [jArrayGen n, jObjectGen n]
-- end snippet valuegen

-- start snippet value-arbitrary
instance Arbitrary JValue where
  arbitrary = sized jValueGen
  shrink    = genericShrink
-- end snippet value-arbitrary

-- start snippet stringify
jsonWhitespaceGen :: Gen String
jsonWhitespaceGen =
  scale (round . sqrt . fromIntegral)
  . listOf
  . elements
  $ [' ' , '\n' , '\r' , '\t']

stringify :: JValue -> Gen String
stringify = pad . go
  where
    surround l r j = l ++ j ++ r
    pad gen = surround <$> jsonWhitespaceGen <*> jsonWhitespaceGen <*> gen
    commaSeparated = pad . pure . intercalate ","

    go value = case value of
      JArray elements ->
        mapM (pad . stringify) elements
          >>= fmap (surround "[" "]") . commaSeparated
      JObject kvs ->
        mapM stringifyKV kvs >>= fmap (surround "{" "}") . commaSeparated
      _           -> return $ show value

    stringifyKV (k, v) =
      surround <$> pad (pure $ showJSONString k) <*> stringify v <*> pure ":"
-- end snippet stringify

-- start snippet parse-result
data ParseResult a = Error [String] | Result a
-- end snippet parse-result

-- start snippet parse-result-instances
instance Show a => Show (ParseResult a) where
  show (Result res) = show res
  show (Error errs) = formatErrors (reverse errs)
    where
      formatErrors []         = error "No errors to format"
      formatErrors [err]      = err
      formatErrors (err:errs) = err <> "\n> " <> intercalate "\n> "
        (map (concatMap (\c -> if c == '\n' then "\n  " else [c])) errs)

instance Functor ParseResult where
  fmap _ (Error errs) = Error errs
  fmap f (Result res) = Result (f res)

instance Applicative ParseResult where
  pure = Result
  Error  errs <*> _   = Error errs
  Result f <*> result = fmap f result

instance Alternative ParseResult where
  empty = Error ["Unknown error."]
  Error _ <|> result = result
  Result res <|> _   = Result res
-- end snippet parse-result-instances

-- start snippet parser1
newtype Parser1 i o = Parser1 { runParser1 :: i -> ParseResult (i, o) }
-- end snippet parser1

-- start snippet parser1-instances
instance Functor (Parser1 i) where
  fmap f parser = Parser1 $ fmap (fmap f) . runParser1 parser

instance Applicative (Parser1 i) where
  pure x    = Parser1 $ pure . (, x)
  pf <*> pa = Parser1 $ \input -> case runParser1 pf input of
    Error err        -> Error err
    Result (rest, f) -> fmap f <$> runParser1 pa rest

instance Alternative (Parser1 i) where
  empty = Parser1 $ const empty
  parser1 <|> parser2 = Parser1 $ \input ->
    runParser1 parser1 input <|> runParser1 parser2 input

instance Monad (Parser1 i) where
  parser >>= f = Parser1 $ \input -> case runParser1 parser input of
    Error errs            -> Error errs
    Result (rest, output) -> runParser1 (f output) rest
-- end snippet parser1-instances

-- start snippet parser-error1
parseError1 :: String -> ParseResult a
parseError1 err = Error [err]

throw1 :: String -> Parser1 String o
throw1 = Parser1 . const . parseError1
-- end snippet parser-error1

-- start snippet char-string-parsers1
satisfy1 :: (Char -> Bool) -> (Char -> String) -> Parser1 String Char
satisfy1 predicate mkError = Parser1 $ \case
  (c:cs) | predicate c -> Result (cs, c)
  (c:_)                -> parseError1 (mkError c)
  _                    -> parseError1 "Empty input"

char1 :: Char -> Parser1 String Char
char1 c = satisfy1 (== c) $ printf "Expected '%v', got '%v'" c

string1 :: String -> Parser1 String String
string1 ""     = pure ""
string1 (c:cs) = (:) <$> char1 c <*> string1 cs
-- end snippet char-string-parsers1

-- start snippet bool-parser-prev
jBoolPrev :: Parser1 String JValue
jBoolPrev =   string1 "true"  $> JBool True
          <|> string1 "false" $> JBool False
-- end snippet bool-parser-prev

-- start snippet bool-parser1
lookahead1 :: Parser1 String Char
lookahead1 = Parser1 $ \case
  input@(c:_) -> Result (input, c)
  _           -> parseError1 "Empty input"

jBool1 :: Parser1 String JValue
jBool1 = do
  c <- lookahead1
  JBool <$> case c of
    't' -> string1 "true"  $> True
    'f' -> string1 "false" $> False
    _   -> throw1 $
      printf "Expected: 't' for true or 'f' for false; got '%v'" c
-- end snippet bool-parser1

-- start snippet zipper
data TextZipper a =
  TextZipper {
    toLeft  :: a
  , toRight :: a
  , above   :: [a]
  , below   :: [a]
  }

instance Show a => Show (TextZipper a) where
  show (TextZipper left right above below) =
    "TextZipper{left=" <> show left
      <> ", right=" <> show right
      <> ", above=" <> show above
      <> ", below=" <> show below

textZipper :: Monoid a => [a] -> TextZipper a
textZipper []           = TextZipper mempty mempty [] []
textZipper (first:rest) = TextZipper mempty first [] rest

currentPosition :: TextZipper String -> (Int, Int)
currentPosition zipper =
  (length $ above zipper, length $ toLeft zipper)

currentChar :: TextZipper String -> Maybe Char
currentChar zipper = case toRight zipper of
  []    -> Nothing
  (c:_) -> Just c

moveByOne :: TextZipper String -> TextZipper String
moveByOne zipper
  | not $ null (toRight zipper) =
      zipper { toLeft  = toLeft zipper <> take 1 (toRight zipper)
             , toRight = drop 1 (toRight zipper)
             }
  | not $ null (below zipper) =
      zipper { above   = above zipper ++ [toLeft zipper]
             , below   = tail $ below zipper
             , toLeft  = ""
             , toRight = head (below zipper)
             }
  | otherwise = zipper

move :: TextZipper String -> TextZipper String
move zipper =
  let zipper' = moveByOne zipper
  in if toRight zipper' == "" && not (isLastLine zipper')
     then move zipper'
     else zipper'
  where
    isLastLine = null . below
-- end snippet zipper

-- start snippet parser
newtype Parser i o = Parser {
    runParser_ :: TextZipper i -> ParseResult (TextZipper i, o)
  }

runParser :: Parser String o -> String -> ParseResult (String, o)
runParser parser input = case runParser_ parser inputZipper of
  Error errs                  -> Error errs
  Result (restZipper, output) ->
    Result ( intercalate "\n" (toRight restZipper : below restZipper)
            , output )
  where
    inputZipper =
      let zipper = textZipper $ splitOn '\n' input
      in case currentChar zipper of
           Nothing -> move zipper
           Just _  -> zipper

    splitOn delim string =
      let (splits, rest) = foldr (\c (acc, string) ->
              if c == delim then ((c:string) : acc, "") else (acc, c:string))
            ([], "") string
      in if rest == "" then splits else rest : splits

instance Functor (Parser i) where
  fmap f parser = Parser $ fmap (fmap f) . runParser_ parser

instance Applicative (Parser i) where
  pure x    = Parser $ pure . (, x)
  pf <*> pa = Parser $ \input -> case runParser_ pf input of
    Error err        -> Error err
    Result (rest, f) -> fmap f <$> runParser_ pa rest

instance Alternative (Parser i) where
  empty = Parser $ const empty
  parser1 <|> parser2 = Parser $ \input ->
    runParser_ parser1 input <|> runParser_ parser2 input

instance Monad (Parser i) where
  parser >>= f = Parser $ \input -> case runParser_ parser input of
    Error err        -> Error err
    Result (rest, o) -> runParser_ (f o) rest
-- end snippet parser

-- start snippet add-line-col
addLineCol :: String -> TextZipper String -> String
addLineCol err zipper =
  let (ln, cn) = currentPosition zipper
      err'     = printf (err <> " at line %d, column %d: ") (ln + 1) (cn + 1)
      left     = toLeft zipper
      right    = toRight zipper
      left'    = replaceWhitespace $ drop (length left - 5) left
      right'   = replaceWhitespace $ take 5 right
      line     = left' <> right'
  in printf (err' <> "%s\n%s^")
            line
            (replicate (length err' + length left') ' ')
  where
    replaceWhitespace = concatMap (\c ->
      fromMaybe (if isControl c then "\\" <> show (ord c) else [c])
                (lookup c replacements))
    replacements = [('\t', "\\t"), ('\r', "\\r"), ('\n', "\\n")]
-- end snippet add-line-col

-- start snippet parser-error
parseError :: String -> TextZipper String -> ParseResult a
parseError err zipper = Error [addLineCol err zipper]

throw :: String -> Parser String o
throw = Parser . parseError

orThrow :: Parser String o -> String -> Parser String o
orThrow parser err = Parser $ \input ->
  case runParser_ parser input of
    Result (rest, a) -> Result (rest, a)
    Error errs       -> Error (addLineCol err input : errs)
-- end snippet parser-error

-- start snippet basic-parsers
lookahead :: Parser String Char
lookahead = Parser $ \input -> case currentChar input of
  Just c  -> Result (input, c)
  Nothing -> parseError "Empty input" input

satisfy :: (Char -> Bool) -> (Char -> String) -> Parser String Char
satisfy predicate mkError = Parser $ \input -> case currentChar input of
  Just c | predicate c -> Result (move input, c)
  Just c               -> parseError (mkError c) input
  _                    -> parseError "Empty input" input
-- end snippet basic-parsers

-- start snippet char-digit-string-parsers
char :: Char -> Parser String Char
char c = satisfy (== c) $ printf "Expected '%v', got '%v'" c

digit :: Parser String Int
digit = digitToInt <$> satisfy isDigit (printf "Expected a digit, got '%v'")

string :: String -> Parser String String
string ""     = pure ""
string (c:cs) = (:) <$> char c <*> string cs
-- end snippet char-digit-string-parsers

-- start snippet null-parser
jNull :: Parser String JValue
jNull = string "null" $> JNull
-- end snippet null-parser

-- start snippet bool-parser
jBool :: Parser String JValue
jBool = do
  c <- lookahead
  JBool <$> case c of
    't' -> string "true"  $> True
    'f' -> string "false" $> False
    _   -> throw $ printf "Expected: 't' for true or 'f' for false; got '%v'" c
-- end snippet bool-parser

-- start snippet char-parser
jsonChar :: Parser String Char
jsonChar = lookahead >>= \case
  '\\' -> char '\\' *> escapedChar
  _    -> otherChar
  where
    escapedChar = lookahead >>= \case
      '"'  -> char '"'
      '\\' -> char '\\'
      '/'  -> char  '/'
      'b'  -> char 'b' $> '\b'
      'f'  -> char 'f' $> '\f'
      'n'  -> char 'n' $> '\n'
      'r'  -> char 'r' $> '\r'
      't'  -> char 't' $> '\t'
      'u'  -> char 'u' *> unicodeChar
      c    -> throw $ printf "Invalid escaped character: '%v'." c

    unicodeChar =
      chr . fromIntegral . digitsToNumber 16 0 <$> replicateM 4 hexDigit

    hexDigit = digitToInt <$>
      satisfy isHexDigit (printf "Expected a hex digit, got '%v'")

    otherChar = satisfy
      (\c -> not (c == '\"' || c == '\\' || isControl c))
      (\c -> let c' = if isControl c then "\\" <> show (ord c) else [c]
        in printf "Did not expected '\"', '\\' and control chars, got '%v'" c')

digitsToNumber :: Int -> Integer -> [Int] -> Integer
digitsToNumber base =
  foldl (\num digit -> num * fromIntegral base + fromIntegral digit)
-- end snippet char-parser

-- start snippet string-parser
jString :: Parser String JValue
jString = JString <$> (char '"' *> jString')
  where
    jString' = lookahead `orThrow` "Expected rest of a string" >>= \case
      '"' -> "" <$ char '"'
      _   -> do
        first <- jsonChar
        if not (isSurrogate first)
        then (:) <$> pure first <*> jString'
        else do
          second <- jsonChar `orThrow` "Expected second character of a surrogate pair"
          if isHighSurrogate first
          then if isLowSurrogate second
               then (combineSurrogates first second :) <$> jString'
               else throw $ printf "Expected a low surrogate, got '%v'" second
          else throw $ printf "Expected a high surrogate, got '%v'" first
-- end snippet string-parser

-- start snippet string-parser-helper
highSurrogateLowerBound, highSurrogateUpperBound :: Int
highSurrogateLowerBound = 0xD800
highSurrogateUpperBound = 0xDBFF

lowSurrogateLowerBound, lowSurrogateUpperBound :: Int
lowSurrogateLowerBound  = 0xDC00
lowSurrogateUpperBound  = 0xDFFF

isHighSurrogate, isLowSurrogate, isSurrogate :: Char -> Bool
isHighSurrogate a =
  ord a >= highSurrogateLowerBound && ord a <= highSurrogateUpperBound
isLowSurrogate a  =
  ord a >= lowSurrogateLowerBound && ord a <= lowSurrogateUpperBound
isSurrogate a     = isHighSurrogate a || isLowSurrogate a

combineSurrogates :: Char -> Char -> Char
combineSurrogates a b = chr $
  ((ord a - highSurrogateLowerBound) `shiftL` 10)
  + (ord b - lowSurrogateLowerBound) + 0x10000
-- end snippet string-parser-helper

-- start snippet string-parser-test
prop_genParseJString :: Property
prop_genParseJString =
  forAllShrink jStringGen shrink $ \js ->
    case runParser jString (show js) of
      Error _       -> False
      Result (_, o) -> o == js
-- end snippet string-parser-test

-- start snippet int-parser
jUInt :: Parser String Integer
jUInt = flip orThrow "Expected unsigned integer" $
  lookahead >>= \case
    '0'           -> fromIntegral <$> digit
    c | isDigit c -> digitsToNumber 10 0 <$> digits
    c             -> throw $ printf "Expected a digit, got '%v'" c

digits :: Parser String [Int]
digits = some digit `orThrow` "Expected digits"

jInt' :: Parser String Integer
jInt' = signInt <$> optional (char '-') <*> jUInt

signInt :: Maybe Char -> Integer -> Integer
signInt (Just '-') i = negate i
signInt _          i = i
-- end snippet int-parser

-- start snippet fracexp-parser
jFrac :: Parser String [Int]
jFrac = char '.' *> digits

jExp :: Parser String Integer
jExp = (char 'e' <|> char 'E')
  *> (signInt <$> optional (char '+' <|> char '-') <*> jUInt)
-- end snippet fracexp-parser

-- start snippet number-parser
jNumber :: Parser String JValue
jNumber = do
  i <- jInt'
  optional lookahead >>= \case
    Just c | c == 'e' || c == 'E' ->
      JNumber <$> pure i <*> pure [] <*> jExp         -- int+exp
    Just '.' -> do
      f <- jFrac
      optional lookahead >>= \case
        Just c' | c' == 'e' || c' == 'E' ->
          JNumber <$> pure i <*> pure f <*> jExp      -- int+frac+exp
        _ -> JNumber <$> pure i <*> pure f <*> pure 0 -- int+frac
    _ -> JNumber <$> pure i <*> pure [] <*> pure 0    -- int
-- end snippet number-parser

-- start snippet number-parser-test
prop_genParseJNumber :: Property
prop_genParseJNumber =
  forAllShrink jNumberGen shrink $ \jn ->
    case runParser jNumber (show jn) of
      Error _       -> False
      Result (_, o) -> o == jn
-- end snippet number-parser-test

-- start snippet array-parser-helper
surroundedBy :: Parser i a -> Parser i b -> Parser i a
surroundedBy parser1 parser2 = parser2 *> parser1 <* parser2

separatedBy :: Parser i v -> Parser i s -> Parser i [v]
separatedBy parser sepParser =
  (:) <$> parser <*> many (sepParser *> parser)

space :: Parser String Char
space = do
  c <- lookahead
  if c == ' ' || c == '\n' || c == '\r' || c == '\t'
  then char c
  else throw $ printf "Expected space, got '%v'" c

spaces :: Parser String String
spaces = many space
-- end snippet array-parser-helper

-- start snippet array-parser
jArray :: Parser String JValue
jArray = JArray <$> do
  _ <- char '[' <* spaces
  lookahead >>= \case
    ']' -> [] <$ (spaces <* char ']')
    _   -> (jValue `separatedBy` char ',') <* char ']'
-- end snippet array-parser

-- start snippet array-parser-test
prop_genParseJArray :: Property
prop_genParseJArray =
  forAllShrink (sized jArrayGen) shrink $ \ja -> do
    jas <- dropWhile isSpace <$> stringify ja
    return . counterexample (show jas) $ case runParser jArray jas of
      Error _       -> False
      Result (_, o) -> o == ja
-- end snippet array-parser-test

-- start snippet object-parser
jObject :: Parser String JValue
jObject = JObject <$> do
  _ <- char '{' <* spaces
  lookahead >>= \case
    '}' -> [] <$ (spaces <* char '}')
    _   -> (pair `separatedBy` char ',') <* char '}'
  where
    pair = (\ ~(JString s) j -> (s, j))
      <$> ((jString `surroundedBy` spaces) `orThrow` "Expected object key")
      <*  char ':'
      <*> (jValue `orThrow` "Expected object value")
-- end snippet object-parser

-- start snippet object-parser-test
prop_genParseJObject :: Property
prop_genParseJObject =
  forAllShrink (sized jObjectGen) shrink $ \jo -> do
    jos <- dropWhile isSpace <$> stringify jo
    return . counterexample (show jos) $ case runParser jObject jos of
      Error _       -> False
      Result (_, o) -> o == jo
-- end snippet object-parser-test

-- start snippet value-parser
jValue :: Parser String JValue
jValue = jValue' `surroundedBy` spaces
  where
    jValue' = lookahead >>= \case
      'n'  -> jNull   `orThrow` "Expected null"
      't'  -> jBool   `orThrow` "Expected true"
      'f'  -> jBool   `orThrow` "Expected false"
      '\"' -> jString `orThrow` "Expected string"
      '['  -> jArray  `orThrow` "Expected array"
      '{'  -> jObject `orThrow` "Expected object"
      c    -> if c == '-' || isDigit c
         then jNumber `orThrow` "Expected number"
         else throw $ printf "Unexpected char: '%v'." c
-- end snippet value-parser

-- start snippet parsejson
parseJSON :: String -> Either String JValue
parseJSON s = case runParser jValue s of
  Result ("", j) -> Right j
  Result (i, _)  -> Left $ "Leftover input: " <> i
  err@(Error _)  -> Left $ show err
-- end snippet parsejson

-- start snippet parsejson-test
prop_genParseJSON :: Property
prop_genParseJSON = forAllShrink (sized jValueGen) shrink $ \value -> do
  json <- stringify value
  return . counterexample (show json) . (== Right value) . parseJSON $ json
-- end snippet parsejson-test

-- start snippet run-test
runTests :: IO ()
runTests = do
  putStrLn "== prop_genParseJString =="
  quickCheck prop_genParseJString

  putStrLn "== prop_genParseJNumber =="
  quickCheck prop_genParseJNumber

  putStrLn "== prop_genParseJArray =="
  quickCheck prop_genParseJArray

  putStrLn "== prop_genParseJObject =="
  quickCheck prop_genParseJObject

  putStrLn "== prop_genParseJSON =="
  quickCheck prop_genParseJSON
-- end snippet run-test

-- start snippet print-result
printResult :: Either String JValue -> IO ()
printResult = putStrLn . either ("ERROR:\n" <>) (("RESULT:\n" <>) . show)
-- end snippet print-result