{-----------------------------------------------------------------------------------------
Module name: BufferedSocket
Made by:     Tomas Möre 2015
 

Usage:
    This module is mean to be imported qualified 

    ex: import Headers qualified BS 

Notes:    
    Buffered sockets are a data type that is a kind of overlay on normal sockets. 

    
    All the exported read / write operations are build such that they ALLWAYS read / write the ammount of bytes requested 
    
    
    This package allows some cases of lazy IO. Some people see Lazy io as the devil incarné. However it is excpected that anyone using this module 
    is cabale of understanding any possible side effects.
    

WARNINGS:
    This module uses a ton of IO and non functional ways of solving problems. 
    This is because we want to be as spacetime efficient as possible. 
    
    This module does NOT contain "beatiful" haskell code
------------------------------------------------------------------------------------------}


{-# LANGUAGE OverloadedStrings #-}

module Network.BufferedSocket.Core
(
  readRaw
, readLazy
, readByte
, readToByte
, readToByteMax
, sendByteString
, readToByteStringMax
, flush
, BufferedSocket 
, makeBufferedSocket
, MaxLength
, ReadSize

, inBuffer
, waitForRead
, isReadable
, isWriteable
, closeRead
, closeWrite
, nativeSocket
  )
where 

import Prelude hiding (getLine, read)



import Control.Monad
import Control.Applicative
import qualified Data.ByteString           as B
import qualified Data.ByteString.Internal  as BI
import qualified Data.ByteString.Lazy      as BL
import qualified Data.ByteString.Builder   as BB 
import qualified Data.ByteString.Char8     as BC8 
import qualified Data.Text as T 
import qualified Data.Text.Lazy as TL 

import Data.List
import Data.IORef
import Data.Functor
import Data.Maybe (isJust)

import qualified Network.Socket as NS -- ns for native socket  
import Foreign.Storable
import Foreign.Ptr
import Foreign.C.Types
import System.IO.Unsafe -- I'm so sorry for this
import System.Posix.Types
import System.Timeout

import Foreign.ForeignPtr
import Data.Monoid
import Data.Word 
import Data.String
import qualified Data.Text.Encoding as ENCSTRICT 
import qualified Data.Text.Lazy.Encoding as ENCLAZY

import GHC.Conc.IO

type MaxLength        = Int

type BufferSize       = Int

type InputBufferSize  = BufferSize
type OutputBufferSize = BufferSize

type BytesInBuffer   = IORef Int 
type ByteOffset      = IORef Int 

type RemainingBytes  = Int   

type ReadSize         = Int 

type Read = Int  

type Timeout = Int 

type SocketData = (NS.Socket,NS.SockAddr,Fd)

type ByteString = B.ByteString
-- Since newtype is disregarded at compile time we make these newtypes to make sure no errors occur 

-- The output buffer is simply a buffer that attemps to store a certain ammount of data before we send the package.
-- This is to make minimalize the TCP overhead. Some operating systems might attemp to do the same. 
type OutputBuffer = (ForeignPtr Word8, BytesInBuffer,BufferSize)
type InputBuffer  = (ForeignPtr Word8, ByteOffset,BytesInBuffer,BufferSize)

newtype BufferedSocket = BufferedSocket (SocketData, InputBuffer, OutputBuffer)

{-
UTIL 
-}


nativeSocket:: BufferedSocket -> NS.Socket
nativeSocket (BufferedSocket ((socket,_,_),_,_))  = socket 

isReadable:: BufferedSocket -> IO Bool 
isReadable (BufferedSocket ((socket,_,_),_,_)) = NS.sIsReadable socket

isWriteable:: BufferedSocket -> IO Bool 
isWriteable (BufferedSocket ((socket,_,_),_,_)) = NS.isWritable socket

closeRead :: BufferedSocket -> IO ()
closeRead (BufferedSocket ((socket,_,_),_,_)) = NS.shutdown socket NS.ShutdownReceive

closeWrite :: BufferedSocket -> IO ()
closeWrite (BufferedSocket ((socket,_,_),_,_)) = NS.shutdown socket NS.ShutdownSend
-- INPUT BUFFER 
inBuffer   :: BufferedSocket -> InputBuffer
inBuffer (BufferedSocket (_,inBuf,_)) = inBuf 

inBufferClear :: InputBuffer -> IO () 
inBufferClear (_, offsetRef, bytesRef, _) = writeIORef offsetRef 0 >> writeIORef bytesRef 0

inBufferReadAll :: InputBuffer -> IO ByteString
inBufferReadAll inBuf@(buf, offsetRef, bytesRef, _) =
    do
        offset <- readIORef offsetRef 
        bytesN <- readIORef bytesRef
        let unreadBytes = bytesN - offset
        if bytesN == 0
            then return "" 
            else (inBufferClear inBuf) >> (return $ BI.fromForeignPtr buf offset unreadBytes)


-- moves the bytes in the buffer to the start 
inBufferRealign :: InputBuffer -> IO ()
inBufferRealign (buf, offsetRef, bytesRef, _) = 
    withForeignPtr buf  $\ptr -> 
        do  putStrLn "realigned"
            offset <- readIORef offsetRef
            when (offset > 0) $
               do   bytesN <- readIORef bytesRef
                    let unreadBytes = bytesN - offset
                        offsetPtr   = plusPtr ptr offset
                    BI.memcpy  ptr offsetPtr unreadBytes
                    writeIORef offsetRef 0 

inBufferFindByteReal :: Ptr Word8 -> Int -> Word8 -> IO (Maybe Int)
inBufferFindByteReal _ 0 _ = return Nothing 
inBufferFindByteReal ptr bytesLeft matchByte = 
    do  currentChar <- peek ptr 
        if isMatch currentChar  
            then return $Just bytesLeft
            else inBufferFindByteReal (plusPtr ptr 1) (bytesLeft - 1) matchByte
    where 
        isMatch = (==matchByte)


inBufferFindByte :: InputBuffer -> Word8 -> IO (Maybe Int)
inBufferFindByte (buf, offsetRef, bytesRef, _) byte =
    do  offset <- readIORef offsetRef 
        bytesN <- readIORef bytesRef 
        let unreadBytes = bytesN - offset

        withForeignPtr buf $\ ptr ->    do  let startPtr = plusPtr ptr offset
                            maybeBytesLeft <- inBufferFindByteReal startPtr unreadBytes byte
                            case maybeBytesLeft of 
                                Just n  -> return $Just  (unreadBytes - n) 
                                Nothing -> return Nothing  

        
-- OUTPUT BUFFER 
outBuffer  :: BufferedSocket -> OutputBuffer  
outBuffer (BufferedSocket (_,_,outBuf)) = outBuf


-- SOCKET DATA 
socketData :: BufferedSocket -> SocketData
socketData (BufferedSocket (sockData,_,_)) = sockData


-- Reads up to max the available bytes of data
bSocketRecv :: BufferedSocket -> IO Int 
bSocketRecv (BufferedSocket ((sock,_,_),(buf,offsetRef,bytesNRef,bufSize),_)) = 
    withForeignPtr buf $\ptr -> 
        do  bytesN <- readIORef bytesNRef 
            let offsetBuf  = plusPtr ptr bytesN
                maxRead    = bufSize - bytesN
            if maxRead > 0 
                then 
                    do 
                        (recievedBytes, _) <- NS.recvBufFrom sock offsetBuf maxRead
                        writeIORef bytesNRef (bytesN + recievedBytes)
                        return recievedBytes
                else 
                    return 0 

-- Unsafe. Do not use if you don't know what you're doing 
bSocketRecvMin :: BufferedSocket -> Int -> IO ()
bSocketRecvMin bSocket n = 
    do  bytesRead <- bSocketRecv bSocket 
        when (bytesRead < n) $ bSocketRecvMin bSocket (n - bytesRead)

waitForRead :: BufferedSocket -> Timeout -> IO Bool 
waitForRead (BufferedSocket ((sock,_,fideDesc),_,_)) timeoutDuration =
    do maybeSucess <- timeout timeoutDuration (threadWaitRead  fideDesc)
       if isJust maybeSucess 
        then return True 
        else return False

{-
CREATING A SOCKET

This puts together a new "BufferedSocket". or commonly refered as "bSock"
A buffered socket is a socket which is acommodated by a Ptr. Because of the need to split headers etc up at certain points but still
safet he rest of the data to be whatever the server requires it to be.

A buffered socket should be called with freeBSocket before disregarded. 
The server thunk function does so automaticly.

-}
--BufferedSocket ((sock, sockAddr, socketFileDesc), inputBuffer, outPutBuffer)
makeBufferedSocket :: (NS.Socket, NS.SockAddr) -> InputBufferSize -> OutputBufferSize -> IO BufferedSocket
makeBufferedSocket (sock, sockAddr)  inBufferSize outBufferSize = 
    do
      inputBuffer      <- makeInputBuffer inBufferSize
      outPutBuffer     <- makeOutputBuffer outBufferSize
      let socketFileDesc = NS.fdSocket sock
      return $BufferedSocket ((sock, sockAddr, Fd socketFileDesc), inputBuffer, outPutBuffer)

makeInputBuffer :: InputBufferSize -> IO InputBuffer
makeInputBuffer bufferSize =
    do  offset      <- newIORef 0 
        bytesCount  <- newIORef 0
        buffer      <- BI.mallocByteString bufferSize
        return (buffer, offset, bytesCount, bufferSize)

makeOutputBuffer :: OutputBufferSize -> IO OutputBuffer
makeOutputBuffer bufferSize = 
    do  bytesCount  <- newIORef 0
        buffer      <- BI.mallocByteString bufferSize
        return (buffer, bytesCount, bufferSize)

{-
USING THE BUFFERED SOCKETS:
-}

{- Strict reading 
If the buffer contains the required ammount of bytes it will simply read it.
If not it will read all the bytes. Clear the buffer and attempt to read more 

-} 
readRaw :: BufferedSocket -> Int -> IO ByteString
readRaw _ 0 = return ""
readRaw bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) readSize = 
    do  bytesN <- readIORef bytesNRef
        offset <- readIORef offsetRef

        let unreadBytes               = bytesN - offset
            availableBytesAfterBytesN = bufSize - bytesN
            availableBytesTotal       = availableBytesAfterBytesN + offset

            missingBytes              = readSize - unreadBytes

        if missingBytes <= 0 
        then writeIORef offsetRef (offset +  readSize) >> (return $BI.fromForeignPtr buf offset readSize)
        else if missingBytes <= availableBytesAfterBytesN 
        then bSocketRecvMin bSock missingBytes >> 
                loop
        else if missingBytes <=  availableBytesTotal
        then inBufferRealign inBuf >> 
                fillBuffer availableBytesTotal >>
                    loop
        else do  
            let strFragment = if unreadBytes > 0 
                                then BI.fromForeignPtr buf offset unreadBytes
                                else ""
            inBufferClear inBuf 
            (strFragment<>) <$> readRaw  bSock (readSize - B.length strFragment)
    where 
        fillBuffer  = bSocketRecvMin bSock
        loop        = readRaw bSock readSize

