{- | This module contains low-level utilities around file seeking.

later this module can be renamed / moved to an internal module.
-}
module DataFrame.IO.Parquet.Seeking (
    SeekableHandle (getSeekableHandle),
    SeekMode (..),
    FileBufferedOrSeekable (..),
    ForceNonSeekable,
    mkFileBufferedOrSeekable,
    mkSeekableHandle,
    readLastBytes,
    withFileBufferedOrSeekable,
    fSeek,
    fGet,
) where

import Control.Monad
import qualified Data.ByteString as BS
import Data.ByteString.Unsafe (unsafeDrop, unsafeTake)
import Data.IORef
import Data.Int
import System.IO

{- | This handle carries a proof that it must be seekable.
Note: Handle and SeekableHandle are not thread safe, should not be
shared across threads, beaware when running parallel/concurrent code.

Not seekable:
  - stdin / stdout
  - pipes / FIFOs

But regular files are always seekable. Parquet fundamentally wants random
access, a non-seekable source will not support effecient access without
buffering the entire file.
-}
newtype SeekableHandle = SeekableHandle {SeekableHandle -> Handle
getSeekableHandle :: Handle}

{- | If we truely want to support non-seekable files, we need to also consider the case
to buffer the entire file in memory.

Not thread safe, contains mutable reference (as Handle already is).

If we need concurrent / parallel parsing or something, we need to read into ByteString
first, not sharing the same handle.
-}
data FileBufferedOrSeekable
    = FileBuffered !(IORef Int64) !BS.ByteString
    | FileSeekable !SeekableHandle

-- | Smart constructor for SeekableHandle
mkSeekableHandle :: Handle -> IO (Maybe SeekableHandle)
mkSeekableHandle :: Handle -> IO (Maybe SeekableHandle)
mkSeekableHandle Handle
h = do
    Bool
seekable <- Handle -> IO Bool
hIsSeekable Handle
h
    Maybe SeekableHandle -> IO (Maybe SeekableHandle)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe SeekableHandle -> IO (Maybe SeekableHandle))
-> Maybe SeekableHandle -> IO (Maybe SeekableHandle)
forall a b. (a -> b) -> a -> b
$ if Bool
seekable then SeekableHandle -> Maybe SeekableHandle
forall a. a -> Maybe a
Just (Handle -> SeekableHandle
SeekableHandle Handle
h) else Maybe SeekableHandle
forall a. Maybe a
Nothing

-- | For testing only
type ForceNonSeekable = Maybe Bool

{- | Smart constructor for FileBufferedOrSeekable, tries to keep in the seekable case
if possible.
-}
mkFileBufferedOrSeekable ::
    ForceNonSeekable -> Handle -> IO FileBufferedOrSeekable
mkFileBufferedOrSeekable :: ForceNonSeekable -> Handle -> IO FileBufferedOrSeekable
mkFileBufferedOrSeekable ForceNonSeekable
forceNonSeek Handle
h = do
    Bool
seekable <- Handle -> IO Bool
hIsSeekable Handle
h
    if Bool -> Bool
not Bool
seekable Bool -> Bool -> Bool
|| ForceNonSeekable
forceNonSeek ForceNonSeekable -> ForceNonSeekable -> Bool
forall a. Eq a => a -> a -> Bool
== Bool -> ForceNonSeekable
forall a. a -> Maybe a
Just Bool
True
        then IORef Int64 -> ByteString -> FileBufferedOrSeekable
FileBuffered (IORef Int64 -> ByteString -> FileBufferedOrSeekable)
-> IO (IORef Int64) -> IO (ByteString -> FileBufferedOrSeekable)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int64 -> IO (IORef Int64)
forall a. a -> IO (IORef a)
newIORef Int64
0 IO (ByteString -> FileBufferedOrSeekable)
-> IO ByteString -> IO FileBufferedOrSeekable
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Handle -> IO ByteString
BS.hGetContents Handle
h
        else FileBufferedOrSeekable -> IO FileBufferedOrSeekable
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (FileBufferedOrSeekable -> IO FileBufferedOrSeekable)
-> FileBufferedOrSeekable -> IO FileBufferedOrSeekable
forall a b. (a -> b) -> a -> b
$ SeekableHandle -> FileBufferedOrSeekable
FileSeekable (SeekableHandle -> FileBufferedOrSeekable)
-> SeekableHandle -> FileBufferedOrSeekable
forall a b. (a -> b) -> a -> b
$ Handle -> SeekableHandle
SeekableHandle Handle
h

