{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveGeneric #-}
module Codec.Xlsx.Writer.Internal.PivotTable
  ( PivotTableFiles(..)
  , renderPivotTableFiles
  ) where

import Data.ByteString.Lazy (ByteString)
import Data.List (elemIndex, transpose)
import Data.List.Extra (nubOrd)
import qualified Data.Map as M
import Data.Maybe (catMaybes, fromMaybe, mapMaybe)
import Data.Text (Text)
import GHC.Generics (Generic)
import Safe (fromJustNote)
import Text.XML

import Codec.Xlsx.Types.Cell
import Codec.Xlsx.Types.Common
import Codec.Xlsx.Types.Internal
import Codec.Xlsx.Types.Internal.Relationships (odr)
import Codec.Xlsx.Types.PivotTable
import Codec.Xlsx.Types.PivotTable.Internal
import Codec.Xlsx.Writer.Internal

data PivotTableFiles = PivotTableFiles
  { pvtfTable :: ByteString
  , pvtfCacheDefinition :: ByteString
  , pvtfCacheRecords :: ByteString
  } deriving (Eq, Show, Generic)

data CacheDefinition = CacheDefinition
  { cdSourceRef :: CellRef
  , cdSourceSheet :: Text
  , cdFields :: [CacheField]
  } deriving (Eq, Show, Generic)

renderPivotTableFiles :: CellMap -> Int -> PivotTable -> PivotTableFiles
renderPivotTableFiles cm cacheId t = PivotTableFiles {..}
  where
    pvtfTable = renderLBS def $ ptDefinitionDocument cacheId cache t
    cache = generateCache cm t
    (cacheDoc, cacheRecordsDoc) = writeCache cache
    pvtfCacheDefinition = renderLBS def cacheDoc
    pvtfCacheRecords = renderLBS def cacheRecordsDoc

ptDefinitionDocument :: Int -> CacheDefinition -> PivotTable -> Document
ptDefinitionDocument cacheId cache t =
    documentFromElement "Pivot table generated by xlsx" $
    ptDefinitionElement "pivotTableDefinition" cacheId cache t

ptDefinitionElement :: Name -> Int -> CacheDefinition -> PivotTable -> Element
ptDefinitionElement nm cacheId cache PivotTable {..} =
  elementList nm attrs elements
  where
    attrs =
      catMaybes
        [ "colGrandTotals" .=? justFalse _pvtColumnGrandTotals
        , "rowGrandTotals" .=? justFalse _pvtRowGrandTotals
        , "outline" .=? justTrue _pvtOutline
        , "outlineData" .=? justTrue _pvtOutlineData
        ] ++
      [ "name" .= _pvtName
      , "dataCaption" .= _pvtDataCaption
      , "cacheId" .= cacheId
      , "dataOnRows" .= (DataPosition `elem` _pvtRowFields)
      ]
    elements = [location, pivotFields, rowFields, colFields, dataFields]
    location =
      leafElement
        "location"
        [ "ref" .= _pvtLocation
          -- TODO : set proper
        , "firstHeaderRow" .= (1 :: Int)
        , "firstDataRow" .= (2 :: Int)
        , "firstDataCol" .= (1 :: Int)
        ]
    name2x = M.fromList $ zip (mapMaybe _pfiName _pvtFields) [0 ..]
    mapFieldToX f = fromJustNote "no field" $ M.lookup f name2x
    pivotFields = elementListSimple "pivotFields" $ map pFieldEl _pvtFields
    maybeFieldIn Nothing _ = False
    maybeFieldIn (Just name) positions = FieldPosition name `elem` positions
    pFieldEl PivotFieldInfo { _pfiName = fName
                            , _pfiOutline = outline
                            , _pfiSortType = sortType
                            , _pfiHiddenItems = hidden
                            }
      | fName `maybeFieldIn` _pvtRowFields =
        pFieldEl' fName outline ("axisRow" :: Text) hidden sortType
      | fName `maybeFieldIn` _pvtColumnFields =
        pFieldEl' fName outline ("axisCol" :: Text) hidden sortType
      | otherwise =
        leafElement "pivotField" $
        [ "dataField" .= True
        , "showAll" .= False
        , "outline" .= outline] ++
        catMaybes ["name" .=? fName]
    pFieldEl' fName outline axis hidden sortType =
      elementList
        "pivotField"
        ([ "axis" .= axis
         , "showAll" .= False
         , "outline" .= outline
         ] ++
         catMaybes [ "name" .=? fName
                   , "sortType" .=? justNonDef FieldSortManual sortType])
        [ elementListSimple "items" $
          items fName hidden ++
          [leafElement "item" ["t" .= ("default" :: Text)]]
        ]
    items Nothing _ = []
    items (Just fName) hidden =
      [ itemEl x item hidden
      | (x, item) <- zip [0 ..] . fromMaybe [] $ M.lookup fName cachedItems
      ]
    itemEl x item hidden =
      leafElement "item" $
      ["x" .= (x :: Int)] ++ catMaybes ["h" .=? justTrue (item `elem` hidden)]
    cachedItems =
      M.fromList $ [(cfName, cfItems) | CacheField {..} <- cdFields cache]
    rowFields =
      elementListSimple "rowFields" . map fieldEl $
      if length _pvtDataFields > 1
        then _pvtRowFields
        else filter (/= DataPosition) _pvtRowFields
    colFields = elementListSimple "colFields" $ map fieldEl _pvtColumnFields
    fieldEl p = leafElement "field" ["x" .= fieldPos p]
    fieldPos DataPosition = (-2) :: Int
    fieldPos (FieldPosition f) = mapFieldToX f
    dataFields = elementListSimple "dataFields" $ map dFieldEl _pvtDataFields
    dFieldEl DataField {..} =
      leafElement "dataField" $
      catMaybes
        [ "name" .=? Just _dfName
        , "fld" .=? Just (mapFieldToX _dfField)
        , "subtotal" .=? justNonDef ConsolidateSum _dfFunction
        ]

generateCache :: CellMap -> PivotTable -> CacheDefinition
generateCache cm PivotTable {..} =
  CacheDefinition
  { cdSourceRef = _pvtSrcRef
  , cdSourceSheet = _pvtSrcSheet
  , cdFields = cachedFields
  }
  where
    cachedFields = mapMaybe (fmap cache . _pfiName) _pvtFields
    cache name =
      CacheField
      { cfName = name
      , cfItems =
          fromJustNote "specified pivot table field does not exist" $
          M.lookup name itemsByName
      }
    ((r1, c1), (r2, c2)) =
      fromJustNote "Invalid src ref of pivot table " $ fromRange _pvtSrcRef
    getCellValue ix = M.lookup ix cm >>= _cellValue
    itemsByName =
      M.fromList $
      flip mapMaybe [c1 .. c2] $ \c -> do
        CellText nm <- getCellValue (r1, c)
        let values = mapMaybe (\r -> getCellValue (r, c)) [(r1 + 1) .. r2]
        return (PivotFieldName nm, nubOrd values)

writeCache :: CacheDefinition -> (Document, Document)
writeCache CacheDefinition {..} = (cacheDefDoc, cacheRecordsDoc)
  where
    cacheDefDoc =
      documentFromElement "Pivot cache definition generated by xlsx" $
      elementList "pivotCacheDefinition" attrs elements
    attrs = ["invalid" .= True, "refreshOnLoad" .= True, odr "id" .= unsafeRefId 1]
    elements = [worksheetSource, cacheFields]
    worksheetSource =
      elementList
        "cacheSource"
        ["type" .= ("worksheet" :: Text)]
        [ leafElement
            "worksheetSource"
            ["ref" .= cdSourceRef, "sheet" .= cdSourceSheet]
        ]
    cacheFields =
      elementListSimple "cacheFields" $ map (toElement "cacheField") cdFields
    cacheRecordsDoc =
      documentFromElement "Pivot cache records generated by xlsx" .
      elementListSimple "pivotCacheRecords" $
      map (elementListSimple "r" . map recordValueToEl) cacheRecords
    recordValueToEl (CacheText t) = leafElement "s" ["v" .= t]
    recordValueToEl (CacheNumber n) = leafElement "n" ["v" .= n]
    recordValueToEl (CacheIndex i) = leafElement "x" ["v" .= i]
    cacheRecords = transpose $ map (itemsToRecordValues . cfItems) cdFields
    itemsToRecordValues vals =
      if all isText vals
        then indexes vals
        else map itemToRecordValue vals
    isText (CellText _) = True
    isText _ = False
    indexes vals =
      [ CacheIndex . fromJustNote "inconsistend definition" $ elemIndex v vals
      | v <- vals
      ]
    itemToRecordValue (CellDouble d) = CacheNumber d
    itemToRecordValue (CellText t) = CacheText t
    itemToRecordValue v = error $ "Unsupported value for pivot tables: " ++ show v