-- 
lazyReader :: BufferedSocket -> [Int] -> IO [ByteString]
lazyReader _ [] = return []
lazyReader bSock ~(chunkSize:rest) = 
    do  chunk <- unsafeInterleaveIO $ readRaw bSock chunkSize
        next  <- unsafeInterleaveIO $ lazyReader bSock rest 
        return (chunk:next)

-- readLazy will read will read in chunks of the same size of the buffered socket 
readLazy :: BufferedSocket -> Int -> IO BL.ByteString
readLazy _ 0 = return ""
readLazy bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) readSize = 
    let chunkSizes = unfoldr (\b -> if b == 0 
                                    then Nothing 
                                    else if b > bufSize 
                                        then Just (bufSize, b - bufSize) 
                                        else Just (b,0)) readSize
    in  BL.fromChunks <$>lazyReader bSock chunkSizes


readByte :: BufferedSocket -> IO Word8
readByte bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) =
    do  bytesN <- readIORef bytesNRef 
        offset <- readIORef offsetRef
        let unreadBytes = bytesN - offset

        if unreadBytes > 0 
            then do 
                writeIORef offsetRef (offset + 1)
                withForeignPtr buf (`peekByteOff` offset)
            else do 
                inBufferClear inBuf 
                bSocketRecvMin bSock 1 
                readByte bSock

