{-# LANGUAGE FlexibleContexts #-}

module Data.Text.Punycode.Encode (encode) where

import Control.Monad
import Control.Monad.State hiding (state)
import Control.Monad.Writer
import qualified Data.ByteString as BS
import Data.Char
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import Data.Text.Punycode.Shared
import Data.Word

data PunycodeState = PunycodeState
  { PunycodeState -> Int
n :: Int,
    PunycodeState -> Int
delta :: Int,
    PunycodeState -> Int
bias :: Int,
    PunycodeState -> Int
h :: Int
  }

-- | Encode a string into its ascii form
encode :: T.Text -> BS.ByteString
encode :: Text -> ByteString
encode = Writer ByteString () -> ByteString
forall w a. Writer w a -> w
execWriter (Writer ByteString () -> ByteString)
-> (Text -> Writer ByteString ()) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Writer ByteString ()
forall (m :: * -> *). MonadWriter ByteString m => Text -> m ()
initialWriter

initialWriter :: (MonadWriter BS.ByteString m) => T.Text -> m ()
initialWriter :: forall (m :: * -> *). MonadWriter ByteString m => Text -> m ()
initialWriter Text
input = do
  ByteString -> m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell ByteString
basics
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ ByteString -> m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$ Word8 -> ByteString
BS.singleton (Word8 -> ByteString) -> Word8 -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ Char -> Int
ord Char
'-'
  StateT PunycodeState m () -> PunycodeState -> m ()
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT ([Int] -> Int -> StateT PunycodeState m ()
forall (m :: * -> *).
(MonadState PunycodeState m, MonadWriter ByteString m) =>
[Int] -> Int -> m ()
inner3 ((Char -> Int) -> [Char] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Char -> Int
ord ([Char] -> [Int]) -> [Char] -> [Int]
forall a b. (a -> b) -> a -> b
$ Text -> [Char]
T.unpack Text
input) Int
b) (PunycodeState -> m ()) -> PunycodeState -> m ()
forall a b. (a -> b) -> a -> b
$
    PunycodeState
      { n :: Int
n = Int
initial_n,
        delta :: Int
delta = Int
0,
        bias :: Int
bias = Int
initial_bias,
        h :: Int
h = Int
b
      }
  where
    basics :: ByteString
basics = Text -> ByteString
TE.encodeUtf8 (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> Text -> Text
T.filter Char -> Bool
isBasic Text
input
    b :: Int
b = ByteString -> Int
BS.length ByteString
basics

inner3 :: (MonadState PunycodeState m, MonadWriter BS.ByteString m) => [Int] -> Int -> m ()
inner3 :: forall (m :: * -> *).
(MonadState PunycodeState m, MonadWriter ByteString m) =>
[Int] -> Int -> m ()
inner3 [Int]
input Int
b = do
  PunycodeState
state <- m PunycodeState
forall s (m :: * -> *). MonadState s m => m s
get
  PunycodeState -> m ()
forall {m :: * -> *}.
(MonadState PunycodeState m, MonadWriter ByteString m) =>
PunycodeState -> m ()
helper PunycodeState
state
  where
    helper :: PunycodeState -> m ()
helper PunycodeState
state
      | Int
h' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
input = do
          PunycodeState -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (PunycodeState -> m ()) -> PunycodeState -> m ()
forall a b. (a -> b) -> a -> b
$ PunycodeState
state {n = m, delta = delta'}
          (Int -> m ()) -> [Int] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (Int -> Int -> m ()
forall (m :: * -> *).
(MonadState PunycodeState m, MonadWriter ByteString m) =>
Int -> Int -> m ()
inner2 Int
b) [Int]
input
          PunycodeState
state' <- m PunycodeState
forall s (m :: * -> *). MonadState s m => m s
get
          PunycodeState -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (PunycodeState -> m ()) -> PunycodeState -> m ()
forall a b. (a -> b) -> a -> b
$ PunycodeState
state' {delta = (delta state') + 1, n = (n state') + 1}
          [Int] -> Int -> m ()
forall (m :: * -> *).
(MonadState PunycodeState m, MonadWriter ByteString m) =>
[Int] -> Int -> m ()
inner3 [Int]
input Int
b
      | Bool
otherwise = () -> m ()
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
      where
        m :: Int
m = [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n') [Int]
input
        n' :: Int
n' = PunycodeState -> Int
n PunycodeState
state
        h' :: Int
h' = PunycodeState -> Int
h PunycodeState
state
        delta' :: Int
delta' = (PunycodeState -> Int
delta PunycodeState
state) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
m Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n') Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
h' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

inner2 :: (MonadState PunycodeState m, MonadWriter BS.ByteString m) => Int -> Int -> m ()
inner2 :: forall (m :: * -> *).
(MonadState PunycodeState m, MonadWriter ByteString m) =>
Int -> Int -> m ()
inner2 Int
b Int
c = do
  PunycodeState
state <- m PunycodeState
forall s (m :: * -> *). MonadState s m => m s
get
  PunycodeState -> m ()
forall {m :: * -> *}.
(MonadWriter ByteString m, MonadState PunycodeState m) =>
PunycodeState -> m ()
helper PunycodeState
state
  where
    helper :: PunycodeState -> m ()
helper PunycodeState
state
      | Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n' = do
          Int
q <- Int -> Int -> Int -> m Int
forall (m :: * -> *).
MonadWriter ByteString m =>
Int -> Int -> Int -> m Int
inner Int
delta' Int
base Int
bias'
          ByteString -> m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$ Word8 -> ByteString
BS.singleton (Word8 -> ByteString) -> Word8 -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> Word8
baseToAscii Int
q
          PunycodeState -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (PunycodeState -> m ()) -> PunycodeState -> m ()
forall a b. (a -> b) -> a -> b
$ PunycodeState
state {bias = adapt delta' (h' + 1) (h' == b), delta = 0, h = (h state) + 1}
      | Bool
otherwise = PunycodeState -> m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (PunycodeState -> m ()) -> PunycodeState -> m ()
forall a b. (a -> b) -> a -> b
$ PunycodeState
state {delta = delta'}
      where
        delta' :: Int
delta' = (PunycodeState -> Int
delta PunycodeState
state) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
d
          where
            d :: Int
d
              | Int
c Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n' = Int
1
              | Bool
otherwise = Int
0
        n' :: Int
n' = PunycodeState -> Int
n PunycodeState
state
        bias' :: Int
bias' = PunycodeState -> Int
bias PunycodeState
state
        h' :: Int
h' = PunycodeState -> Int
h PunycodeState
state

inner :: (MonadWriter BS.ByteString m) => Int -> Int -> Int -> m Int
inner :: forall (m :: * -> *).
MonadWriter ByteString m =>
Int -> Int -> Int -> m Int
inner Int
q Int
k Int
bias'
  | Int
q Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
t = Int -> m Int
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
q
  | Bool
otherwise = do
      ByteString -> m ()
forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$ Word8 -> ByteString
BS.singleton (Word8 -> ByteString) -> Word8 -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> Word8
baseToAscii (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ Int
t Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ((Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
t) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` (Int
base Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
t))
      Int -> Int -> Int -> m Int
forall (m :: * -> *).
MonadWriter ByteString m =>
Int -> Int -> Int -> m Int
inner ((Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
t) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` (Int
base Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
t)) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
base) Int
bias'
  where
    t :: Int
t
      | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
bias' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
tmin = Int
tmin
      | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
bias' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
tmax = Int
tmax
      | Bool
otherwise = Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
bias'

baseToAscii :: Int -> Word8
baseToAscii :: Int -> Word8
baseToAscii Int
i
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
26 = Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Char -> Int
ord Char
'a')
  | Bool
otherwise = Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
26) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Char -> Int
ord Char
'0')