--
-- INFOB3CC Concurrency
-- Practical 1: IBAN calculator
--
-- http://ics.uu.nl/docs/vakken/b3cc/assessment.html
--
module IBAN
  ( Mode (..),
    Config (..),
    count,
    list,
    search,
  )
where

import Control.Concurrent
import Crypto.Hash.SHA1
import Data.Atomics (casIORef, peekTicket, readForCAS)
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B8
import Data.IORef
import System.IO

-- -----------------------------------------------------------------------------
-- 0. m-test
-- -----------------------------------------------------------------------------

digits :: Int -> [Int]
digits n
  | n < 10 = [n]
  | otherwise = n `mod` 10 : digits (n `div` 10)

-- Perform the m-test on 'number'
mtest :: Int -> Int -> Bool
mtest m number = finalSum `mod` m == 0
  where
    finalSum = sum $ zipWith (*) (digits number) [1 ..]

-- -----------------------------------------------------------------------------
-- 0.5 My own helper functions
-- -----------------------------------------------------------------------------

-- This function is used the split the range in nearly equal sections
sectionBorder :: Int -> Int -> Int -> Int -> Int
sectionBorder b e n i = (min nHigh i) * (s + 1) + (max 0 (i - nHigh)) * s + b
  where
    -- Here we try to find the section this thread
    -- has to go through. We want that the sections
    -- to be divided as equally as possible between
    -- the threads and that there are no gaps or
    -- overlap. Often it won't be possible to divide
    -- it into equal sections. Some section will be
    -- of size s and some sections will be of size
    -- s+1. We denote nLow as te amount of sections
    -- that have size s and nHigh as the amount of
    -- sections with size s+1. Note that nHigh+nLow=n
    -- and we have the following equation
    --
    -- nLow*s+nHigh*(s+1) = e-b+1
    --
    -- This we can rewrite as
    --
    -- n*s+nHigh = e-b +1
    --
    -- hence
    s = (e - b) `div` n
    nHigh = e - b - n * s

addValue :: IORef Int -> Int -> IO ()
addValue ioref value = do
  t <- readForCAS ioref
  addValue' t
  where
    addValue' t = do
      (success, current) <- casIORef ioref t ((peekTicket t) + value)
      if success
        then return ()
        else -- Try again untill it succeeds
          addValue' current

-- MTRO stands for multi threaded range operation. I wrote this because there
-- are a lot of similarities between the count function and list function.
-- Without this function I would have to copy code.
mtro :: Config -> (Int -> IO ()) -> IO ()
mtro config op = do forkThreads n mtroThread
  where
    b = cfgLower config
    m = cfgModulus config
    e = cfgUpper config
    n = cfgThreads config

    mtroThread i = do
      -- bSec and eSec represent the range
      -- that it will be looking at
      evaluate bSec
      evaluate eSec

      mtroSection bSec eSec
      where
        bSec = sectionBorder b e n i
        eSec = sectionBorder b e n (i + 1)
    mtroSection bsec esec = sequence_ $ map f [bsec .. (esec - 1)]
      where
        f number
          | mtest m number = op number
          | otherwise = return ()

-- -----------------------------------------------------------------------------
-- 1. Counting mode (3pt)
-- -----------------------------------------------------------------------------

count :: Config -> IO Int
count config = do
  countref <- newIORef 0
  mtro config (const (addValue countref 1))
  readIORef countref

-- -----------------------------------------------------------------------------
-- 2. List mode (3pt)
-- -----------------------------------------------------------------------------

list :: Handle -> Config -> IO ()
list handle config = do
  countm <- newMVar 0
  mtro config (listOp countm)
  where
    listOp :: MVar Int -> Int -> IO ()
    listOp countm number = do
      counter <- takeMVar countm
      hPutStrLn handle $ show (counter + 1) ++ " " ++ show number
      putMVar countm (counter + 1)

-- -----------------------------------------------------------------------------
-- 3. Search mode (4pt)
-- -----------------------------------------------------------------------------

chunkSize :: Int
chunkSize = 100

data Range = Range Int Int
  deriving (Show, Eq)

-- This value will be put in the beginning
-- of the queue to make clear where
-- the queue ends
dummyRange :: Range
dummyRange = Range (-1) (-1)

data Queue a = Queue (MVar (List a)) (MVar (List a))

type List a = MVar (Item a)

data Item a = Item a (List a)

