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

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

{- | 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 HashMap Word8 TrieByte
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 (Word16 -> Word16 -> Word16
forall a. Ord a => a -> a -> a
min Word16
v1 Word16
v2) ((TrieByte -> TrieByte -> TrieByte)
-> HashMap Word8 TrieByte
-> HashMap Word8 TrieByte
-> HashMap Word8 TrieByte
forall k v.
Eq 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 (Word8 -> TrieByte -> HashMap Word8 TrieByte
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 HashMap Word8 TrieByte
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 = TrieByte -> TrieByte -> TrieByte
forall a. Semigroup a => a -> a -> a
(SG.<>)

trieByteFromList :: [(ByteString, Word16)] -> TrieByte
trieByteFromList :: [(ByteString, Word16)] -> TrieByte
trieByteFromList = ((ByteString, Word16) -> TrieByte)
-> [(ByteString, Word16)] -> TrieByte
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((ByteString -> Word16 -> TrieByte)
-> (ByteString, Word16) -> TrieByte
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 Word8 -> HashMap Word8 TrieByte -> Maybe TrieByte
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 -> String -> Parser Word16
forall a. String -> Parser ByteString a
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 Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
placeholder
      then Parser Word16
keepGoing
      else Parser Word16
keepGoing Parser Word16 -> Parser Word16 -> Parser Word16
forall a.
Parser ByteString a -> Parser ByteString a -> Parser ByteString a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Word16 -> Parser Word16
forall a. a -> Parser ByteString a
forall (m :: * -> *) a. Monad m => a -> m a
return Word16
value