{-# LANGUAGE OverloadedStrings #-}
import           Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as BS
import qualified Data.Vector as Vector
import           Data.Vector (Vector)
import qualified Data.Vector.Unboxed as U
import Data.List(intercalate)
import Data.Maybe(fromMaybe)
import System.Environment(getArgs)

data CSV = CSV { bytes :: !ByteString, rows :: !(U.Vector Int) }

csvFromBytes :: ByteString -> CSV
csvFromBytes bs =
  CSV { bytes = bs
      , rows  = U.fromList (0 : map (1 +) (BS.elemIndices '\n' bs))
                -- one extra
      }

csvRowBytes :: CSV -> Int -> BS.ByteString
csvRowBytes csv n =
  fromMaybe BS.empty $
  do let ix i = rows csv U.!? i
     startIx <- ix n
     nextIx  <- ix (n + 1)
     let len = nextIx - 1 - startIx
     return $! BS.take len (BS.drop startIx (bytes csv))

csvRow :: CSV -> Int -> Vector ByteString
csvRow csv n = Vector.fromList (BS.split ',' (csvRowBytes csv n))

csvRowNum :: CSV -> Int
csvRowNum csv = U.length (rows csv) - 1

-- XXX: one could write a much better fielter to save work.
csvFilter :: (Vector ByteString -> Bool) -> CSV -> CSV
csvFilter p csv = csv { rows = U.ifilter keep (rows csv) }
  where
  len      = U.length (rows csv)
  keep l _ = l == len - 1 || p (csvRow csv l)

main :: IO ()
main = do args <- getArgs
          case args of
            ["-g",file] -> generateSample file
            [file]      -> do csv <- csvFromBytes <$> BS.readFile file
                              print (csvRowNum csv)
                              let rockspring fs = fs Vector.!? 8 == Just "RKS"
                              print (csvRowNum (csvFilter rockspring csv))
            _ -> putStrLn "USAGE: FILE, or -g FILE to generate a sample"


generateSample :: FilePath -> IO ()
generateSample file = writeFile file (unlines (replicate recNum sampleRecord))
  where
  recNum       = 200 * 1000 * 10   -- ~200Mb
  sampleRecord = intercalate "," (replicate 25 "RKS")  -- ~100 bytes