-- The newQueue is the same as in the
-- slides except that it adds the dummyRange
-- in the beginning
newQueue :: IO (Queue Range)
newQueue = do
  hole <- newEmptyMVar
  dummy <- newMVar (Item (dummyRange) hole)
  readLock <- newMVar dummy
  writeLock <- newMVar hole
  return (Queue readLock writeLock)

-- Enqueue is basically
-- the same as in the slides
enqueue :: a -> Queue a -> IO ()
enqueue x (Queue _ writeLock) =
  do
    writeEnd <- takeMVar writeLock
    newWriteEnd <- newEmptyMVar
    putMVar writeEnd (Item x newWriteEnd)
    putMVar writeLock newWriteEnd

-- Returns Nothing if the queue is empty
dequeue :: Queue Range -> IO (Maybe Range)
dequeue queue@(Queue readLock _) =
  do
    readEnd <- takeMVar readLock
    Item x newReadEnd <- takeMVar readEnd
    if x == dummyRange
      then do
        enqueue x queue
        putMVar readLock newReadEnd
        return Nothing
      else do
        putMVar readLock newReadEnd
        return (Just x)

searchThread :: Int -> IORef Int -> Queue Range -> ByteString -> MVar Int -> Int -> IO ()
searchThread m amountToDo queue query searchNumber _ = searchThread'
  where
    searchThread' = do
      amToDo <- readIORef amountToDo
      notFound <- isEmptyMVar searchNumber
      if (notFound && amToDo /= 0)
        then do
          maybeRange <- dequeue queue
          case maybeRange of
            -- One might think that when the queue is empty
            -- that the program is finished. However, this is
            -- not the case. It is possible that another thread
            -- has temporarily taken the last element out of the
            -- queue.
            Nothing -> pure ()
            Just r -> do
              searchrange <- splitRange r
              searchRange searchrange
          searchThread'
        else return ()
    splitRange :: Range -> IO Range
    splitRange (oldrange@(Range begin end)) = do
      let searchrange
            | end - begin > chunkSize = Range begin (begin + chunkSize)
            | otherwise = oldrange
      if end - begin > chunkSize
        then do
          let range1 = Range (begin + chunkSize) ((begin + chunkSize + end) `div` 2)
          let range2 = Range ((begin + chunkSize + end) `div` 2) end

          evaluate range1
          evaluate range2

          enqueue range1 queue
          enqueue range2 queue
          addValue amountToDo 1
        else do
          addValue amountToDo (-1)
      return searchrange

    searchRange :: Range -> IO ()
    searchRange (Range begin end) = sequence_ (map checkAcountnumber [begin .. (end - 1)])
      where
        checkAcountnumber number
          | mtest m number && checkHash query (show number) = putMVar searchNumber number
          | otherwise = pure ()

search :: Config -> ByteString -> IO (Maybe Int)
search config query =
  do
    queue <- newQueue

    -- Keeps track how far to program has come with
    -- traversing the queue. The program stops when
    -- amountToDo == 0
    amountToDo <- newIORef 1

    enqueue (Range b e) queue
    searchNumber <- newEmptyMVar
    forkThreads n (searchThread m amountToDo queue query searchNumber)
    tryTakeMVar searchNumber
  where
    b = cfgLower config
    e = cfgUpper config
    m = cfgModulus config
    n = cfgThreads config

-- -----------------------------------------------------------------------------
-- Starting framework
-- -----------------------------------------------------------------------------

data Mode = Count | List | Search ByteString
  deriving (Show)

data Config = Config
  { cfgLower :: !Int,
    cfgUpper :: !Int,
    cfgModulus :: !Int,
    cfgThreads :: !Int
  }
  deriving (Show)

-- Evaluates a term, before continuing with the next IO operation.
--
evaluate :: a -> IO ()
evaluate x = x `seq` return ()

-- Forks 'n' threads. Waits until those threads have finished. Each thread
-- runs the supplied function given its thread ID in the range [0..n).
--
forkThreads :: Int -> (Int -> IO ()) -> IO ()
forkThreads n work = do
  -- Fork the threads and create a list of the MVars which
  -- per thread tell whether the work has finished.
  finishVars <- mapM work' [0 .. n - 1]
  -- Wait on all MVars
  mapM_ takeMVar finishVars
  where
    work' :: Int -> IO (MVar ())
    work' index = do
      var <- newEmptyMVar
      _ <- forkOn index (work index >> putMVar var ())
      return var

-- Checks whether 'value' has the expected hash.
--
checkHash :: ByteString -> String -> Bool
checkHash expected value = expected == hash (B8.pack value)