{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE NamedFieldPuns #-}

module Moonlight.Triangulation.Internal.BoxedPaged
  ( BoxedPaged
  , MutableBoxedPaged
  , BoxedStorageError (..)
  , emptyBoxedPaged
  , boxedFill
  , boxedPagedLength
  , boxedUnsafeIndex
  , boxedFromVector
  , boxedToVector
  , thawBoxedPaged
  , readBoxedPaged
  , writeBoxedPaged
  , resetBoxedRange
  , boxedThawPristine
  , freezeBoxedPaged
  , boxedUpdate
  ) where

import Control.DeepSeq (NFData)
import Control.Monad.ST (ST)
import Data.Bits (shiftL, shiftR, (.&.))
import qualified Data.IntMap.Strict as IntMap
import qualified Data.List as List
import Data.STRef
import qualified Data.Vector as V
import qualified Data.Vector.Mutable as MV
import Moonlight.Triangulation.Internal.PageDirectory

import GHC.Generics (Generic)

-- | A paged payload store in which an absent page means every slot in it holds
-- the fill. The fill belongs to the store rather than to a thaw of it: a
-- payload component either has a default that new elements inherit — the three
-- element payloads do — or it has none and every slot must be written before it
-- is read, which is the vertex component, whose payloads arrive with it.
--
-- The near-universal instantiation is @() () ()@, and under an unconditional
-- representation that costs one pointer per directed edge, undirected edge and
-- face, all aimed at the same closure. Here it costs nothing: an element
-- created at its default writes no slot, so no page is ever materialized.
data BoxedStorage a
  = DenseBoxedPages !(V.Vector (V.Vector a))
  | DefaultedBoxedPages !a !(PageDirectory (V.Vector a))
  deriving stock (Generic, Functor)
  deriving anyclass (NFData)

data BoxedPaged a = BoxedPaged
  { boxedLength :: {-# UNPACK #-} !Int
  , boxedStorage :: !(BoxedStorage a)
  }
  deriving stock (Generic, Functor)
  deriving anyclass (NFData)

data BoxedStorageError
  = BoxedFreezeLengthNegative {-# UNPACK #-} !Int
  | BoxedFreezeDensePageMissing {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  deriving stock (Eq, Show, Generic)
  deriving anyclass (NFData)

boxedFill :: BoxedPaged a -> Maybe a
boxedFill BoxedPaged{boxedStorage = DenseBoxedPages _} = Nothing
boxedFill BoxedPaged{boxedStorage = DefaultedBoxedPages fill _} = Just fill

-- Structural equality would distinguish a store that materialized a page of
-- defaults from one that did not, and the two are the same store. Equality is
-- therefore what can be observed: the length, the fill every unwritten slot
-- will report, and the sequence itself.
instance Eq a => Eq (BoxedPaged a) where
  left == right =
    boxedLength left == boxedLength right
      && boxedFill left == boxedFill right
      && boxedToVector left == boxedToVector right

instance Show a => Show (BoxedPaged a) where
  showsPrec precedence paged =
    showParen (precedence > 10) $
      showString "BoxedPaged "
        . showsPrec 11 (boxedFill paged)
        . showString " "
        . showsPrec 11 (boxedToVector paged)

data MutableBoxedPaged s a = MutableBoxedPaged
  { mutableBoxedBase :: !(BoxedPaged a)
  , mutableBoxedDirty :: !(STRef s (IntMap.IntMap (MV.MVector s a)))
  , mutableBoxedFill :: !(Maybe a)
  , -- | The thaw found no materialized page. Together with an empty dirty set
    -- this says every slot in the store reports the fill, which is the state a
    -- payload component nobody writes stays in for the whole transaction —
    -- the near-universal @() () ()@ case. Read through 'boxedThawPristine'.
    mutableBoxedPristineBase :: !Bool
  }

pageBits :: Int
pageBits = 8
{-# INLINE pageBits #-}

pageSize :: Int
pageSize = 1 `shiftL` pageBits
{-# INLINE pageSize #-}

-- Every index this store decomposes is an element identifier or a walk over
-- one, so the signed division a 'quotRem' would emit corrects for a sign that
-- cannot occur.
pageOf :: Int -> Int
pageOf index = index `shiftR` pageBits
{-# INLINE pageOf #-}

offsetOf :: Int -> Int
offsetOf index = index .&. (pageSize - 1)
{-# INLINE offsetOf #-}

emptyBoxedPaged :: Maybe a -> BoxedPaged a
emptyBoxedPaged fill =
  BoxedPaged
    0
    (case fill of
       Nothing -> DenseBoxedPages V.empty
       Just value -> DefaultedBoxedPages value emptyDirectory)

boxedPagedLength :: BoxedPaged a -> Int
boxedPagedLength = boxedLength
{-# INLINE boxedPagedLength #-}

-- | Index with a handle already admitted by the owning triangulation. Dense
-- stores carry every live page explicitly; defaulted stores can answer an
-- absent or short page from their stored default. The handle domain is owned
-- by the DCEL, so this raw kernel performs no second, disagreeing bounds check.
boxedUnsafeIndex :: BoxedPaged a -> Int -> a
boxedUnsafeIndex BoxedPaged{boxedStorage} index =
  let !page = pageOf index
      !offset = offsetOf index
   in case boxedStorage of
        DenseBoxedPages pages ->
          V.unsafeIndex (V.unsafeIndex pages page) offset
        DefaultedBoxedPages fill pages ->
          case lookupDirectory page pages of
            Just values | offset < V.length values -> V.unsafeIndex values offset
            _ -> fill
{-# INLINE boxedUnsafeIndex #-}

boxedFromVector :: Maybe a -> V.Vector a -> BoxedPaged a
boxedFromVector fill values =
  BoxedPaged
    { boxedLength = V.length values
    , boxedStorage =
        case fill of
          Nothing -> DenseBoxedPages (V.fromList (map snd (go 0)))
          Just value -> DefaultedBoxedPages value (directoryFromAscList (go 0))
    }
 where
  go !offset
    | offset >= V.length values = []
    | otherwise =
        let !page = pageOf offset
            !count = min pageSize (V.length values - offset)
            !chunk = V.slice offset count values
         in (page, chunk) : go (offset + count)

-- Resolved a page at a time. Asking the directory per element paid a descent
-- for every slot of a run the descent had already found.
boxedToVector :: BoxedPaged a -> V.Vector a
boxedToVector BoxedPaged{boxedLength, boxedStorage}
  | boxedLength <= 0 = V.empty
  | otherwise = case boxedStorage of
      DenseBoxedPages pages -> V.take boxedLength (V.concat (V.toList pages))
      DefaultedBoxedPages fill pages ->
        V.concat (map (pageRun fill pages) [0 .. pageOf (boxedLength - 1)])
 where
  pageRun fill pages page =
    let !base = page `shiftL` pageBits
        !width = min pageSize (boxedLength - base)
     in case lookupDirectory page pages of
          Just values
            | V.length values >= width -> V.slice 0 width values
            | otherwise -> values V.++ V.replicate (width - V.length values) fill
          Nothing -> V.replicate width fill

-- The fill comes from the store, not from the caller. A thaw that took it as an
-- argument obliged every call site to name the right one, and 'boxedUpdate' —
-- which passed the value being written — named the wrong one. That was harmless
-- only while pages were never absent.
thawBoxedPaged :: BoxedPaged a -> ST s (MutableBoxedPaged s a)
thawBoxedPaged base = do
  mutableBoxedDirty <- newSTRef IntMap.empty
  pure
    MutableBoxedPaged
      { mutableBoxedBase = base
      , mutableBoxedDirty
      , mutableBoxedFill = boxedFill base
      , mutableBoxedPristineBase = case boxedStorage base of
          DenseBoxedPages pages -> V.null pages
          DefaultedBoxedPages _ pages -> directorySize pages == 0
      }

readBoxedPaged :: MutableBoxedPaged s a -> Int -> ST s a
readBoxedPaged MutableBoxedPaged{mutableBoxedBase, mutableBoxedDirty} index = do
  dirty <- readSTRef mutableBoxedDirty
  let !page = pageOf index
      !offset = offsetOf index
  case IntMap.lookup page dirty of
    Just values -> MV.unsafeRead values offset
    Nothing -> case boxedStorage mutableBoxedBase of
      DenseBoxedPages pages ->
        pure (V.unsafeIndex (V.unsafeIndex pages page) offset)
      DefaultedBoxedPages fill pages ->
        case lookupDirectory page pages of
          Just values | offset < V.length values -> pure (V.unsafeIndex values offset)
          _ -> pure fill
{-# INLINE readBoxedPaged #-}

writeBoxedPaged :: MutableBoxedPaged s a -> Int -> a -> ST s ()
writeBoxedPaged paged index value = do
  let !page = pageOf index
      !offset = offsetOf index
  values <- ensureMutableBoxedPage paged page
  MV.unsafeWrite values offset value
{-# INLINE writeBoxedPaged #-}

-- | Whether the thaw found no materialized page. With no writer into the
-- element payload planes inside a transaction — there is none; they are
-- written only through the persistent setters, outside one — this answers for
-- the whole transaction that every slot reports the fill, so a rewrite has
-- nothing to return and a relocation has nothing to move. It is why the
-- near-universal @() () ()@ mesh pays a predictable branch and no page.
boxedThawPristine :: MutableBoxedPaged s a -> Bool
boxedThawPristine = mutableBoxedPristineBase
{-# INLINE boxedThawPristine #-}

-- | Return a contiguous run of slots to the store's fill. A slot no write has
-- reached already reports the fill, so a page absent from both the base and the
-- dirty set is left absent: recycling an element at its default keeps the store
-- sparse. Only a page holding written values is materialized and overwritten.
resetBoxedRange :: MutableBoxedPaged s a -> a -> Int -> Int -> ST s ()
resetBoxedRange paged@MutableBoxedPaged{mutableBoxedBase, mutableBoxedDirty} fill start count
  | count <= 0 = pure ()
  | otherwise = go start
 where
  !end = start + count
  go !index
    | index >= end = pure ()
    | otherwise = do
        dirty <- readSTRef mutableBoxedDirty
        let !page = pageOf index
            !offset = offsetOf index
            !width = min (pageSize - offset) (end - index)
            !written =
              IntMap.member page dirty
                || case boxedStorage mutableBoxedBase of
                  DenseBoxedPages pages ->
                    page < V.length pages
                      && offset < V.length (V.unsafeIndex pages page)
                  DefaultedBoxedPages _ pages ->
                    case lookupDirectory page pages of
                      Just values -> offset < V.length values
                      Nothing -> False
        if written
          then do
            values <- ensureMutableBoxedPage paged page
            MV.set (MV.slice offset width values) fill
            go (index + width)
          else go (index + width)

ensureMutableBoxedPage :: MutableBoxedPaged s a -> Int -> ST s (MV.MVector s a)
ensureMutableBoxedPage MutableBoxedPaged{mutableBoxedBase, mutableBoxedDirty, mutableBoxedFill} page = do
  dirty <- readSTRef mutableBoxedDirty
  case IntMap.lookup page dirty of
    Just values -> pure values
    Nothing -> do
      -- Copying a page slot by slot writes every element twice — once to
      -- establish the fill and once to overwrite it — and pays a write barrier
      -- per store. A full page is one thaw; a short tail page fills only the
      -- slots the base does not reach. This is the same shape the unboxed
      -- store already uses.
      values <- case boxedStorage mutableBoxedBase of
        DenseBoxedPages pages
          | page < V.length pages -> copyExisting (V.unsafeIndex pages page)
          | otherwise -> MV.new pageSize
        DefaultedBoxedPages fill pages ->
          case lookupDirectory page pages of
            Just old -> copyExisting old
            Nothing -> MV.replicate pageSize fill
      writeSTRef mutableBoxedDirty $! IntMap.insert page values dirty
      pure values
 where
  copyExisting old
    | V.length old == pageSize = V.thaw old
    | otherwise = do
        vector <- case mutableBoxedFill of
          Just fill -> MV.replicate pageSize fill
          Nothing -> MV.new pageSize
        V.copy (MV.slice 0 (V.length old) vector) old
        pure vector

freezeBoxedPaged :: Int -> MutableBoxedPaged s a -> ST s (Either BoxedStorageError (BoxedPaged a))
freezeBoxedPaged length' MutableBoxedPaged{mutableBoxedBase, mutableBoxedDirty}
  | length' < 0 = pure (Left (BoxedFreezeLengthNegative length'))
  | otherwise = do
      dirty <- readSTRef mutableBoxedDirty
      frozenDirty <- traverse V.unsafeFreeze dirty
      pure $ case boxedStorage mutableBoxedBase of
        DenseBoxedPages basePages -> do
          pages <-
            traverse
              (densePage frozenDirty basePages)
              [0 .. pageCount - 1]
          pure
            BoxedPaged
              { boxedLength = length'
              , boxedStorage = DenseBoxedPages (V.fromList pages)
              }
        DefaultedBoxedPages fill basePages ->
          let !merged =
                List.foldl'
                  (\pages (page, values) -> insertDirectory page values pages)
                  basePages
                  (IntMap.toAscList frozenDirty)
              !kept = directoryRestrict pageCount merged
              !trimmed = case directoryLookupMax kept of
                Nothing -> kept
                Just (lastPage, values) ->
                  let !lastLength = length' - lastPage * pageSize
                   in if lastLength < V.length values
                        then insertDirectory lastPage (V.take lastLength values) kept
                        else kept
           in Right
                BoxedPaged
                  { boxedLength = length'
                  , boxedStorage = DefaultedBoxedPages fill trimmed
                  }
 where
  !pageCount = if length' <= 0 then 0 else pageOf (length' + pageSize - 1)

  densePage frozenDirty basePages page =
    let !width = min pageSize (length' - page * pageSize)
     in case IntMap.lookup page frozenDirty of
          Just values -> Right (V.take width values)
          Nothing -> case basePages V.!? page of
            Just values -> Right (V.take width values)
            Nothing -> Left (BoxedFreezeDensePageMissing page pageCount)

boxedUpdate :: Int -> a -> BoxedPaged a -> BoxedPaged a
boxedUpdate index value source@BoxedPaged{boxedStorage} =
  source{boxedStorage = updateStorage boxedStorage}
 where
  !page = pageOf index
  !offset = offsetOf index

  updateStorage (DenseBoxedPages pages) =
    let !values = V.unsafeIndex pages page
        !updated = V.modify (\mutable -> MV.unsafeWrite mutable offset value) values
     in DenseBoxedPages (V.modify (\mutable -> MV.unsafeWrite mutable page updated) pages)
  updateStorage (DefaultedBoxedPages fill pages) =
    let !values = case lookupDirectory page pages of
          Just existing
            | V.length existing == pageSize -> existing
            | otherwise -> existing V.++ V.replicate (pageSize - V.length existing) fill
          Nothing -> V.replicate pageSize fill
        !updated = V.modify (\mutable -> MV.unsafeWrite mutable offset value) values
     in DefaultedBoxedPages fill (insertDirectory page updated pages)