{- | With / bracket pattern for FileBufferedOrSeekable

Warning: do not return the FileBufferedOrSeekable outside the scope of the action as
it will be closed.
-}
withFileBufferedOrSeekable ::
    ForceNonSeekable ->
    FilePath ->
    IOMode ->
    (FileBufferedOrSeekable -> IO a) ->
    IO a
withFileBufferedOrSeekable :: forall a.
ForceNonSeekable
-> FilePath -> IOMode -> (FileBufferedOrSeekable -> IO a) -> IO a
withFileBufferedOrSeekable ForceNonSeekable
forceNonSeek FilePath
path IOMode
ioMode FileBufferedOrSeekable -> IO a
action = FilePath -> IOMode -> (Handle -> IO a) -> IO a
forall r. FilePath -> IOMode -> (Handle -> IO r) -> IO r
withFile FilePath
path IOMode
ioMode ((Handle -> IO a) -> IO a) -> (Handle -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Handle
h -> do
    FileBufferedOrSeekable
fbos <- ForceNonSeekable -> Handle -> IO FileBufferedOrSeekable
mkFileBufferedOrSeekable ForceNonSeekable
forceNonSeek Handle
h
    FileBufferedOrSeekable -> IO a
action FileBufferedOrSeekable
fbos

{- | Read the last @n@ bytes, useful for reading metadata without loading the
entire file. Uses 'BS.hGet' (not @hGetContents@) so the handle stays open for
the subsequent column-chunk reads.
-}
readLastBytes :: Integer -> FileBufferedOrSeekable -> IO BS.ByteString
readLastBytes :: Integer -> FileBufferedOrSeekable -> IO ByteString
readLastBytes Integer
n (FileSeekable SeekableHandle
sh) = do
    let h :: Handle
h = SeekableHandle -> Handle
getSeekableHandle SeekableHandle
sh
    Handle -> SeekMode -> Integer -> IO ()
hSeek Handle
h SeekMode
SeekFromEnd (Integer -> Integer
forall a. Num a => a -> a
negate Integer
n)
    Handle -> Int -> IO ByteString
BS.hGet Handle
h (Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n)
readLastBytes Integer
n (FileBuffered IORef Int64
i ByteString
bs) = do
    IORef Int64 -> Int64 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int64
i (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int64) -> Int -> Int64
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
bs)
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
> Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.length ByteString
bs)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> IO ()
forall a. HasCallStack => FilePath -> a
error FilePath
"lastBytes: n > length bs"
    ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
BS.drop (ByteString -> Int
BS.length ByteString
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n) ByteString
bs

fSeek :: FileBufferedOrSeekable -> SeekMode -> Integer -> IO ()
fSeek :: FileBufferedOrSeekable -> SeekMode -> Integer -> IO ()
fSeek (FileSeekable (SeekableHandle Handle
h)) SeekMode
seekMode Integer
seekTo = Handle -> SeekMode -> Integer -> IO ()
hSeek Handle
h SeekMode
seekMode Integer
seekTo
fSeek (FileBuffered IORef Int64
i ByteString
_bs) SeekMode
AbsoluteSeek Integer
seekTo = IORef Int64 -> Int64 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int64
i (Integer -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
seekTo)
fSeek (FileBuffered IORef Int64
i ByteString
_bs) SeekMode
RelativeSeek Integer
seekTo = IORef Int64 -> (Int64 -> Int64) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef Int64
i (Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Integer -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
seekTo)
fSeek (FileBuffered IORef Int64
i ByteString
bs) SeekMode
SeekFromEnd Integer
seekTo = IORef Int64 -> Int64 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Int64
i (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int64) -> Int -> Int64
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Integer -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
seekTo)

fGet :: FileBufferedOrSeekable -> Int -> IO BS.ByteString
fGet :: FileBufferedOrSeekable -> Int -> IO ByteString
fGet (FileSeekable (SeekableHandle Handle
h)) Int
n = Handle -> Int -> IO ByteString
BS.hGet Handle
h Int
n
fGet (FileBuffered IORef Int64
iRef ByteString
bs) Int
n
    | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = do
        Int
i <- Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> IO Int64 -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef Int64 -> IO Int64
forall a. IORef a -> IO a
readIORef IORef Int64
iRef
        if (ByteString -> Int
BS.length ByteString
bs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
            then if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= ByteString -> Int
BS.length ByteString
bs then ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
unsafeDrop Int
i ByteString
bs else ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
            else ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString)
-> (ByteString -> ByteString) -> ByteString -> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
unsafeTake Int
n (ByteString -> ByteString)
-> (ByteString -> ByteString) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ByteString -> ByteString
unsafeDrop Int
i (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString
bs
    | Bool
otherwise = FilePath -> IO ByteString
forall a. HasCallStack => FilePath -> a
error FilePath
"Can't read a negative number of bytes"