{-# LANGUAGE OverloadedStrings #-}

module Dhall.Csv.Util (encodeCsvDefault, decodeCsvDefault) where

import Data.Csv  (NamedRecord, Record)
import Data.List (sort)
import Data.Text (Text)

import qualified Data.ByteString.Lazy       as ByteString
import qualified Data.ByteString.Lazy.Char8 as ByteString.Char8
import qualified Data.Csv
import qualified Data.HashMap.Strict        as HashMap
import qualified Data.Text.Encoding
import qualified Data.Vector                as Vector

{-| Utility to convert a list of @NamedRecord@ to Text formatted as a CSV.
-}
encodeCsvDefault :: [NamedRecord] -> Text
encodeCsvDefault :: [NamedRecord] -> Text
encodeCsvDefault [NamedRecord]
csv = ByteString -> Text
Data.Text.Encoding.decodeUtf8 (ByteString -> Text) -> ByteString -> Text
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
ByteString.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Header -> [NamedRecord] -> ByteString
forall a. ToNamedRecord a => Header -> [a] -> ByteString
Data.Csv.encodeByName Header
header [NamedRecord]
csv
  where
    header :: Header
header = case [NamedRecord]
csv of
        [] -> Header
forall a. Vector a
Vector.empty
        (NamedRecord
m:[NamedRecord]
_) -> [ByteString] -> Header
forall a. [a] -> Vector a
Vector.fromList ([ByteString] -> Header) -> [ByteString] -> Header
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
forall a. Ord a => [a] -> [a]
sort ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall a b. (a -> b) -> a -> b
$ NamedRecord -> [ByteString]
forall k v. HashMap k v -> [k]
HashMap.keys NamedRecord
m

{-| Utility to decode a CSV into a list of records.
    Must specify whether the CSV to decode has header or not.
-}
decodeCsvDefault :: Bool -> Text -> Either String [NamedRecord]
decodeCsvDefault :: Bool -> Text -> Either String [NamedRecord]
decodeCsvDefault Bool
hasHeader
    | Bool
hasHeader = Text -> Either String [NamedRecord]
decodeCsvWithHeader
    | Bool
otherwise = Text -> Either String [NamedRecord]
decodeCsvNoHeader

decodeCsvWithHeader :: Text -> Either String [NamedRecord]
decodeCsvWithHeader :: Text -> Either String [NamedRecord]
decodeCsvWithHeader Text
txt = do
    (Header
_, Vector NamedRecord
vec) <- ByteString -> Either String (Header, Vector NamedRecord)
forall a.
FromNamedRecord a =>
ByteString -> Either String (Header, Vector a)
Data.Csv.decodeByName (ByteString -> Either String (Header, Vector NamedRecord))
-> ByteString -> Either String (Header, Vector NamedRecord)
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
ByteString.fromStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
Data.Text.Encoding.encodeUtf8 Text
txt
    [NamedRecord] -> Either String [NamedRecord]
forall (m :: * -> *) a. Monad m => a -> m a
return ([NamedRecord] -> Either String [NamedRecord])
-> [NamedRecord] -> Either String [NamedRecord]
forall a b. (a -> b) -> a -> b
$ Vector NamedRecord -> [NamedRecord]
forall a. Vector a -> [a]
Vector.toList Vector NamedRecord
vec

decodeCsvNoHeader :: Text -> Either String [NamedRecord]
decodeCsvNoHeader :: Text -> Either String [NamedRecord]
decodeCsvNoHeader Text
txt = do
    Vector Header
vec <- HasHeader -> ByteString -> Either String (Vector Header)
forall a.
FromRecord a =>
HasHeader -> ByteString -> Either String (Vector a)
Data.Csv.decode HasHeader
Data.Csv.NoHeader (ByteString -> Either String (Vector Header))
-> ByteString -> Either String (Vector Header)
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
ByteString.fromStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Text -> ByteString
Data.Text.Encoding.encodeUtf8 Text
txt
    [NamedRecord] -> Either String [NamedRecord]
forall (m :: * -> *) a. Monad m => a -> m a
return ([NamedRecord] -> Either String [NamedRecord])
-> [NamedRecord] -> Either String [NamedRecord]
forall a b. (a -> b) -> a -> b
$ (Header -> NamedRecord) -> [Header] -> [NamedRecord]
forall a b. (a -> b) -> [a] -> [b]
map Header -> NamedRecord
addDefaultHeader ([Header] -> [NamedRecord]) -> [Header] -> [NamedRecord]
forall a b. (a -> b) -> a -> b
$ Vector Header -> [Header]
forall a. Vector a -> [a]
Vector.toList Vector Header
vec

addDefaultHeader :: Record -> NamedRecord
addDefaultHeader :: Header -> NamedRecord
addDefaultHeader = [(ByteString, ByteString)] -> NamedRecord
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HashMap.fromList ([(ByteString, ByteString)] -> NamedRecord)
-> (Header -> [(ByteString, ByteString)]) -> Header -> NamedRecord
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([ByteString] -> [ByteString] -> [(ByteString, ByteString)]
forall a b. [a] -> [b] -> [(a, b)]
zip [ByteString]
headerBS) ([ByteString] -> [(ByteString, ByteString)])
-> (Header -> [ByteString]) -> Header -> [(ByteString, ByteString)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Header -> [ByteString]
forall a. Vector a -> [a]
Vector.toList
  where
    header :: [String]
header = (Int -> String) -> [Int] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map ((Char
'_' Char -> String -> String
forall a. a -> [a] -> [a]
:) (String -> String) -> (Int -> String) -> Int -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> String
forall a. Show a => a -> String
show) ([Int
1..] :: [Int])
    headerBS :: [ByteString]
headerBS = (String -> ByteString) -> [String] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (ByteString -> ByteString
ByteString.toStrict (ByteString -> ByteString)
-> (String -> ByteString) -> String -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ByteString
ByteString.Char8.pack) [String]
header