{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE ViewPatterns          #-}
import           Data.Attoparsec.Text
import           Data.Coerce
import           Data.List              (sortOn)
import qualified Data.RadixTree         as R
import qualified Data.Text              as T
import           Test.SmallCheck
import           Test.SmallCheck.Series
import           Test.Tasty
import           Test.Tasty.SmallCheck

sortOnLengthDesc :: [T.Text] -> [T.Text]
sortOnLengthDesc = sortOn (negate . T.length)

mkAlternatives :: NonEmpty (NonEmpty Char) -> [T.Text]
mkAlternatives = sortOnLengthDesc . map T.pack . coerce

naiiveSameAsRadix :: Monad m => Bool -> Property m
naiiveSameAsRadix doCompress =
  forAll $ \(mkAlternatives -> alternatives) ->
    let
      rparse
        | doCompress = R.parse_ $! R.fromFoldable_ alternatives
        | otherwise = case R.compressBy (T.concat alternatives) (R.fromFoldable_ alternatives) of
            Just crt -> R.parse_ $! crt
            Nothing  -> error "could not compress radixtree!"
    in forAll $ \(textS :: NonEmpty Char) ->
      let text = T.pack (coerce textS)
      in parseOnly (choice (map string alternatives)) text ==
         parseOnly rparse text

lookupSameishAsParse :: Property IO
lookupSameishAsParse =
  forAll $ \alternatives ->
  let
    !rt = R.fromFoldable_ (mkAlternatives alternatives)
    !p = R.parse_ rt <* endOfInput
  in forAll $ \(textS :: NonEmpty Char) ->
    let text = T.pack (coerce textS)
    in case (parseOnly p text, R.lookup_ rt text) of
      (Right x1, Just x2) -> x1 == x2
      (Left  _ , Nothing) -> True
      _                   -> False

main :: IO ()
main = defaultMain $ testGroup "radixtree"
  [ testProperty "naiive parsing has the same result as radix tree parsing" (naiiveSameAsRadix False)
  , testProperty "naiive parsing has the same result as radix tree parsing when compressed"
    (naiiveSameAsRadix True)
  , testProperty "lookup ~ parse" lookupSameishAsParse
  ]