module Country.Unexposed.TrieByte
  ( TrieByte
  , trieByteFromList
  , trieByteParser
  ) where

import Data.HashMap.Strict (HashMap)
import Data.Word (Word16,Word8)
import Data.ByteString (ByteString)
import Data.Semigroup (Semigroup)
import Control.Applicative ((<|>))
import qualified Data.ByteString as B
import qualified Data.HashMap.Strict as HM
import qualified Data.Attoparsec.ByteString as AB
import qualified Data.Semigroup as SG

-- | If the value is not the max Word16 (65535), there 
--   is a match. This means that 65535 cannot be used, which 
--   is fine for this since 65535 is not used as a country code.
data TrieByte = TrieByte
  { TrieByte -> Word16
trieValue :: {-# UNPACK #-} !Word16
  , TrieByte -> HashMap Word8 TrieByte
trieChildren :: !(HashMap Word8 TrieByte)
  }

empty :: TrieByte
empty :: TrieByte
empty = Word16 -> HashMap Word8 TrieByte -> TrieByte
TrieByte Word16
placeholder forall k v. HashMap k v
HM.empty

append :: TrieByte -> TrieByte -> TrieByte
append :: TrieByte -> TrieByte -> TrieByte
append (TrieByte Word16
v1 HashMap Word8 TrieByte
c1) (TrieByte Word16
v2 HashMap Word8 TrieByte
c2) = Word16 -> HashMap Word8 TrieByte -> TrieByte
TrieByte (forall a. Ord a => a -> a -> a
min Word16
v1 Word16
v2) (forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
HM.unionWith TrieByte -> TrieByte -> TrieByte
append HashMap Word8 TrieByte
c1 HashMap Word8 TrieByte
c2)

placeholder :: Word16
placeholder :: Word16
placeholder = Word16
0xFFFF

singleton :: ByteString -> Word16 -> TrieByte
singleton :: ByteString -> Word16 -> TrieByte
singleton ByteString
fullName Word16
code = ByteString -> TrieByte
go ByteString
fullName where
  go :: ByteString -> TrieByte
  go :: ByteString -> TrieByte
go ByteString
name = case ByteString -> Maybe (Word8, ByteString)
B.uncons ByteString
name of
    Just (Word8
char,ByteString
nameNext) -> Word16 -> HashMap Word8 TrieByte -> TrieByte
TrieByte Word16
placeholder (forall k v. Hashable k => k -> v -> HashMap k v
HM.singleton Word8
char (ByteString -> TrieByte
go ByteString
nameNext))
    Maybe (Word8, ByteString)
Nothing -> Word16 -> HashMap Word8 TrieByte -> TrieByte
TrieByte Word16
code forall k v. HashMap k v
HM.empty

instance Semigroup TrieByte where
  <> :: TrieByte -> TrieByte -> TrieByte
(<>) = TrieByte -> TrieByte -> TrieByte
append

instance Monoid TrieByte where
  mempty :: TrieByte
mempty = TrieByte
empty
  mappend :: TrieByte -> TrieByte -> TrieByte
mappend = forall a. Semigroup a => a -> a -> a
(SG.<>)

trieByteFromList :: [(ByteString,Word16)] -> TrieByte
trieByteFromList :: [(ByteString, Word16)] -> TrieByte
trieByteFromList = forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ByteString -> Word16 -> TrieByte
singleton)

-- it seems like attoparsec should have some kind of convenience
-- for being able to commit to consuming a certain amount of
-- input once your certain that it will be consumed, but I cannot
-- find a way to use the api to do this.
trieByteParser :: TrieByte -> AB.Parser Word16
trieByteParser :: TrieByte -> Parser Word16
trieByteParser = TrieByte -> Parser Word16
go where
  go :: TrieByte -> AB.Parser Word16
  go :: TrieByte -> Parser Word16
go (TrieByte Word16
value HashMap Word8 TrieByte
children) = do
    let keepGoing :: Parser Word16
keepGoing = do
          Word8
c <- Parser Word8
AB.anyWord8
          case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup Word8
c HashMap Word8 TrieByte
children of
            Maybe TrieByte
Nothing -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"did not recognize country name"
            Just TrieByte
trieNext -> TrieByte -> Parser Word16
go TrieByte
trieNext
    if Word16
value forall a. Eq a => a -> a -> Bool
== Word16
placeholder
      then Parser Word16
keepGoing
      else Parser Word16
keepGoing forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> forall (m :: * -> *) a. Monad m => a -> m a
return Word16
value