content: module DataFrame.IO.Parquet.Levels where
  
  import Data.Int
  import Data.List
  import qualified Data.Text as T
  import Data.Word
  
  import DataFrame.IO.Parquet.Binary
  import DataFrame.IO.Parquet.Encoding
  import DataFrame.IO.Parquet.Thrift
  import DataFrame.IO.Parquet.Types
  
  readLevelsV1 :: Int -> Int -> Int -> [Word8] -> ([Int], [Int], [Word8])
  readLevelsV1 n maxDef maxRep bs =
      let bwDef = bitWidthForMaxLevel maxDef
          bwRep = bitWidthForMaxLevel maxRep
  
          (repLvlsU32, afterRep) =
              if bwRep == 0
                  then (replicate n 0, bs)
                  else
                      let repLength = littleEndianWord32 (take 4 bs)
                          repData = take (fromIntegral repLength) (drop 4 bs)
                          afterRepData = drop (4 + fromIntegral repLength) bs
                          (repVals, _) = decodeRLEBitPackedHybrid bwRep n repData
                       in (repVals, afterRepData)
  
          (defLvlsU32, afterDef) =
              if bwDef == 0
                  then (replicate n 0, afterRep)
                  else
                      let defLength = littleEndianWord32 (take 4 afterRep)
                          defData = take (fromIntegral defLength) (drop 4 afterRep)
                          afterDefData = drop (4 + fromIntegral defLength) afterRep
                          (defVals, _) = decodeRLEBitPackedHybrid bwDef n defData
                       in (defVals, afterDefData)
       in (map fromIntegral defLvlsU32, map fromIntegral repLvlsU32, afterDef)
  
  readLevelsV2 ::
      Int -> Int -> Int -> Int32 -> Int32 -> [Word8] -> ([Int], [Int], [Word8])
  readLevelsV2 n maxDef maxRep defLen repLen bs =
      let (repBytes, afterRepBytes) = splitAt (fromIntegral repLen) bs
          (defBytes, afterDefBytes) = splitAt (fromIntegral defLen) afterRepBytes
          bwDef = bitWidthForMaxLevel maxDef
          bwRep = bitWidthForMaxLevel maxRep
          (repLvlsU32, _) =
              if bwRep == 0
                  then (replicate n 0, repBytes)
                  else decodeRLEBitPackedHybrid bwRep n repBytes
          (defLvlsU32, _) =
              if bwDef == 0
                  then (replicate n 0, defBytes)
                  else decodeRLEBitPackedHybrid bwDef n defBytes
       in (map fromIntegral defLvlsU32, map fromIntegral repLvlsU32, afterDefBytes)
  
  stitchNullable :: Int -> [Int] -> [a] -> [Maybe a]
  stitchNullable maxDef defLvls vals = go defLvls vals
    where
      go [] _ = []
      go (d : ds) vs
          | d == maxDef = case vs of
              (v : vs') -> Just v : go ds vs'
              [] -> error "value stream exhausted"
          | otherwise = Nothing : go ds vs
  
  data SNode = SNode
      { sName :: String
      , sRep :: RepetitionType
      , sChildren :: [SNode]
      }
      deriving (Show, Eq)
  
  parseOne :: [SchemaElement] -> (SNode, [SchemaElement])
  parseOne [] = error "parseOne: empty schema list"
  parseOne (se : rest) =
      let childCount = fromIntegral (numChildren se)
          (kids, rest') = parseMany childCount rest
       in ( SNode
              { sName = T.unpack (elementName se)
              , sRep = repetitionType se
              , sChildren = kids
              }
          , rest'
          )
  
  parseMany :: Int -> [SchemaElement] -> ([SNode], [SchemaElement])
  parseMany 0 xs = ([], xs)
  parseMany n xs =
      let (node, xs') = parseOne xs
          (nodes, xs'') = parseMany (n - 1) xs'
       in (node : nodes, xs'')
  
  parseAll :: [SchemaElement] -> [SNode]
  parseAll [] = []
  parseAll xs = let (n, xs') = parseOne xs in n : parseAll xs'
  
  levelsForPath :: [SchemaElement] -> [String] -> (Int, Int)
  levelsForPath schemaTail path = go 0 0 (parseAll schemaTail) path
    where
      go defC repC _ [] = (defC, repC)
      go defC repC nodes (p : ps) =
          case find (\n -> sName n == p) nodes of
              Nothing -> (defC, repC)
              Just n ->
                  let defC' = defC + (if sRep n == OPTIONAL then 1 else 0)
                      repC' = repC + (if sRep n == REPEATED then 1 else 0)
                   in go defC' repC' (sChildren n) ps
  