module Country.Unexposed.Trie
  ( Trie
  , trieFromList
  , trieParser
  ) where

import Data.HashMap.Strict (HashMap)
import Data.Word (Word16)
import Data.Text (Text)
import Data.Semigroup (Semigroup)
import Control.Applicative ((<|>))
import qualified Data.Text as T
import qualified Data.HashMap.Strict as HM
import qualified Data.Attoparsec.Text as AT
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 Trie = Trie
  { Trie -> Word16
trieValue :: {-# UNPACK #-} !Word16
  , Trie -> HashMap Char Trie
trieChildren :: !(HashMap Char Trie)
  }

empty :: Trie
empty :: Trie
empty = Word16 -> HashMap Char Trie -> Trie
Trie Word16
placeholder forall k v. HashMap k v
HM.empty

append :: Trie -> Trie -> Trie
append :: Trie -> Trie -> Trie
append (Trie Word16
v1 HashMap Char Trie
c1) (Trie Word16
v2 HashMap Char Trie
c2) = Word16 -> HashMap Char Trie -> Trie
Trie (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 Trie -> Trie -> Trie
append HashMap Char Trie
c1 HashMap Char Trie
c2)

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

singleton :: Text -> Word16 -> Trie
singleton :: Text -> Word16 -> Trie
singleton Text
fullName Word16
code = Text -> Trie
go Text
fullName where
  go :: Text -> Trie
  go :: Text -> Trie
go Text
name = case Text -> Maybe (Char, Text)
T.uncons Text
name of
    Just (Char
char,Text
nameNext) -> Word16 -> HashMap Char Trie -> Trie
Trie Word16
placeholder (forall k v. Hashable k => k -> v -> HashMap k v
HM.singleton Char
char (Text -> Trie
go Text
nameNext))
    Maybe (Char, Text)
Nothing -> Word16 -> HashMap Char Trie -> Trie
Trie Word16
code forall k v. HashMap k v
HM.empty

instance Semigroup Trie where
  <> :: Trie -> Trie -> Trie
(<>) = forall a. Monoid a => a -> a -> a
mappend

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

trieFromList :: [(Text,Word16)] -> Trie
trieFromList :: [(Text, Word16)] -> Trie
trieFromList = 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 Text -> Word16 -> Trie
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.
trieParser :: Trie -> AT.Parser Word16
trieParser :: Trie -> Parser Word16
trieParser = Trie -> Parser Word16
go where
  go :: Trie -> AT.Parser Word16
  go :: Trie -> Parser Word16
go (Trie Word16
value HashMap Char Trie
children) = do
    let keepGoing :: Parser Word16
keepGoing = do
          Char
c <- Parser Char
AT.anyChar
          case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup Char
c HashMap Char Trie
children of
            Maybe Trie
Nothing -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"did not recognize country name"
            Just Trie
trieNext -> Trie -> Parser Word16
go Trie
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