Stream.hs

{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StrictData          #-}
{-# LANGUAGE TemplateHaskell     #-}

-- | Writes Excel files from a stream, which allows creation of
--   large Excel files while remaining in constant memory.
module Codec.Xlsx.Writer.Stream
  ( writeXlsx
  , writeXlsxMultipleSheets
  , writeXlsxWithSharedStrings
  , writeXlsxWithConfig
  , SheetWriteSettings(..)
  , defaultSettings
  , wsSheetView
  , wsZip
  , wsColumnProperties
  , wsRowProperties
  , wsStyles
  , WriteXlsxConfig
  , defWriteXlsxConfig
  , setWriteSettings
  , setSharedStringsMap
  , setNamesAndConduits
  , setNameInfosAndConduits
  -- *** Shared strings
  , sharedStrings
  , sharedStringsStream
  ) where

import Codec.Archive.Zip.Conduit.UnZip
import Codec.Archive.Zip.Conduit.Zip
import Codec.Xlsx.Parser.Internal (addSmlNamespace)
import Codec.Xlsx.Parser.Stream
import Codec.Xlsx.Types (ColumnsProperties (..), RowProperties (..),
                         Styles (..), _AutomaticHeight, _CustomHeight,
                         emptyStyles, rowHeightLens, SheetState (..))
import Codec.Xlsx.Types.Cell
import Codec.Xlsx.Types.Common
import Codec.Xlsx.Types.Internal.Relationships (odr, pr, odRelNs)
import Codec.Xlsx.Types.SheetViews
import Codec.Xlsx.Writer.Internal (nonEmptyElListSimple, toAttrVal, toElement,
                                   txtd, txti)
import Codec.Xlsx.Writer.Internal.Stream
import Conduit (PrimMonad, yield, (.|))
import qualified Conduit as C
import Codec.Xlsx.LensCompat
import Control.Monad
import Control.Monad.Catch
import Control.Monad.Reader.Class
import Control.Monad.State.Strict
import Data.ByteString (ByteString)
import Data.ByteString.Builder (Builder)
import Data.Coerce
import Data.Conduit (ConduitT)
import qualified Data.Conduit.List as CL
import Data.Foldable (fold, traverse_)
import Data.List
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Maybe
import Data.Text (Text)
import qualified Data.Text as Text
import Data.Time
import Data.Word
import Data.XML.Types
import Text.Printf
import Text.XML (toXMLElement)
import qualified Text.XML as TXML
import Text.XML.Stream.Render
import Text.XML.Unresolved (elementToEvents)
import Data.Bifunctor (Bifunctor(second))


upsertSharedStrings :: MonadState SharedStringState m => Row -> m [(Text,Int)]
upsertSharedStrings row =
  traverse upsertSharedString items
  where
    items :: [Text]
    items = row ^.. ri_cell_row . traversed . cellValue . _Just . _CellText

-- | Process sheetItems into shared strings structure to be put into
--   'writeXlsxWithSharedStrings'
sharedStrings :: Monad m  => ConduitT Row b m (Map Text Int)
sharedStrings = void sharedStringsStream .| CL.foldMap (uncurry Map.singleton)

-- | creates a unique number for every encountered string in the stream
--   This is used for creating a required structure in the xlsx format
--   called shared strings. Every string get's transformed into a number
--
--   exposed to allow further processing, we also know the map after processing
--   but I don't think conduit provides a way of getting that out.
--   use 'sharedStrings' to just get the map
sharedStringsStream :: Monad m  =>
  ConduitT Row (Text, Int) m (Map Text Int)
sharedStringsStream = fmap (view string_map) $ C.execStateC initialSharedString $
  CL.mapFoldableM upsertSharedStrings

-- | Settings for writing a single sheet.
data SheetWriteSettings = MkSheetWriteSettings
  { _wsSheetView        :: [SheetView]
  , _wsZip              :: ZipOptions -- ^ Enable zipOpt64=True if you intend writing large xlsx files, zip needs 64bit for files over 4gb.
  , _wsColumnProperties :: [ColumnsProperties]
  , _wsRowProperties    :: Map Int RowProperties
  , _wsStyles           :: Styles
  }
instance Show  SheetWriteSettings where
  -- ZipOptions lacks a show instance-}
  show (MkSheetWriteSettings s _ y r _) = printf "MkSheetWriteSettings{ _wsSheetView=%s, _wsColumnProperties=%s, _wsZip=defaultZipOptions, _wsRowProperties=%s }" (show s) (show y) (show r)
makeLenses ''SheetWriteSettings

defaultSettings :: SheetWriteSettings
defaultSettings = MkSheetWriteSettings
  { _wsSheetView = []
  , _wsColumnProperties = []
  , _wsRowProperties = mempty
  , _wsStyles = emptyStyles
  , _wsZip = defaultZipOptions {
  zipOpt64 = False
  -- There is a magick number in the zip archive package,
  -- https://hackage.haskell.org/package/zip-archive-0.4.1/docs/src/Codec.Archive.Zip.html#local-6989586621679055672
  -- if we enable 64bit the number doesn't align causing the test to fail.
  }
  }



-- | Transform a 'Row' stream into a stream that creates the xlsx file format
--   (to be consumed by sinkfile for example)
--  This first runs 'sharedStrings' and then 'writeXlsxWithSharedStrings'.
--  If you want xlsx files this is the most obvious function to use.
--  the others are exposed in case you can cache the shared strings for example.
-- See also 'writeXlsxMultipleSheets'.
writeXlsx :: MonadThrow m
    => PrimMonad m
    => SheetWriteSettings -- ^ use 'defaultSettings'
    -> ConduitT () Row m () -- ^ the conduit producing sheetitems
    -> ConduitT () ByteString m Word64 -- ^ result conduit producing xlsx files
writeXlsx settings sheetC = do
    writeXlsxWithConfig (defWriteXlsxConfig
                        & setWriteSettings settings
                        & setNamesAndConduits [("Sheet1", sheetC)])

-- | Same as 'writeXlsx' but write to multiple sheets.
writeXlsxMultipleSheets :: MonadThrow m
    => PrimMonad m
    => SheetWriteSettings
    -- ^ use 'defaultSettings'.
    -- Currently, all sheets will use the same setting.
    -> [(Text, ConduitT () Row m ())]
    -- ^ the conduits producing sheetitems for each sheet
    -> ConduitT () ByteString m Word64 -- ^ result conduit producing xlsx files
writeXlsxMultipleSheets settings sheets = do
    writeXlsxWithConfig (defWriteXlsxConfig
                        & setWriteSettings settings
                        & setNamesAndConduits sheets)

-- TODO maybe should use bimap instead: https://hackage.haskell.org/package/bimap-0.4.0/docs/Data-Bimap.html
-- it guarantees uniqueness of both text and int
-- | This write Excel file with a shared strings lookup table.
--   It appears that it is optional.
--   Failed lookups will result in valid xlsx.
--   There are several conditions on shared strings,
--
--      1. Every text to int is unique on both text and int.
--      2. Every Int should have a gap no greater than 1. [("xx", 3), ("yy", 4)] is okay, whereas [("xx", 3), ("yy", 5)] is not.
--      3. It's expected this starts from 0.
--
--   Use 'sharedStringsStream' to get a good shared strings table.
--   This is provided because the user may have a more efficient way of
--   constructing this table than the library can provide,
--   for example through database operations.
writeXlsxWithSharedStrings :: MonadThrow m => PrimMonad m
    => SheetWriteSettings
    -> Map Text Int -- ^ shared strings table
    -> [(Text, ConduitT () Row m ())]
    -> ConduitT () ByteString m Word64
writeXlsxWithSharedStrings settings sharedStrings' sheets =
  writeXlsxWithConfig (defWriteXlsxConfig
                      & setWriteSettings settings
                      & setSharedStringsMap (Just sharedStrings')
                      & setNamesAndConduits sheets)

-- | Configure how to write to an xlsx. Provides setters but no getters to
-- hide internals for future compatibility
data WriteXlsxConfig m = MkWriteXlsxConfig
  { _writeSettings :: SheetWriteSettings
  , _sharedStringsMap :: Maybe (Map Text Int)
  , _infosAndConduits :: [((Text, SheetState), ConduitT () Row m ())]
  }

-- | Create a blank config. Uses the default write setting config, derives shared
-- strings from the conduits and has no conduits or sheet informations within.
defWriteXlsxConfig :: WriteXlsxConfig m
defWriteXlsxConfig = MkWriteXlsxConfig defaultSettings mempty []

-- | Overwrite the existing sheet write settings.
setWriteSettings :: SheetWriteSettings -> WriteXlsxConfig m -> WriteXlsxConfig m
setWriteSettings ws (MkWriteXlsxConfig _ ssm iac) = MkWriteXlsxConfig ws ssm iac

-- | Set the shared strings map. If set to Nothing, derives the shared strings
-- from the conduits.
setSharedStringsMap :: Maybe (Map Text Int) -> WriteXlsxConfig m -> WriteXlsxConfig m
setSharedStringsMap ssm (MkWriteXlsxConfig ws _ iac) = MkWriteXlsxConfig ws ssm iac

-- | Set the sheet infos and conduits for each sheet.
setNameInfosAndConduits :: [((Text, SheetState), ConduitT () Row m ())] -> WriteXlsxConfig m -> WriteXlsxConfig m
setNameInfosAndConduits iac (MkWriteXlsxConfig ws ssm _) = MkWriteXlsxConfig ws ssm iac

-- | Set the names and conduits for each sheet. The sheet info is derived from
-- the ordering.
setNamesAndConduits :: [(Text, ConduitT () Row m ())] -> WriteXlsxConfig m -> WriteXlsxConfig m
setNamesAndConduits nacs (MkWriteXlsxConfig ws ssm _) = MkWriteXlsxConfig ws ssm infosAndConduits'
  where
  infosAndConduits' = nacs <&> \(n, c) -> ((n, Visible), c)

-- | Write to an XLSX using the 'WriteXlsxConfig' structure.
writeXlsxWithConfig :: MonadThrow m => PrimMonad m
    => WriteXlsxConfig m
    -> ConduitT () ByteString m Word64
writeXlsxWithConfig (MkWriteXlsxConfig ws ssm iac) = do
  let rowConduits = foldr ((>>) . snd) mempty iac
  sstrings <- maybe (rowConduits .| sharedStrings) pure ssm
  combinedFiles ws sstrings iac .| zipStream (ws ^. wsZip)

combinedFiles :: PrimMonad m
  => SheetWriteSettings
  -> Map Text Int
  -> [((Text, SheetState), ConduitT () Row m ())]
  -> ConduitT () (ZipEntry, ZipData m) m ()
combinedFiles settings sharedStrings' sheets =
  let indicesAndSheets = zip [1..] sheets
      zippedSheets = map (\(sheetId, (_, rowConduit)) ->
        ( zipEntry ("xl/worksheets/sheet" <> Text.pack (show sheetId) <> ".xml")
        , ZipDataSource $ rowConduit .| C.runReaderC settings (writeWorkSheet sharedStrings') .| eventsToBS
        )) indicesAndSheets
  in
  C.yieldMany $
    [ (zipEntry "xl/sharedStrings.xml", ZipDataSource $ writeSst sharedStrings' .| eventsToBS)
    , (zipEntry "[Content_Types].xml", ZipDataSource $ writeContentTypes .| eventsToBS)
    , (zipEntry "xl/workbook.xml", ZipDataSource $ writeWorkbook (map (second fst) indicesAndSheets)
    .| renderBuilder ( def {rsNamespaces=[("r", odRelNs)]})
    .| C.builderToByteString)
    , (zipEntry "xl/styles.xml", ZipDataByteString $ coerce $ settings ^. wsStyles)
    , (zipEntry "xl/_rels/workbook.xml.rels", ZipDataSource $ writeWorkbookRels (length sheets) .| eventsToBS)
    , (zipEntry "_rels/.rels", ZipDataSource $ writeRootRels .| eventsToBS)
    ] <> zippedSheets

el :: Monad m => Name -> Monad m => forall i.  ConduitT i Event m () -> ConduitT i Event m ()
el x = tag x mempty

--   Clark notation is used a lot for xml namespaces in this module:
--   <https://hackage.haskell.org/package/xml-types-0.3.8/docs/Data-XML-Types.html#t:Name>
--   Name has an IsString instance which parses it
override :: Monad m => Text -> Text -> forall i.  ConduitT i Event m ()
override content' part =
    tag "{http://schemas.openxmlformats.org/package/2006/content-types}Override"
      (attr "ContentType" content'
       <> attr "PartName" part) $ pure ()

-- | required by Excel.
writeContentTypes :: Monad m => forall i.  ConduitT i Event m ()
writeContentTypes = doc "{http://schemas.openxmlformats.org/package/2006/content-types}Types" $ do
    override "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet.main+xml" "/xl/workbook.xml"
    override "application/vnd.openxmlformats-officedocument.spreadsheetml.sharedStrings+xml" "/xl/sharedStrings.xml"
    override "application/vnd.openxmlformats-officedocument.spreadsheetml.styles+xml" "/xl/styles.xml"
    override "application/vnd.openxmlformats-officedocument.spreadsheetml.worksheet+xml" "/xl/worksheets/sheet1.xml"
    override "application/vnd.openxmlformats-package.relationships+xml" "/xl/_rels/workbook.xml.rels"
    override "application/vnd.openxmlformats-package.relationships+xml" "/_rels/.rels"

-- | required by Excel.
writeWorkbook :: Monad m => [(Int, (Text, SheetState))] -> forall i.  ConduitT i Event m ()
writeWorkbook sheetInfos =
  let addSheet (sheetId, (sheetName, visibility)) = tag (addSmlNamespace "sheet")
          (attr "name" sheetName
          <> attr "sheetId" (Text.pack $ show sheetId)
          <> attr (odr "id") ("rId" <> (Text.pack $ show (sheetId + 2)))
          <> attr "state" (case visibility of
                              Visible    -> "visible"
                              Hidden     -> "hidden"
                              VeryHidden -> "veryHidden")
          ) $ pure ()
  in doc (addSmlNamespace "workbook") $
    el (addSmlNamespace "sheets") $ mapM_ addSheet sheetInfos

doc :: Monad m => Name ->  forall i.  ConduitT i Event m () -> ConduitT i Event m ()
doc root docM = do
  yield EventBeginDocument
  el root docM
  yield EventEndDocument

relationship :: Monad m => Text -> Int -> Text ->  forall i.  ConduitT i Event m ()
relationship target id' type' =
  tag (pr "Relationship")
    (attr "Type" type'
      <> attr "Id" (Text.pack $ "rId" <> show id')
      <> attr "Target" target
    ) $ pure ()

writeWorkbookRels :: Monad m => forall i. Int -> ConduitT i Event m ()
writeWorkbookRels sheetCount = doc (pr "Relationships") $ do
  relationship "sharedStrings.xml" 1 "http://schemas.openxmlformats.org/officeDocument/2006/relationships/sharedStrings"
  relationship "styles.xml" 2 "http://schemas.openxmlformats.org/officeDocument/2006/relationships/styles"
  let schema = "http://schemas.openxmlformats.org/officeDocument/2006/relationships/worksheet"
  mapM_ (\sheetId -> relationship ("worksheets/sheet" <> (Text.pack (show sheetId)) <> ".xml") (sheetId + 2) schema) [1..sheetCount]

writeRootRels :: Monad m => forall i.  ConduitT i Event m ()
writeRootRels = doc (pr "Relationships") $
  relationship "xl/workbook.xml" 1 "http://schemas.openxmlformats.org/officeDocument/2006/relationships/officeDocument"


zipEntry :: Text -> ZipEntry
zipEntry x = ZipEntry
  { zipEntryName = Left x
  , zipEntryTime = LocalTime (toEnum 0) midnight
  , zipEntrySize = Nothing
  , zipEntryExternalAttributes = Nothing
  }

eventsToBS :: PrimMonad m  => ConduitT Event ByteString m ()
eventsToBS = writeEvents .| C.builderToByteString

writeSst ::  Monad m  => Map Text Int  -> forall i.  ConduitT i Event m ()
writeSst sharedStrings' = doc (addSmlNamespace "sst") $
    void $ traverse (el (addSmlNamespace "si") .  el (addSmlNamespace "t") . content . fst
                  ) $ sortBy (\(_, i) (_, y :: Int) -> compare i y) $ Map.toList sharedStrings'

writeEvents ::  PrimMonad m => ConduitT Event Builder m ()
writeEvents = renderBuilder def

sheetViews :: forall m . MonadReader SheetWriteSettings m => forall i . ConduitT i Event m ()
sheetViews = do
  sheetView <- view wsSheetView

  unless (null sheetView) $ el (addSmlNamespace "sheetViews") $ do
    let
        view' :: [Element]
        view' = setNameSpaceRec spreadSheetNS . toXMLElement .  toElement (addSmlNamespace "sheetView") <$> sheetView

    C.yieldMany $ elementToEvents =<< view'

spreadSheetNS :: Text
spreadSheetNS = fold $ nameNamespace $ addSmlNamespace ""

setNameSpaceRec :: Text -> Element -> Element
setNameSpaceRec space xelm =
    xelm {elementName = ((elementName xelm ){nameNamespace =
                                    Just space })
      , elementNodes = elementNodes xelm <&> \case
                                    NodeElement x -> NodeElement (setNameSpaceRec space x)
                                    y -> y
    }

columns :: MonadReader SheetWriteSettings m => ConduitT Row Event m ()
columns = do
  colProps <- view wsColumnProperties
  let cols :: Maybe TXML.Element
      cols = nonEmptyElListSimple (addSmlNamespace "cols") $ map (toElement (addSmlNamespace "col")) colProps
  traverse_ (C.yieldMany . elementToEvents . toXMLElement) cols

writeWorkSheet :: MonadReader SheetWriteSettings  m => Map Text Int  -> ConduitT Row Event m ()
writeWorkSheet sharedStrings' = doc (addSmlNamespace "worksheet") $ do
    sheetViews
    columns
    el (addSmlNamespace "sheetData") $ C.awaitForever (mapRow sharedStrings')

mapRow :: MonadReader SheetWriteSettings m => Map Text Int -> Row -> ConduitT Row Event m ()
mapRow sharedStrings' sheetItem = do
  mRowProp <- preview $ wsRowProperties . ix (unRowIndex rowIx) . rowHeightLens . _Just . failing _CustomHeight _AutomaticHeight
  let rowAttr :: Attributes
      rowAttr = ixAttr <> fold (attr "ht" . txtd <$> mRowProp)
  tag (addSmlNamespace "row") rowAttr $
    void $ itraverse (mapCell sharedStrings' rowIx) (sheetItem ^. ri_cell_row)
  where
    rowIx = sheetItem ^. ri_row_index
    ixAttr = attr "r" $ toAttrVal rowIx

mapCell ::
  Monad m => Map Text Int -> RowIndex -> Int -> Cell -> ConduitT Row Event m ()
mapCell sharedStrings' rix cix' cell =
  when (has (cellValue . _Just) cell || has (cellStyle . _Just) cell) $
  tag (addSmlNamespace "c") celAttr $
    when (has (cellValue . _Just) cell) $
    el (addSmlNamespace "v") $
      content $ renderCell sharedStrings' cell
  where
    cix = ColumnIndex cix'
    celAttr  = attr "r" ref <>
      renderCellType sharedStrings' cell
      <> foldMap (attr "s" . txti) (cell ^. cellStyle)
    ref :: Text
    ref = coerce $ singleCellRef (rix, cix)

renderCellType :: Map Text Int -> Cell -> Attributes
renderCellType sharedStrings' cell =
  maybe mempty
  (attr "t" . renderType sharedStrings')
  $ cell ^? cellValue . _Just

renderCell :: Map Text Int -> Cell -> Text
renderCell sharedStrings' cell =  renderValue sharedStrings' val
  where
    val :: CellValue
    val = fromMaybe (CellText mempty) $ cell ^? cellValue . _Just

renderValue :: Map Text Int -> CellValue -> Text
renderValue sharedStrings' = \case
  CellText x ->
    -- if we can't find it in the sst, print the string
    maybe x toAttrVal $ sharedStrings' ^? ix x
  CellDouble x -> toAttrVal x
  CellBool b -> toAttrVal b
  CellRich _ -> error "rich text is not supported yet"
  CellError err  -> toAttrVal err


renderType :: Map Text Int -> CellValue -> Text
renderType sharedStrings' = \case
  CellText x ->
    maybe "str" (const "s") $ sharedStrings' ^? ix x
  CellDouble _ -> "n"
  CellBool _ -> "b"
  CellRich _ -> "r"
  CellError _ -> "e"