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

import Control.Applicative ((<|>))
import qualified Data.Attoparsec.Text as AT
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import qualified Data.Semigroup as SG
import Data.Text (Text)
import qualified Data.Text as T
import Data.Word (Word16)

{- | 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 HashMap Char Trie
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 (Word16 -> Word16 -> Word16
forall a. Ord a => a -> a -> a
min Word16
v1 Word16
v2) ((Trie -> Trie -> Trie)
-> HashMap Char Trie -> HashMap Char Trie -> HashMap Char Trie
forall k v.
Eq 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 (Char -> Trie -> HashMap Char Trie
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 HashMap Char Trie
forall k v. HashMap k v
HM.empty

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

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

trieFromList :: [(Text, Word16)] -> Trie
trieFromList :: [(Text, Word16)] -> Trie
trieFromList = ((Text, Word16) -> Trie) -> [(Text, Word16)] -> Trie
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap ((Text -> Word16 -> Trie) -> (Text, Word16) -> Trie
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 Char -> HashMap Char Trie -> Maybe Trie
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 -> String -> Parser Word16
forall a. String -> Parser Text a
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 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 Text a -> Parser Text a -> Parser Text a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Word16 -> Parser Word16
forall a. a -> Parser Text a
forall (m :: * -> *) a. Monad m => a -> m a
return Word16
value