readToByte :: BufferedSocket -> Word8 -> IO ByteString
readToByte bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) byte = 
    do  maybeByteIndex <- inBufferFindByte inBuf byte
        case maybeByteIndex of 
            Nothing -> (<>) <$> inBufferReadAll inBuf <*> readToByte bSock byte --  >>= (<>) fmap (<>(readToByte bSock byte)) $ inBufferReadAll inBuf
            Just n  -> do 
                        offset <- readIORef offsetRef 
                        writeIORef offsetRef (offset + n + 1)
                        return $BI.fromForeignPtr buf offset n

readToByteMax :: BufferedSocket -> Word8 -> MaxLength -> IO (Maybe ByteString)
readToByteMax bSock@(BufferedSocket ((sock,_,_),inBuf@(buf,offsetRef,bytesNRef,bufSize),_)) byte maxLen 
    | hasNoMaxLen = return Nothing 
    | otherwise = 
        do  maybeByteIndex <- inBufferFindByte inBuf byte
            bytesN         <- readIORef bytesNRef
            offset         <- readIORef offsetRef 

            let unreadBytes = bytesN - offset

            case maybeByteIndex of
                Nothing ->  if unreadBytes >= maxLen
                            then return Nothing  
                            else do let maxBytesLeft = maxLen - unreadBytes
                                    thisBufData <- inBufferReadAll inBuf -- empties the current buffer data
                                    bSocketRecv bSock                    -- attempts to rebuffer the socket 
                                    maybeData   <- readToByteMax bSock byte maxBytesLeft -- repeat 
                                    case maybeData of 
                                        Nothing -> return Nothing 
                                        Just a  -> return $Just (thisBufData <> a)
                Just n -> do writeIORef offsetRef (offset + n + 1)
                             return $Just $BI.fromForeignPtr buf offset n 
    where 
        hasNoMaxLen = maxLen <= 0

