-- SPDX-FileCopyrightText: Copyright (c) 2025 Objectionary.com
-- SPDX-License-Identifier: MIT

module Random (randomString, shuffle) where

import Control.Exception (throwIO)
import Control.Monad (forM_, replicateM)
import Data.Char (intToDigit)
import Data.IORef (IORef, atomicModifyIORef', newIORef)
import Data.Set (Set)
import qualified Data.Set as Set
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as M
import GHC.IO (unsafePerformIO)
import System.Random (newStdGen, randomRIO)
import System.Random.Stateful (newIOGenM, uniformRM)
import Text.Printf (printf)

strings :: IORef (Set String)
{-# NOINLINE strings #-}
strings :: IORef (Set [Char])
strings = IO (IORef (Set [Char])) -> IORef (Set [Char])
forall a. IO a -> a
unsafePerformIO (Set [Char] -> IO (IORef (Set [Char]))
forall a. a -> IO (IORef a)
newIORef Set [Char]
forall a. Set a
Set.empty)

generate :: String -> IO String
generate :: [Char] -> IO [Char]
generate [] = [Char] -> IO [Char]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
generate (Char
'%' : Char
ch : [Char]
rest) = do
  [Char]
rep <- case Char
ch of
    Char
'x' -> Int -> IO Char -> IO [Char]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
8 (IO Char -> IO [Char]) -> IO Char -> IO [Char]
forall a b. (a -> b) -> a -> b
$ do
      Int
v <- (Int, Int) -> IO Int
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Int
0, Int
15)
      Char -> IO Char
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Char
intToDigit Int
v)
    Char
'd' -> [Char] -> Int -> [Char]
forall r. PrintfType r => [Char] -> r
printf [Char]
"%04d" (Int -> [Char]) -> IO Int -> IO [Char]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> IO Int
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
randomRIO (Int
0 :: Int, Int
9999)
    Char
_ -> [Char] -> IO [Char]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char
'%', Char
ch]
  [Char]
next <- [Char] -> IO [Char]
generate [Char]
rest
  [Char] -> IO [Char]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char]
rep [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
next)
generate (Char
ch : [Char]
rest) = do
  [Char]
rest' <- [Char] -> IO [Char]
generate [Char]
rest
  [Char] -> IO [Char]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
ch Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char]
rest')

maxAttempts :: Int
maxAttempts :: Int
maxAttempts = Int
100000

regenerate :: String -> IO String
regenerate :: [Char] -> IO [Char]
regenerate [Char]
pat = Int -> IO [Char]
go Int
maxAttempts
  where
    go :: Int -> IO String
    go :: Int -> IO [Char]
go Int
0 = IOError -> IO [Char]
forall e a. Exception e => e -> IO a
throwIO ([Char] -> IOError
userError ([Char] -> [Char] -> [Char]
forall r. PrintfType r => [Char] -> r
printf [Char]
"randomString() cannot produce a unique value for pattern '%s': the value space is exhausted" [Char]
pat))
    go Int
attempts = do
      [Char]
next <- [Char] -> IO [Char]
generate [Char]
pat
      Bool
fresh <- IORef (Set [Char]) -> (Set [Char] -> (Set [Char], Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' IORef (Set [Char])
strings ((Set [Char] -> (Set [Char], Bool)) -> IO Bool)
-> (Set [Char] -> (Set [Char], Bool)) -> IO Bool
forall a b. (a -> b) -> a -> b
$ \Set [Char]
set ->
        if [Char]
next [Char] -> Set [Char] -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set [Char]
set
          then (Set [Char]
set, Bool
False)
          else ([Char] -> Set [Char] -> Set [Char]
forall a. Ord a => a -> Set a -> Set a
Set.insert [Char]
next Set [Char]
set, Bool
True)
      if Bool
fresh then [Char] -> IO [Char]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Char]
next else Int -> IO [Char]
go (Int
attempts Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

randomString :: String -> IO String
randomString :: [Char] -> IO [Char]
randomString [Char]
pat
  | [Char] -> Bool
randomized [Char]
pat = [Char] -> IO [Char]
regenerate [Char]
pat
  | Bool
otherwise = [Char] -> IO [Char]
generate [Char]
pat
  where
    randomized :: String -> Bool
    randomized :: [Char] -> Bool
randomized [] = Bool
False
    randomized (Char
'%' : Char
ch : [Char]
rest) = Char
ch Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'd' Bool -> Bool -> Bool
|| Char
ch Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'x' Bool -> Bool -> Bool
|| [Char] -> Bool
randomized [Char]
rest
    randomized (Char
_ : [Char]
rest) = [Char] -> Bool
randomized [Char]
rest

shuffle :: [a] -> IO [a]
shuffle :: forall a. [a] -> IO [a]
shuffle [a]
xs = do
  IOGenM StdGen
gen <- StdGen -> IO (IOGenM StdGen)
forall (m :: * -> *) g. MonadIO m => g -> m (IOGenM g)
newIOGenM (StdGen -> IO (IOGenM StdGen)) -> IO StdGen -> IO (IOGenM StdGen)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO StdGen
forall (m :: * -> *). MonadIO m => m StdGen
newStdGen
  let n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
  MVector RealWorld a
v <- Vector a -> IO (MVector (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Vector a -> m (MVector (PrimState m) a)
V.thaw ([a] -> Vector a
forall a. [a] -> Vector a
V.fromList [a]
xs)
  [Int] -> (Int -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 .. Int
1] ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
    Int
j <- (Int, Int) -> IOGenM StdGen -> IO Int
forall a g (m :: * -> *).
(UniformRange a, StatefulGen g m) =>
(a, a) -> g -> m a
forall g (m :: * -> *). StatefulGen g m => (Int, Int) -> g -> m Int
uniformRM (Int
0, Int
i) IOGenM StdGen
gen
    MVector (PrimState IO) a -> Int -> Int -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> Int -> Int -> m ()
M.swap MVector RealWorld a
MVector (PrimState IO) a
v Int
i Int
j
  Vector a -> [a]
forall a. Vector a -> [a]
V.toList (Vector a -> [a]) -> IO (Vector a) -> IO [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState IO) a -> IO (Vector a)
forall (m :: * -> *) a.
PrimMonad m =>
MVector (PrimState m) a -> m (Vector a)
V.freeze MVector RealWorld a
MVector (PrimState IO) a
v