-- Dangerous will read forever if nothing is found 
-- You will most probably want to use the readToByteStringMax
readToByteString :: BufferedSocket -> ByteString -> IO ByteString
readToByteString _ "" = error "readToByteString can not take an empty bytestring" 
readToByteString bSock searchString =
    do  dataString <- readToByte bSock firstByte
        trail      <- readRaw bSock $B.length restSearchString
        if trail == restSearchString
        then return dataString
        else ((dataString<>trail)<>) <$>readToByteString bSock searchString
    where 
        firstByte        = B.head searchString 
        restSearchString = B.tail searchString

-- Same as above but Will read to a max length. If this is reached It returns "Nothing"
readToByteStringMax :: BufferedSocket -> ByteString -> MaxLength -> IO (Maybe ByteString)
readToByteStringMax _ _ 0  = return Nothing
readToByteStringMax _ "" _ = error "readToByteString max got an empty string"
readToByteStringMax bSock searchString maxLength = 
    do  maybeDataString <- readToByteMax bSock firstByte maxLength
        case maybeDataString of 
            Nothing         -> return Nothing 
            Just dataString -> do   trail <- readRaw bSock $B.length restSearchString
                                    if trail == restSearchString
                                    then return $ Just dataString
                                    else do maybeMoreString <- readToByteStringMax bSock searchString (maxLength - B.length dataString - B.length trail)
                                            case maybeMoreString of 
                                                Nothing -> return Nothing 
                                                Just a  -> return $Just $ dataString <> trail <> a
    where 
        firstByte        = B.head searchString 
        restSearchString = B.tail searchString


-- Simple and clean. force pushed the data to the actual socket. Afther this 
flush :: BufferedSocket -> IO()
flush (BufferedSocket ((socket,sockAddr,_),_,(outBuf, bytesNRef, bufferSize))) =
    do  bytesInBuffer <- readIORef bytesNRef
        when (bytesInBuffer > 0) $
            do  let sender ptr bytes = 
                        do  sent <- NS.sendBufTo socket ptr bytes sockAddr
                            when (sent < bytes) $ sender (plusPtr ptr sent) (bytes - sent)

                    str = BI.fromForeignPtr outBuf 0 bytesInBuffer
                withForeignPtr outBuf (`sender` bytesInBuffer)
                writeIORef bytesNRef 0
    

-- if the string sent to this function is less or equal to the ammount of bytes left in the buffer 
-- we will simply fill the buffer 
-- (if equal) we will also flush this buffer.
-- If the string is greater then the ammount of bytes left 
-- we will attempt to fill the last bytes in the buffer then 
sendByteString:: BufferedSocket -> ByteString -> IO ()
sendByteString _ "" = return ()
sendByteString bSock@(BufferedSocket ((sock,sockAddr,_),_,(outBuf, bytesNRef, bufferSize))) outputStr = 
    do  bytesInBuffer <- readIORef bytesNRef
        let bytesLeftInBuffer = bufferSize - bytesInBuffer

        if srcLength <= bytesLeftInBuffer
            then do
                withForeignPtr outBuf $ \bufPtr -> 
                    withForeignPtr sourceFrgnPtr $ \srcPtr -> 
                        BI.memcpy (plusPtr bufPtr bytesInBuffer) (plusPtr srcPtr srcOffet) srcLength
                          
                writeIORef bytesNRef (bytesInBuffer + srcLength)
                when (srcLength + bytesInBuffer == bufferSize) $ flush bSock
            else do 

                let (current,rest)  = B.splitAt bytesLeftInBuffer outputStr
                    overflowRest    = div (B.length rest) bufferSize
                    directSendChunk = B.take overflowRest rest 
                    directSendLen   = B.length directSendChunk
                    (directSend,_,_)= BI.toForeignPtrdirectSendChunk

                    toBufferPart    = B.drop overflowRest rest

                sendByteString bSock current 
                when (directSendChunk /= "") $withForeignPtr sourceFrgnPtr $\ptr -> void $ NS.sendBufTo sock (plusPtr ptr srcOffet) directSendLen sockAddr
                sendByteString bSock toBufferPart

    where 
        (sourceFrgnPtr, srcOffet, srcLength) = BI.toForeignPtr outputStr